Update missing model config lookup

This commit is contained in:
filipstrand 2025-01-19 23:23:19 +01:00
parent 976573c837
commit 6fbbc9e5f7

View File

@ -1,6 +1,6 @@
import mlx.core.random as random
from mflux import Config, Flux1, ModelConfig
from mflux import Config, Flux1, ModelLookup
from mflux.config.runtime_config import RuntimeConfig
from mflux.dreambooth.dataset.dataset import Dataset
from mflux.dreambooth.dataset.iterator import Iterator
@ -28,12 +28,13 @@ class DreamBoothInitializer:
random.seed(training_spec.seed)
# Load the model
model_config = ModelLookup.from_name(training_spec.model)
flux = Flux1(
model_config=ModelConfig.from_alias(training_spec.model),
model_config=model_config,
quantize=training_spec.quantize,
)
runtime_config = RuntimeConfig(
model_config=ModelConfig.from_alias(training_spec.model),
model_config=model_config,
config=Config(
num_inference_steps=training_spec.steps,
width=training_spec.width,