Refactor model config

src/mflux/__init__.py

@@ -1,5 +1,5 @@
 from mflux.config.config import Config
-from mflux.config.model_config import ModelConfig, ModelLookup
+from mflux.config.model_config import ModelConfig
 from mflux.controlnet.flux_controlnet import Flux1Controlnet
 from mflux.error.exceptions import StopImageGenerationException
 from mflux.flux.flux import Flux1
@@ -10,7 +10,6 @@ __all__ = [
     "Flux1Controlnet",
     "Config",
     "ModelConfig",
-    "ModelLookup",
     "ImageUtil",
     "StopImageGenerationException",
 ]

src/mflux/callbacks/callbacks.py

@@ -41,12 +41,12 @@ class Callbacks:
 
     @staticmethod
     def interruption(
-            seed: int,
+        seed: int,
-            prompt: str,
+        prompt: str,
-            step: int,
+        step: int,
-            latents: mx.array,
+        latents: mx.array,
-            config: RuntimeConfig,
+        config: RuntimeConfig,
-            time_steps: tqdm
+        time_steps: tqdm
     ):  # fmt: off
         for subscriber in CallbackRegistry.interrupt_callbacks():
             subscriber.call_interrupt(

src/mflux/config/model_config.py

@@ -1,98 +1,92 @@
-import warnings
+from functools import lru_cache
-from dataclasses import dataclass
 from typing import Literal
 
-DEFAULT_TRAIN_STEPS = 1000
+from mflux.error.error import InvalidBaseModel, ModelConfigError
-
-KNOWN_SEQUENCE_LENGTH_BY_BASE_MODEL = {"dev": 512, "schnell": 256}
 
 
-class ModelConfigError(ValueError):
-    """User error in model config."""
-
-
-class InvalidBaseModel(ModelConfigError):
-    """Invalid base model, cannot infer model properties."""
-
-
-@dataclass
 class ModelConfig:
-    model_name: str
+    def __init__(
-    num_train_steps: int
+        self,
-    max_sequence_length: int
+        alias: str | None,
-    supports_guidance: bool
+        model_name: str,
-    base_model: str | None
+        base_model: str | None,
+        num_train_steps: int,
+        max_sequence_length: int,
+        supports_guidance: bool,
+    ):
+        self.alias = alias
+        self.model_name = model_name
+        self.base_model = base_model
+        self.num_train_steps = num_train_steps
+        self.max_sequence_length = max_sequence_length
+        self.supports_guidance = supports_guidance
 
-
-DefaultModelConfigs = {
-    "dev": ModelConfig(
-        model_name="black-forest-labs/FLUX.1-dev",
-        num_train_steps=DEFAULT_TRAIN_STEPS,
-        max_sequence_length=KNOWN_SEQUENCE_LENGTH_BY_BASE_MODEL["dev"],
-        supports_guidance=True,
-        base_model=None,
-    ),
-    "schnell": ModelConfig(
-        model_name="black-forest-labs/FLUX.1-schnell",
-        num_train_steps=DEFAULT_TRAIN_STEPS,
-        max_sequence_length=KNOWN_SEQUENCE_LENGTH_BY_BASE_MODEL["schnell"],
-        supports_guidance=False,
-        base_model=None,
-    ),
-}
-
-
-class ModelLookup:
     @staticmethod
-    def from_alias(alias: str) -> ModelConfig:
+    @lru_cache
-        warnings.warn(
+    def dev() -> "ModelConfig":
-            "from_alias is deprecated and will be removed in a future release. Please use from_name instead.",
+        return ModelConfig(
-            DeprecationWarning,
+            alias="dev",
-            stacklevel=2,
+            model_name="black-forest-labs/FLUX.1-dev",
+            base_model=None,
+            num_train_steps=1000,
+            max_sequence_length=512,
+            supports_guidance=True,
+        )
+
+    @staticmethod
+    @lru_cache
+    def schnell() -> "ModelConfig":
+        return ModelConfig(
+            alias="schnell",
+            model_name="black-forest-labs/FLUX.1-schnell",
+            base_model=None,
+            num_train_steps=1000,
+            max_sequence_length=256,
+            supports_guidance=False,
         )
-        return ModelLookup.from_name(model_name=alias, base_model=None)
 
     @staticmethod
     def from_name(
         model_name: str,
         base_model: Literal["dev", "schnell"] | None = None,
-    ) -> ModelConfig:
+    ) -> "ModelConfig":
-        if model_name in DefaultModelConfigs:
+        dev = ModelConfig.dev()
-            return DefaultModelConfigs[model_name]
+        schnell = ModelConfig.schnell()
 
-        if all(["dev" not in model_name, "schnell" not in model_name, base_model is None]):
+        # 0. Validate explicit base_model
-            raise ModelConfigError(
+        allowed_names = [dev.alias, dev.model_name, schnell.alias, schnell.model_name]
-                "Cannot infer base model and max_sequence_length "
+        if base_model and base_model not in allowed_names:
-                f"from model reference: {model_name!r}. "
+            raise InvalidBaseModel(f"Invalid base_model. Choose one of {allowed_names}")
-                "Please specify --base-model [dev | schnell]"
-            )
 
-        if base_model is not None and base_model not in ["dev", "schnell"]:
+        # 1. If model_name is "dev" or "schnell" then simply return
-            raise InvalidBaseModel("As of this version, mflux only recognizes base models dev or schnell")
+        if model_name == dev.model_name or model_name == dev.alias:
+            return dev
+        if model_name == schnell.model_name or model_name == schnell.alias:
+            return schnell
 
-        if base_model is None:
+        # 1. Determine the appropriate base model
-            # infer base model on apparent model naming
+        default_base = None
+        if not base_model:
             if "dev" in model_name:
-                base_model = "dev"
+                default_base = dev
             elif "schnell" in model_name:
-                base_model = "schnell"
+                default_base = schnell
-
+            else:
-        if base_model == "dev":
+                raise ModelConfigError(f"Cannot infer base_model from {model_name}. Specify --base-model.")
-            supports_guidance = True
+        elif base_model == dev.model_name or base_model == dev.alias:
-            max_sequence_length = KNOWN_SEQUENCE_LENGTH_BY_BASE_MODEL["dev"]
+            default_base = dev
-        elif base_model == "schnell":
+        elif base_model == schnell.model_name or base_model == schnell.alias:
-            supports_guidance = False
+            default_base = schnell
-            max_sequence_length = KNOWN_SEQUENCE_LENGTH_BY_BASE_MODEL["schnell"]
 
+        # 2. Construct the config based on the model name and base default
         return ModelConfig(
+            alias=default_base.alias,
             model_name=model_name,
-            num_train_steps=DEFAULT_TRAIN_STEPS,
+            base_model=default_base.model_name,
-            max_sequence_length=max_sequence_length,
+            num_train_steps=default_base.num_train_steps,
-            supports_guidance=supports_guidance,
+            max_sequence_length=default_base.max_sequence_length,
-            base_model=base_model,
+            supports_guidance=default_base.supports_guidance,
         )
 
-
+    def is_dev(self) -> bool:
-# keep these class members to be backwards compatible with < 0.5.0 ModelConfig Enum implementation
+        return self.alias == "dev"
-ModelConfig.FLUX1_DEV = DefaultModelConfigs["dev"]
-ModelConfig.FLUX1_SCHNELL = DefaultModelConfigs["schnell"]

src/mflux/config/runtime_config.py

@@ -70,9 +70,9 @@ class RuntimeConfig:
         return None
 
     @staticmethod
-    def _create_sigmas(config, model) -> mx.array:
+    def _create_sigmas(config: Config, model_config: ModelConfig) -> mx.array:
         sigmas = RuntimeConfig._create_sigmas_values(config.num_inference_steps)
-        if model == ModelConfig.FLUX1_DEV:
+        if model_config.is_dev():
             sigmas = RuntimeConfig._shift_sigmas(sigmas=sigmas, width=config.width, height=config.height)
         return sigmas
 

src/mflux/controlnet/flux_controlnet.py

@@ -52,7 +52,7 @@ class Flux1Controlnet(nn.Module):
         seed: int,
         prompt: str,
         controlnet_image_path: str,
-        config: Config = Config(),
+        config: Config,
     ) -> GeneratedImage:
         # 0. Create a new runtime config based on the model type and input parameters
         config = RuntimeConfig(config, self.model_config)

