75 lines
2.8 KiB
Python
75 lines
2.8 KiB
Python
import mlx.core.random as random
|
|
|
|
from mflux import Config, Flux1, ModelConfig
|
|
from mflux.config.runtime_config import RuntimeConfig
|
|
from mflux.dreambooth.dataset.dataset import Dataset
|
|
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.training_state import TrainingState
|
|
from mflux.dreambooth.statistics.statistics import Statistics
|
|
|
|
|
|
class DreamBoothInitializer:
|
|
@staticmethod
|
|
def initialize(
|
|
config_path: str | None,
|
|
checkpoint_path: str | None,
|
|
) -> (Flux1, RuntimeConfig, TrainingSpec, TrainingState):
|
|
# The training specification describing the details of the training process. It is resolved
|
|
# differently depending on if training starts from scratch or resumes from checkpoint.
|
|
training_spec = TrainingSpec.resolve(
|
|
config_path=config_path,
|
|
checkpoint_path=checkpoint_path
|
|
) # fmt: off
|
|
|
|
# Set global random seed to make training deterministic
|
|
random.seed(training_spec.seed)
|
|
|
|
# Load the model
|
|
flux = Flux1(
|
|
model_config=ModelConfig.from_alias(training_spec.model),
|
|
quantize=training_spec.quantize,
|
|
)
|
|
runtime_config = RuntimeConfig(
|
|
model_config=ModelConfig.from_alias(training_spec.model),
|
|
config=Config(
|
|
num_inference_steps=training_spec.steps,
|
|
width=training_spec.width,
|
|
height=training_spec.height,
|
|
guidance=training_spec.guidance,
|
|
),
|
|
)
|
|
|
|
# Create the optimizer
|
|
optimizer = Optimizer.from_spec(training_spec)
|
|
|
|
# Create the LoRA layers by matching them against the corresponding Flux layers
|
|
lora_layers = LoRALayers.from_spec(flux=flux, training_spec=training_spec)
|
|
|
|
# Prepare the fine-tuning dataset and create the iterator
|
|
dataset = Dataset.prepare_dataset(
|
|
flux=flux,
|
|
raw_data=training_spec.examples,
|
|
width=training_spec.width,
|
|
height=training_spec.height,
|
|
)
|
|
iterator = Iterator.from_spec(
|
|
training_spec=training_spec,
|
|
dataset=dataset
|
|
) # fmt: off
|
|
|
|
# Setup loss statistics
|
|
statistics = Statistics.from_spec(training_spec=training_spec)
|
|
|
|
# The training state consisting of everything that moves during training
|
|
training_state = TrainingState(
|
|
optimizer=optimizer,
|
|
lora_layers=lora_layers,
|
|
iterator=iterator,
|
|
statistics=statistics,
|
|
)
|
|
|
|
return flux, runtime_config, training_spec, training_state
|