import datetime import json import tempfile import zipfile from pathlib import Path from mflux.dreambooth.dataset.iterator import Iterator from mflux.dreambooth.lora_layers.lora_layers import LoRALayers from mflux.dreambooth.optimization.optimizer import Optimizer from mflux.dreambooth.state.training_spec import TrainingSpec from mflux.dreambooth.state.zip_util import ZipUtil from mflux.dreambooth.statistics.statistics import Statistics DREAMBOOTH_PATH_CHECKPOINTS = "_checkpoints" DREAMBOOTH_FILE_NAME_CHECKPOINT = "checkpoint" DREAMBOOTH_FILE_NAME_LORA_ADAPTER = "adapter" DREAMBOOTH_FILE_NAME_OPTIMIZER = "optimizer" DREAMBOOTH_FILE_NAME_ITERATOR = "iterator" DREAMBOOTH_FILE_NAME_LOSS_FILE = "loss" DREAMBOOTH_FILE_NAME_CONFIG_FILE = "config" DREAMBOOTH_PATH_VALIDATION_IMAGES = "_validation/images/" DREAMBOOTH_FILE_NAME_VALIDATION_IMAGE = "validation_image" DREAMBOOTH_PATH_VALIDATION_PLOT = "_validation/plots/" DREAMBOOTH_FILE_NAME_VALIDATION_LOSS = "loss" class TrainingState: def __init__( self, iterator: Iterator, lora_layers: LoRALayers, optimizer: Optimizer, statistics: Statistics, ): self.iterator = iterator self.optimizer = optimizer self.lora_layers = lora_layers self.statistics = statistics def save(self, training_spec: TrainingSpec) -> None: # Create a temporary directory to store files before zipping with tempfile.TemporaryDirectory() as temp_dir: # Save files to temporary directory optimizer_path = Path(temp_dir) / f"{self.iterator.num_iterations:07d}_{DREAMBOOTH_FILE_NAME_OPTIMIZER}.safetensors" # fmt:off lora_path = Path(temp_dir) / f"{self.iterator.num_iterations:07d}_{DREAMBOOTH_FILE_NAME_LORA_ADAPTER}.safetensors" # fmt:off iterator_path = Path(temp_dir) / f"{self.iterator.num_iterations:07d}_{DREAMBOOTH_FILE_NAME_ITERATOR}.json" loss_path = Path(temp_dir) / f"{self.iterator.num_iterations:07d}_{DREAMBOOTH_FILE_NAME_LOSS_FILE}.json" config_path = Path(temp_dir) / f"{self.iterator.num_iterations:07d}_{DREAMBOOTH_FILE_NAME_CONFIG_FILE}.json" checkpoint_path = Path(temp_dir) / f"{DREAMBOOTH_FILE_NAME_CHECKPOINT}.json" # fmt: off paths = [optimizer_path, lora_path, iterator_path, loss_path, config_path, checkpoint_path] # Save individual files to temporary directory self.optimizer.save(optimizer_path) self.lora_layers.save(lora_path, training_spec) self.iterator.save(iterator_path) self.statistics.save(loss_path) self._save_train_config(config_path, training_spec) # For completeness, we also save the (initial) config # Create checkpoint data checkpoint_data = self._create_checkpoint_data(training_spec, self.iterator.start_date_time) with open(checkpoint_path, "w") as json_file: json.dump(checkpoint_data, json_file, indent=4) # Create zip file output_path = Path(training_spec.saver.output_path) / DREAMBOOTH_PATH_CHECKPOINTS output_path.mkdir(parents=True, exist_ok=True) zip_path = output_path / f"{self.iterator.num_iterations:07d}_checkpoint.zip" with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zipf: for file_path in paths: if file_path.exists(): zipf.write(file_path, file_path.name) def _create_checkpoint_data(self, training_spec: TrainingSpec, start_date_time: datetime) -> dict: now = datetime.datetime.now() return { "metadata": { "start": start_date_time.strftime("%Y-%m-%d %H:%M:%S"), "end": now.strftime("%Y-%m-%d %H:%M:%S"), "duration": self._format_duration(start_date_time, now), "number_of_training_examples": self.iterator.dataset.size(), }, "files": { "config": f"{self.iterator.num_iterations:07d}_{DREAMBOOTH_FILE_NAME_CONFIG_FILE}.json", "optimizer": f"{self.iterator.num_iterations:07d}_{DREAMBOOTH_FILE_NAME_OPTIMIZER}.safetensors", "lora_adapter": f"{self.iterator.num_iterations:07d}_{DREAMBOOTH_FILE_NAME_LORA_ADAPTER}.safetensors", "iterator": f"{self.iterator.num_iterations:07d}_{DREAMBOOTH_FILE_NAME_ITERATOR}.json", "loss": f"{self.iterator.num_iterations:07d}_{DREAMBOOTH_FILE_NAME_LOSS_FILE}.json", }, } def should_save(self, training_spec: TrainingSpec) -> bool: return self.iterator.num_iterations % training_spec.saver.checkpoint_frequency == 0 def should_plot_loss(self, training_spec: TrainingSpec) -> bool: if training_spec.instrumentation is None: return False return self.iterator.num_iterations % training_spec.instrumentation.plot_frequency == 0 def should_generate_image(self, training_spec: TrainingSpec) -> bool: if training_spec.instrumentation is None: return False return self.iterator.num_iterations % training_spec.instrumentation.generate_image_frequency == 0 def get_current_validation_image_path(self, training_spec: TrainingSpec) -> Path: output_path = Path(training_spec.saver.output_path) / DREAMBOOTH_PATH_VALIDATION_IMAGES output_path.mkdir(parents=True, exist_ok=True) path = output_path / Path(f"{self.iterator.num_iterations :07d}_{DREAMBOOTH_FILE_NAME_VALIDATION_IMAGE}.png") return path def get_current_loss_plot_path(self, training_spec: TrainingSpec) -> Path: output_path = Path(training_spec.saver.output_path) / DREAMBOOTH_PATH_VALIDATION_PLOT output_path.mkdir(parents=True, exist_ok=True) path = output_path / Path(f"{self.iterator.num_iterations :07d}_{DREAMBOOTH_FILE_NAME_VALIDATION_LOSS}.pdf") return path @staticmethod def _get_parent_path(path: str | None) -> str | None: return None if path is None else str(Path(path).resolve()) @staticmethod def _format_duration(start: datetime, end: datetime) -> str: # Calculate the duration duration = end - start total_seconds = int(duration.total_seconds()) # Calculate hours, minutes, and seconds hours, remainder = divmod(total_seconds, 3600) minutes, seconds = divmod(remainder, 60) # Build the string parts = [] if hours > 0: parts.append(f"{hours} hour{'s' if hours > 1 else ''}") if minutes > 0: parts.append(f"{minutes} minute{'s' if minutes > 1 else ''}") if seconds > 0 or not parts: # Include seconds even if zero if no other parts exist parts.append(f"{seconds} second{'s' if seconds > 1 else ''}") return " and ".join(parts) @staticmethod def _save_train_config(path: Path, training_spec: TrainingSpec) -> None: if training_spec.config_path is not None: with open(Path(training_spec.config_path), "r") as f: data = json.load(f) # Since the zip file can be moved to an arbitrary location, we do a slight # modification and override the image location with the absolute path of the image files data["examples"]["path"] = str(Path(training_spec.config_path).parent / data["examples"]["path"]) else: checkpoint = ZipUtil.unzip( training_spec.checkpoint_path, "checkpoint.json", lambda x: json.load(open(x, "r")) ) data = ZipUtil.unzip( training_spec.checkpoint_path, checkpoint["files"]["config"], lambda x: json.load(open(x, "r")) ) with open(path, "w", encoding="utf-8") as file: json.dump(data, file, indent=4)