src/mflux/dreambooth/dreambooth_initializer.py

@@ -1,6 +1,6 @@
 import mlx.core.random as random
 
-from mflux import Config, Flux1, ModelLookup
+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
@@ -28,7 +28,7 @@ class DreamBoothInitializer:
         random.seed(training_spec.seed)
 
         # Load the model
-        model_config = ModelLookup.from_name(training_spec.model)
+        model_config = ModelConfig.from_name(training_spec.model)
         flux = Flux1(
             model_config=model_config,
             quantize=training_spec.quantize,

src/mflux/error/error.py

@@ -0,0 +1,6 @@
+class ModelConfigError(ValueError):
+    """User error in model config."""
+
+
+class InvalidBaseModel(ModelConfigError):
+    """Invalid base model, cannot infer model properties."""

src/mflux/flux/flux.py

@@ -4,7 +4,7 @@ from tqdm import tqdm
 
 from mflux.callbacks.callbacks import Callbacks
 from mflux.config.config import Config
-from mflux.config.model_config import ModelConfig, ModelLookup
+from mflux.config.model_config import ModelConfig
 from mflux.config.runtime_config import RuntimeConfig
 from mflux.flux.flux_initializer import FluxInitializer
 from mflux.latent_creator.latent_creator import Img2Img, LatentCreator
@@ -47,7 +47,7 @@ class Flux1(nn.Module):
         self,
         seed: int,
         prompt: str,
-        config: Config = Config(),
+        config: Config,
     ) -> GeneratedImage:
         # 0. Create a new runtime config based on the model type and input parameters
         config = RuntimeConfig(config, self.model_config)
