Merge pull request #125 from filipstrand/more-quantization

More quantization options

LICENSE

@@ -1,6 +1,6 @@
 MIT License
 
-Copyright (c) 2024 Filip Strand
+Copyright (c) 2025 Filip Strand
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

README.md

@@ -378,7 +378,7 @@ Luxury food photograph of an italian Linguine pasta alle vongole dish with lots
 
 ### 🗜️ Quantization
 
-MFLUX supports running FLUX in 4-bit or 8-bit quantized mode. Running a quantized version can greatly speed up the
+MFLUX supports running FLUX in 3, 4, 6, or 8-bit quantized mode. Running a quantized version can greatly speed up the
 generation process and reduce the memory consumption by several gigabytes. [Quantized models also take up less disk space](#-size-comparisons-for-quantized-models).
 
 ```sh
@@ -404,11 +404,10 @@ running the 8-bit quantized version on this particular machine. Unlike the non-q
 
 The model sizes for both `schnell` and `dev` at various quantization levels are as follows:
 
-| 4 bit  | 8 bit   | Original (16 bit) |
-|--------|---------|-------------------|
-| 9.85GB | 18.16GB | 33.73GB           |
+| 3 bit  | 4 bit  | 6 bit   | 8 bit   | Original (16 bit) |
+|--------|--------|---------|---------|-------------------|
+| 7.52GB | 9.61GB | 13.81GB | 18.01GB | 33.73GB           |
 
-The reason weights sizes are not fully cut in half is because a small number of weights are not quantized and kept at full precision.
 
 #### 💾 Saving a quantized version to disk
 

pyproject.toml

@@ -15,7 +15,7 @@ requires-python = ">=3.10"
 dependencies = [
     "huggingface-hub>=0.24.5,<1.0",
     "matplotlib>=3.9.2,<4.0",
-    "mlx>=0.20.0,<1.0",
+    "mlx>=0.22.0,<1.0",
     "numpy>=2.0.1,<3.0",
     "opencv-python>=4.10.0,<5.0",
     "piexif>=1.1.3,<2.0",

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

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

@@ -11,6 +11,7 @@ from mflux.controlnet.transformer_controlnet import TransformerControlnet
 from mflux.flux.flux_initializer import FluxInitializer
 from mflux.latent_creator.latent_creator import LatentCreator
 from mflux.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder
+from mflux.models.text_encoder.prompt_encoder import PromptEncoder
 from mflux.models.text_encoder.t5_encoder.t5_encoder import T5Encoder
 from mflux.models.transformer.transformer import Transformer
 from mflux.models.vae.vae import VAE
@@ -51,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)
@@ -73,10 +74,14 @@ class Flux1Controlnet(nn.Module):
         )  # fmt: off
 
         # 3. Encode the prompt
-        t5_tokens = self.t5_tokenizer.tokenize(prompt)
-        clip_tokens = self.clip_tokenizer.tokenize(prompt)
-        prompt_embeds = self.t5_text_encoder(t5_tokens)
-        pooled_prompt_embeds = self.clip_text_encoder(clip_tokens)
+        prompt_embeds, pooled_prompt_embeds = PromptEncoder.encode_prompt(
+            prompt=prompt,
+            prompt_cache=self.prompt_cache,
+            t5_tokenizer=self.t5_tokenizer,
+            clip_tokenizer=self.clip_tokenizer,
+            t5_text_encoder=self.t5_text_encoder,
+            clip_text_encoder=self.clip_text_encoder,
+        )
 
         # (Optional) Call subscribers for beginning of loop
         Callbacks.before_loop(

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,11 +4,12 @@ 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
 from mflux.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder
+from mflux.models.text_encoder.prompt_encoder import PromptEncoder
 from mflux.models.text_encoder.t5_encoder.t5_encoder import T5Encoder
 from mflux.models.transformer.transformer import Transformer
 from mflux.models.vae.vae import VAE
@@ -46,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)
@@ -66,10 +67,14 @@ class Flux1(nn.Module):
         )
 
         # 2. Encode the prompt
-        t5_tokens = self.t5_tokenizer.tokenize(prompt)
-        clip_tokens = self.clip_tokenizer.tokenize(prompt)
-        prompt_embeds = self.t5_text_encoder(t5_tokens)
-        pooled_prompt_embeds = self.clip_text_encoder(clip_tokens)
+        prompt_embeds, pooled_prompt_embeds = PromptEncoder.encode_prompt(
+            prompt=prompt,
+            prompt_cache=self.prompt_cache,
+            t5_tokenizer=self.t5_tokenizer,
+            clip_tokenizer=self.clip_tokenizer,
+            t5_text_encoder=self.t5_text_encoder,
+            clip_text_encoder=self.clip_text_encoder,
+        )
 
         # (Optional) Call subscribers for beginning of loop
         Callbacks.before_loop(
@@ -134,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,13 +17,14 @@ class FluxInitializer:
     @staticmethod
     def init(
         flux_model,
-        model_config,
+        model_config: ModelConfig,
         quantize: int | None,
         local_path: str | None,
         lora_paths: list[str] | None,
         lora_scales: list[float] | None,
     ) -> None:
-        # 0. Set paths and config for later
+        # 0. Set paths, configs and prompt_cache for later
+        flux_model.prompt_cache = {}
         flux_model.lora_paths = lora_paths
         flux_model.lora_scales = lora_scales
         flux_model.model_config = model_config
@@ -81,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/models/text_encoder/prompt_encoder.py

@@ -0,0 +1,33 @@
+import mlx.core as mx
+
+from mflux.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder
+from mflux.models.text_encoder.t5_encoder.t5_encoder import T5Encoder
+from mflux.tokenizer.clip_tokenizer import TokenizerCLIP
+from mflux.tokenizer.t5_tokenizer import TokenizerT5
+
+
+class PromptEncoder:
+    @staticmethod
+    def encode_prompt(
+        prompt: str,
+        prompt_cache: dict[str, (mx.array, mx.array)],
+        t5_tokenizer: TokenizerT5,
+        clip_tokenizer: TokenizerCLIP,
+        t5_text_encoder: T5Encoder,
+        clip_text_encoder: CLIPEncoder,
+    ) -> (mx.array, mx.array):
+        # 1. Return prompt encodings if already cached
+        if prompt in prompt_cache:
+            return prompt_cache[prompt]
+
+        # 1. Encode the prompt
+        t5_tokens = t5_tokenizer.tokenize(prompt)
+        clip_tokens = clip_tokenizer.tokenize(prompt)
+        prompt_embeds = t5_text_encoder(t5_tokens)
+        pooled_prompt_embeds = clip_text_encoder(clip_tokens)
+
+        # 2. Cache the encoded prompt
+        prompt_cache[prompt] = (prompt_embeds, pooled_prompt_embeds)
+
+        # 3. Return prompt encodings
+        return prompt_embeds, pooled_prompt_embeds

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):
-        super().__init__(*pargs, **kwargs)
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
         self.supports_metadata_config = False
         self.supports_image_generation = False
         self.supports_controlnet = False

src/mflux/ui/defaults.py

@@ -7,4 +7,4 @@ MODEL_INFERENCE_STEPS = {
     "dev": 14,
     "schnell": 4,
 }
-QUANTIZE_CHOICES = [4, 8]
+QUANTIZE_CHOICES = [3, 4, 6, 8]

src/mflux/weights/quantization_util.py

@@ -21,12 +21,10 @@ class QuantizationUtil:
 
         if quantize is not None or q_level is not None:
             bits = int(q_level) if q_level is not None else quantize
-            # fmt: off
-            nn.quantize(vae, class_predicate=lambda _, m: isinstance(m, nn.Linear), group_size=64, bits=bits)
-            nn.quantize(transformer, class_predicate=lambda _, m: isinstance(m, nn.Linear) and len(m.weight[1]) > 64, group_size=64, bits=bits)
-            nn.quantize(t5_text_encoder, class_predicate=lambda _, m: isinstance(m, nn.Linear), group_size=64, bits=bits)
-            nn.quantize(clip_text_encoder, class_predicate=lambda _, m: isinstance(m, nn.Linear), group_size=64, bits=bits)
-            # fmt: on
+            nn.quantize(vae, bits=bits)
+            nn.quantize(transformer, bits=bits)
+            nn.quantize(t5_text_encoder, bits=bits)
+            nn.quantize(clip_text_encoder, bits=bits)
 
     @staticmethod
     def quantize_controlnet(
@@ -35,6 +33,7 @@ class QuantizationUtil:
         transformer_controlnet: nn.Module,
     ):
         q_level = weights.meta_data.quantization_level
+
         if quantize is not None or q_level is not None:
             bits = int(q_level) if q_level is not None else quantize
-            nn.quantize(transformer_controlnet, class_predicate=lambda _, m: isinstance(m, nn.Linear) and len(m.weight[1]) > 128, group_size=128, bits=bits)  # fmt: off
+            nn.quantize(transformer_controlnet, bits=bits)

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():
-    model_attrs = ModelLookup.from_alias("dev")
+def test_bfl_dev_full_name():
+    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():
-    model_attrs = ModelLookup.from_alias("schnell")
+def test_bfl_schnell_full_name():
+    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],