diff --git a/src/mflux/dreambooth/dreambooth_initializer.py b/src/mflux/dreambooth/dreambooth_initializer.py
index ab7ce10..bd668af 100644
--- a/src/mflux/dreambooth/dreambooth_initializer.py
+++ b/src/mflux/dreambooth/dreambooth_initializer.py
@@ -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,