@@ -139,7 +139,7 @@ class Flux1(nn.Module):
     @staticmethod
     def from_name(model_name: str, quantize: int | None = None) -> "Flux1":
         return Flux1(
-            model_config=ModelLookup.from_name(model_name=model_name, base_model=None),
+            model_config=ModelConfig.from_name(model_name=model_name, base_model=None),
             quantize=quantize,
         )
 

src/mflux/flux/flux_initializer.py

@@ -1,3 +1,4 @@
+from mflux import ModelConfig
 from mflux.controlnet.transformer_controlnet import TransformerControlnet
 from mflux.controlnet.weight_handler_controlnet import WeightHandlerControlnet
 from mflux.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder
@@ -16,7 +17,7 @@ class FluxInitializer:
     @staticmethod
     def init(
         flux_model,
-        model_config,
+        model_config: ModelConfig,
         quantize: int | None,
         local_path: str | None,
         lora_paths: list[str] | None,
@@ -82,7 +83,7 @@ class FluxInitializer:
     @staticmethod
     def init_controlnet(
         flux_model,
-        model_config,
+        model_config: ModelConfig,
         quantize: int | None,
         local_path: str | None,
         lora_paths: list[str] | None,

src/mflux/generate.py

@@ -1,4 +1,4 @@
-from mflux import Config, Flux1, ModelLookup, StopImageGenerationException
+from mflux import Config, Flux1, ModelConfig, StopImageGenerationException
 from mflux.callbacks.callback_registry import CallbackRegistry
 from mflux.callbacks.instances.stepwise_handler import StepwiseHandler
 from mflux.ui.cli.parsers import CommandLineParser
@@ -16,7 +16,7 @@ def main():
 
     # 1. Load the model
     flux = Flux1(
-        model_config=ModelLookup.from_name(model_name=args.model, base_model=args.base_model),
+        model_config=ModelConfig.from_name(model_name=args.model, base_model=args.base_model),
         quantize=args.quantize,
         local_path=args.path,
         lora_paths=args.lora_paths,

src/mflux/generate_controlnet.py

@@ -1,4 +1,4 @@
-from mflux import Config, Flux1Controlnet, ModelLookup, StopImageGenerationException
+from mflux import Config, Flux1Controlnet, ModelConfig, StopImageGenerationException
 from mflux.callbacks.callback_registry import CallbackRegistry
 from mflux.callbacks.instances.canny_saver import CannyImageSaver
 from mflux.callbacks.instances.stepwise_handler import StepwiseHandler
@@ -17,7 +17,7 @@ def main():
 
     # 1. Load the model
     flux = Flux1Controlnet(
-        model_config=ModelLookup.from_name(model_name=args.model, base_model=args.base_model),
+        model_config=ModelConfig.from_name(model_name=args.model, base_model=args.base_model),
         quantize=args.quantize,
         local_path=args.path,
         lora_paths=args.lora_paths,

src/mflux/save.py

@@ -1,4 +1,4 @@
-from mflux import Flux1, ModelLookup
+from mflux import Flux1, ModelConfig
 from mflux.ui.cli.parsers import CommandLineParser
 
 
@@ -11,7 +11,7 @@ def main():
     print(f"Saving model {args.model} with quantization level {args.quantize}\n")
 
     flux = Flux1(
-        model_config=ModelLookup.from_name(args.model, base_model=args.base_model),
+        model_config=ModelConfig.from_name(args.model, base_model=args.base_model),
         quantize=args.quantize,
         lora_paths=args.lora_paths,
         lora_scales=args.lora_scales,

src/mflux/ui/cli/parsers.py

@@ -16,7 +16,7 @@ class ModelSpecAction(argparse.Action):
 
         if values.count("/") != 1:
             raise argparse.ArgumentError(
-                self, ('Value must be either "dev", "schnell", or "' f'in format "org/model". Got: {values}')
+                self, 'Value must be either "dev", "schnell", or "' f'in format "org/model". Got: {values}'
             )
 
         # If we got here, values contains exactly one slash
@@ -26,8 +26,8 @@ class ModelSpecAction(argparse.Action):
 # fmt: off
 class CommandLineParser(argparse.ArgumentParser):
 
-    def __init__(self, *pargs, **kwargs):
+    def __init__(self, *args, **kwargs):
-        super().__init__(*pargs, **kwargs)
+        super().__init__(*args, **kwargs)
         self.supports_metadata_config = False
         self.supports_image_generation = False
         self.supports_controlnet = False

tests/dreambooth/test_train_and_load_weights.py

@@ -42,7 +42,7 @@ class TestTrainAndLoadWeights:
 
             # When: Loading a new Flux instance with the trained LoRA...
             fluxB = Flux1(
-                model_config=ModelConfig.FLUX1_DEV,
+                model_config=ModelConfig.dev(),
                 quantize=4,
                 lora_paths=[LORA_FILE],
                 lora_scales=[1.0],

tests/image_generation/test_generate_image.py

@@ -9,7 +9,7 @@ class TestImageGenerator:
         ImageGeneratorTestHelper.assert_matches_reference_image(
             reference_image_path="reference_schnell.png",
             output_image_path=TestImageGenerator.OUTPUT_IMAGE_FILENAME,
-            model_config=ModelConfig.FLUX1_SCHNELL,
+            model_config=ModelConfig.schnell(),
             steps=2,
             seed=42,
             height=341,
@@ -21,7 +21,7 @@ class TestImageGenerator:
         ImageGeneratorTestHelper.assert_matches_reference_image(
             reference_image_path="reference_dev.png",
             output_image_path=TestImageGenerator.OUTPUT_IMAGE_FILENAME,
-            model_config=ModelConfig.FLUX1_DEV,
+            model_config=ModelConfig.dev(),
             steps=15,
             seed=42,
             height=341,
@@ -33,7 +33,7 @@ class TestImageGenerator:
         ImageGeneratorTestHelper.assert_matches_reference_image(
             reference_image_path="reference_dev_lora.png",
             output_image_path=TestImageGenerator.OUTPUT_IMAGE_FILENAME,
-            model_config=ModelConfig.FLUX1_DEV,
+            model_config=ModelConfig.dev(),
             steps=15,
             seed=42,
             height=341,
@@ -47,7 +47,7 @@ class TestImageGenerator:
         ImageGeneratorTestHelper.assert_matches_reference_image(
             reference_image_path="reference_dev_lora_multiple.png",
             output_image_path=TestImageGenerator.OUTPUT_IMAGE_FILENAME,
-            model_config=ModelConfig.FLUX1_DEV,
+            model_config=ModelConfig.dev(),
             steps=15,
             seed=42,
             height=341,
@@ -63,7 +63,7 @@ class TestImageGenerator:
             init_image_path="reference_dev_lora.png",
             init_image_strength=0.4,
             output_image_path=TestImageGenerator.OUTPUT_IMAGE_FILENAME,
-            model_config=ModelConfig.FLUX1_DEV,
+            model_config=ModelConfig.dev(),
             steps=8,
             seed=44,
             height=341,

tests/image_generation/test_generate_image_controlnet.py

@@ -11,7 +11,7 @@ class TestImageGeneratorControlnet:
             reference_image_path="reference_controlnet_schnell.png",
             output_image_path=TestImageGeneratorControlnet.OUTPUT_IMAGE_FILENAME,
             controlnet_image_path=TestImageGeneratorControlnet.CONTROLNET_REFERENCE_FILENAME,
-            model_config=ModelConfig.FLUX1_SCHNELL,
+            model_config=ModelConfig.schnell(),
             steps=2,
             seed=43,
             prompt="The joker with a hat and a cane",
@@ -23,7 +23,7 @@ class TestImageGeneratorControlnet:
             reference_image_path="reference_controlnet_dev.png",
             output_image_path=TestImageGeneratorControlnet.OUTPUT_IMAGE_FILENAME,
             controlnet_image_path=TestImageGeneratorControlnet.CONTROLNET_REFERENCE_FILENAME,
-            model_config=ModelConfig.FLUX1_DEV,
+            model_config=ModelConfig.dev(),
             steps=15,
             seed=42,
             prompt="The joker with a hat and a cane",
@@ -35,7 +35,7 @@ class TestImageGeneratorControlnet:
             reference_image_path="reference_controlnet_dev_lora.png",
             output_image_path=TestImageGeneratorControlnet.OUTPUT_IMAGE_FILENAME,
             controlnet_image_path=TestImageGeneratorControlnet.CONTROLNET_REFERENCE_FILENAME,
-            model_config=ModelConfig.FLUX1_DEV,
+            model_config=ModelConfig.dev(),
             steps=15,
             seed=43,
             prompt="mkym this is made of wool, The joker with a hat and a cane",

tests/model_config/test_model_config.py

@@ -1,68 +1,69 @@
 import pytest
 
-from mflux.config.model_config import InvalidBaseModel, ModelConfigError, ModelLookup
+from mflux.config.model_config import ModelConfig
+from mflux.error.error import InvalidBaseModel, ModelConfigError
 
 
 def test_bfl_dev():
-    model_attrs = ModelLookup.from_name("dev")
+    model_attrs = ModelConfig.from_name("dev")
     assert model_attrs.model_name.startswith("black-forest-labs/")
     assert model_attrs.max_sequence_length == 512
     assert model_attrs.supports_guidance is True
 
 
-def test_bfl_dev_from_alias():
+def test_bfl_dev_full_name():
-    model_attrs = ModelLookup.from_alias("dev")
+    model_attrs = ModelConfig.from_name("black-forest-labs/FLUX.1-dev")
     assert model_attrs.model_name.startswith("black-forest-labs/")
     assert model_attrs.max_sequence_length == 512
     assert model_attrs.supports_guidance is True
 
 
 def test_bfl_schnell():
-    model_attrs = ModelLookup.from_name("schnell")
+    model_attrs = ModelConfig.from_name("schnell")
     assert model_attrs.model_name.startswith("black-forest-labs/")
     assert model_attrs.max_sequence_length == 256
     assert model_attrs.supports_guidance is False
 
 
-def test_bfl_schnell_from_alias():
+def test_bfl_schnell_full_name():
-    model_attrs = ModelLookup.from_alias("schnell")
+    model_attrs = ModelConfig.from_name("black-forest-labs/FLUX.1-schnell")
     assert model_attrs.model_name.startswith("black-forest-labs/")
     assert model_attrs.max_sequence_length == 256
     assert model_attrs.supports_guidance is False
 
 
 def test_community_dev_implicit_base_model():
-    model_attrs = ModelLookup.from_name("acme-lab/some-awesome-dev-model")
+    model_attrs = ModelConfig.from_name("acme-lab/some-awesome-dev-model")
     assert model_attrs.max_sequence_length == 512
     assert model_attrs.supports_guidance is True
 
 
 def test_community_schnell_implicit_base_model():
-    model_attrs = ModelLookup.from_name("acme-lab/some-quick-schnell-model")
+    model_attrs = ModelConfig.from_name("acme-lab/some-quick-schnell-model")
     assert model_attrs.max_sequence_length == 256
     assert model_attrs.supports_guidance is False
 
 
 def test_community_dev_explicit_base_model():
-    model_attrs = ModelLookup.from_name("acme-lab/some-awesome-model", base_model="dev")
+    model_attrs = ModelConfig.from_name("acme-lab/some-awesome-model", base_model="dev")
     assert model_attrs.max_sequence_length == 512
     assert model_attrs.supports_guidance is True
 
 
 def test_community_schnell_explicit_base_model():
-    model_attrs = ModelLookup.from_name("acme-lab/some-awesome-model", base_model="schnell")
+    model_attrs = ModelConfig.from_name("acme-lab/some-awesome-model", base_model="schnell")
     assert model_attrs.max_sequence_length == 256
     assert model_attrs.supports_guidance is False
 
 
 def test_model_config_error():
-    assert pytest.raises(ModelConfigError, ModelLookup.from_name, "acme-lab/some-model-who-knows-what-its-based-on")
+    assert pytest.raises(ModelConfigError, ModelConfig.from_name, "acme-lab/some-model-who-knows-what-its-based-on")
 
 
 def test_invalid_base_model_error():
     assert pytest.raises(
         InvalidBaseModel,
-        ModelLookup.from_name,
+        ModelConfig.from_name,
         "acme-lab/some-model-who-knows-what-its-based-on",
         base_model="something_unknown",
     )

tests/model_saving/test_model_saving.py

@@ -12,7 +12,7 @@ class TestModelSaving:
         try:
             # given a saved quantized model (and an image from that model)
             fluxA = Flux1(
-                model_config=ModelConfig.FLUX1_DEV,
+                model_config=ModelConfig.dev(),
                 quantize=4,
             )
             image1 = fluxA.generate_image(
@@ -29,7 +29,7 @@ class TestModelSaving:
 
             # when loading the quantized model (also without specifying bits)
             fluxB = Flux1(
-                model_config=ModelConfig.FLUX1_DEV,
+                model_config=ModelConfig.dev(),
                 local_path=PATH,
             )
 

tests/model_saving/test_model_saving_lora.py

@@ -13,7 +13,7 @@ class TestModelSavingLora:
         try:
             # given a saved quantized model on disk (without LoRA)...
             fluxA = Flux1(
-                model_config=ModelConfig.FLUX1_SCHNELL,
+                model_config=ModelConfig.schnell(),
                 quantize=4,
             )
             fluxA.save_model(PATH)
@@ -21,7 +21,7 @@ class TestModelSavingLora:
 
             # ...and given an 'on-the-fly' quantized model which we generate an image from
             fluxB = Flux1(
-                model_config=ModelConfig.FLUX1_SCHNELL,
+                model_config=ModelConfig.schnell(),
                 quantize=4,
                 lora_paths=TestModelSavingLora.get_lora_path(),
                 lora_scales=[1.0],
@@ -39,7 +39,7 @@ class TestModelSavingLora:
 
             # when loading the quantized model from a local path (also without specifying bits) with a LoRA...
             fluxC = Flux1(
-                model_config=ModelConfig.FLUX1_SCHNELL,
+                model_config=ModelConfig.schnell(),
                 local_path=PATH,
                 lora_paths=TestModelSavingLora.get_lora_path(),
                 lora_scales=[1.0],