Update missing model config lookup
This commit is contained in:
parent
976573c837
commit
6fbbc9e5f7
@ -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,
|
||||
|
||||
Loading…
Reference in New Issue
Block a user