163 lines
7.7 KiB
Python
163 lines
7.7 KiB
Python
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)
|