comprehensive ModelConfig refactor to support compatible HuggingFace dev/schnell models

README.md

@@ -263,7 +263,7 @@ Or, with the correct python environment active, create and run a separate script
 from mflux import Flux1, Config
 
 # Load the model
-flux = Flux1.from_alias(
+flux = Flux1.from_name(
    alias="schnell",  # "schnell" or "dev"
    quantize=8,       # 4 or 8
 )

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "mflux"
-version = "0.5.1"
+version = "0.6.0"
 description = "A MLX port of FLUX based on the Huggingface Diffusers implementation."
 readme = "README.md"
 keywords = ["diffusers", "flux", "mlx"]

src/mflux/__init__.py

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

src/mflux/config/model_config.py

@@ -1,27 +1,98 @@
-from enum import Enum
+from dataclasses import dataclass
+from typing import Literal
+
+DEFAULT_TRAIN_STEPS = 1000
+
+KNOWN_SEQUENCE_LENGTH_BY_BASE_MODEL = {"dev": 512, "schnell": 256}
 
 
-class ModelConfig(Enum):
-    FLUX1_DEV = ("black-forest-labs/FLUX.1-dev", "dev", 1000, 512)
-    FLUX1_SCHNELL = ("black-forest-labs/FLUX.1-schnell", "schnell", 1000, 256)
+class ModelConfigError(ValueError):
+    """User error in model config."""
 
-    def __init__(
-        self,
-        model_name: str,
-        alias: str,
-        num_train_steps: int,
-        max_sequence_length: int,
-    ):
-        self.alias = alias
-        self.model_name = model_name
-        self.num_train_steps = num_train_steps
-        self.max_sequence_length = max_sequence_length
 
+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
+
+    @property
+    def alias(self):
+        # maintain compatibility with < 0.4.0 behavior
+        # where alias is the name of an official model
+        if self.model_name.startswith("black-forest-labs/FLUX.1-"):
+            return self.model_name[len("black-forest-labs/FLUX.1-") :].lower()
+        return None
+
+
+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":
-        try:
-            for model in ModelConfig:
-                if model.alias == alias:
-                    return model
-        except KeyError:
-            raise ValueError(f"'{alias}' is not a valid model")
+    def from_name(
+        alias: str,
+        base_model: Literal["dev", "schnell"] | None = None,
+    ) -> ModelConfig:
+        if alias in DefaultModelConfigs:
+            return DefaultModelConfigs[alias]
+
+        if all(["dev" not in alias, "schnell" not in alias, base_model is None]):
+            raise ModelConfigError(
+                "Cannot infer base model and max_sequence_length "
+                f"from model reference: {alias!r}. "
+                "Please specify --base-model [dev | schnell]"
+            )
+
+        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")
+
+        if base_model is None:
+            # infer base model on apparent model namming
+            if "dev" in alias:
+                base_model = "dev"
+            elif "schnell" in alias:
+                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"]
+
+        return ModelConfig(
+            alias,  # actually this arg is model_name
+            DEFAULT_TRAIN_STEPS,
+            max_sequence_length,
+            supports_guidance,
+            base_model,
+        )
+
+
+# maintain old `ModelConfig.from_alias` function name for backwards compatibility in user code and docs
+ModelConfig.from_alias = ModelLookup.from_name
+# 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"]

src/mflux/flux/flux.py

@@ -5,7 +5,7 @@ from mlx import nn
 from tqdm import tqdm
 
 from mflux.config.config import Config
-from mflux.config.model_config import ModelConfig
+from mflux.config.model_config import ModelConfig, ModelLookup
 from mflux.config.runtime_config import RuntimeConfig
 from mflux.error.exceptions import StopImageGenerationException
 from mflux.latent_creator.latent_creator import LatentCreator
@@ -136,12 +136,15 @@ class Flux1(nn.Module):
         )
 
     @staticmethod
-    def from_alias(alias: str, quantize: int | None = None) -> "Flux1":
+    def from_name(model_name: str, quantize: int | None = None) -> "Flux1":
         return Flux1(
-            model_config=ModelConfig.from_alias(alias),
+            model_config=ModelLookup.from_name(model_name),
             quantize=quantize,
         )
 
+    # maintain old `from_alias` function name for backwards compatibility in user code and docs
+    from_alias = from_name
+
     def save_model(self, base_path: str) -> None:
         ModelSaver.save_model(self, self.bits, base_path)
 

src/mflux/generate.py

@@ -1,7 +1,7 @@
 import time
 from pathlib import Path
 
-from mflux import Config, Flux1, ModelConfig, StopImageGenerationException
+from mflux import Config, Flux1, ModelLookup, StopImageGenerationException
 from mflux.ui.cli.parsers import CommandLineParser
 
 
@@ -17,7 +17,7 @@ def main():
 
     # Load the model
     flux = Flux1(
-        model_config=ModelConfig.from_alias(args.model),
+        model_config=ModelLookup.from_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,7 +1,7 @@
 import time
 from pathlib import Path
 
-from mflux import ConfigControlnet, Flux1Controlnet, ModelConfig, StopImageGenerationException
+from mflux import ConfigControlnet, Flux1Controlnet, ModelLookup, StopImageGenerationException
 from mflux.ui.cli.parsers import CommandLineParser
 
 
@@ -16,7 +16,7 @@ def main():
 
     # Load the model
     flux = Flux1Controlnet(
-        model_config=ModelConfig.from_alias(args.model),
+        model_config=ModelLookup.from_name(args.model, base_model=args.base_model),
         quantize=args.quantize,
         local_path=args.path,
         lora_paths=args.lora_paths,

src/mflux/models/transformer/time_text_embed.py

@@ -14,7 +14,7 @@ class TimeTextEmbed(nn.Module):
     def __init__(self, model_config: ModelConfig):
         super().__init__()
         self.text_embedder = TextEmbedder()
-        self.guidance_embedder = GuidanceEmbedder() if model_config == ModelConfig.FLUX1_DEV else None
+        self.guidance_embedder = GuidanceEmbedder() if model_config.supports_guidance else None
         self.timestep_embedder = TimestepEmbedder()
 
     def __call__(

src/mflux/post_processing/generated_image.py

@@ -57,10 +57,13 @@ class GeneratedImage:
             # mflux_version is used by future metadata readers
             # to determine supportability of metadata-derived workflows
             "mflux_version": GeneratedImage.get_version(),
-            "model": str(self.model_config.alias),
+            "model": str(self.model_config.alias)
+            if self.model_config.alias is not None
+            else self.model_config.model_name,
+            "base_model": str(self.model_config.base_model),
             "seed": self.seed,
             "steps": self.steps,
-            "guidance": self.guidance if ModelConfig.FLUX1_DEV else None,  # only the dev model supports guidance
+            "guidance": self.guidance if self.model_config.supports_guidance else None,
             "precision": str(self.precision),
             "quantize": self.quantization,
             "generation_time_seconds": round(self.generation_time, 2),

src/mflux/save.py

@@ -1,4 +1,4 @@
-from mflux import Flux1, ModelConfig
+from mflux import Flux1, ModelLookup
 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=ModelConfig.from_alias(args.model),
+        model_config=ModelLookup.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

@@ -6,6 +6,21 @@ from pathlib import Path
 from mflux.ui import defaults as ui_defaults
 
 
+class ModelSpecAction(argparse.Action):
+    def __call__(self, parser, namespace, values, option_string=None):
+        if values in ["dev", "schnell"]:
+            setattr(namespace, self.dest, values)
+            return
+
+        if values.count("/") != 1:
+            raise argparse.ArgumentError(
+                self, ('Value must be either "dev", "schnell", or "' f'in format "org/model". Got: {values}')
+            )
+
+        # If we got here, values contains exactly one slash
+        setattr(namespace, self.dest, values)
+
+
 # fmt: off
 class CommandLineParser(argparse.ArgumentParser):
 
@@ -19,12 +34,12 @@ class CommandLineParser(argparse.ArgumentParser):
 
     def add_model_arguments(self, path_type: t.Literal["load", "save"] = "load", require_model_arg: bool = True) -> None:
 
-        self.add_argument("--model", "-m", type=str, required=require_model_arg, choices=ui_defaults.MODEL_CHOICES, help=f"The model to use ({' or '.join(ui_defaults.MODEL_CHOICES)}).")
-
+        self.add_argument("--model", "-m", type=str, required=require_model_arg, action=ModelSpecAction, help=f"The model to use ({' or '.join(ui_defaults.MODEL_CHOICES)} or a compatible huggingface repo_id org/model).")
         if path_type == "load":
             self.add_argument("--path", type=str, default=None, help="Local path for loading a model from disk")
         else:
             self.add_argument("--path", type=str, required=True, help="Local path for saving a model to disk.")
+        self.add_argument("--base-model", type=str, required=False, choices=ui_defaults.MODEL_CHOICES, help="When using a third-party huggingface model, explicitly specify whether the base model is dev or schnell")
         self.add_argument("--quantize",  "-q", type=int, choices=ui_defaults.QUANTIZE_CHOICES, default=None, help=f"Quantize the model ({' or '.join(map(str, ui_defaults.QUANTIZE_CHOICES))}, Default is None)")
 
     def add_lora_arguments(self) -> None:
@@ -93,6 +108,9 @@ class CommandLineParser(argparse.ArgumentParser):
                 # when not provided by CLI flag, find it in the config file
                 namespace.model = prior_gen_metadata.get("model", None)
 
+            if namespace.base_model is None:
+                namespace.base_model = prior_gen_metadata.get("base_model", None)
+
             if namespace.prompt is None:
                 namespace.prompt = prior_gen_metadata.get("prompt", None)
 

tests/arg_parser/test_cli_argparser.py

@@ -93,10 +93,12 @@ def test_model_arg_not_in_file(mflux_generate_parser, mflux_generate_minimal_arg
     with patch('sys.argv', mflux_generate_minimal_argv + ['--model', 'dev', '--config-from-metadata', metadata_file.as_posix()]):  # fmt: off
         args = mflux_generate_parser.parse_args()
         assert args.model == "dev"
+        assert args.base_model is None
     # test value read from flag
     with patch('sys.argv', mflux_generate_minimal_argv + ['--model', 'schnell', '--config-from-metadata', metadata_file.as_posix()]):  # fmt: off
         args = mflux_generate_parser.parse_args()
         assert args.model == "schnell"
+        assert args.base_model is None
 
 
 def test_model_arg_in_file(mflux_generate_parser, mflux_generate_minimal_argv, base_metadata_dict, temp_dir):
@@ -114,6 +116,25 @@ def test_model_arg_in_file(mflux_generate_parser, mflux_generate_minimal_argv, b
         assert args.model == "schnell"
 
 
+def test_base_model_arg_in_file(mflux_generate_parser, mflux_generate_minimal_argv, base_metadata_dict, temp_dir):
+    metadata_file = temp_dir / "model.json"
+    with metadata_file.open("wt") as m:
+        base_metadata_dict["model"] = "some-lab/some-model"
+        base_metadata_dict["base_model"] = "dev"
+        json.dump(base_metadata_dict, m, indent=4)
+    # test value read from file
+    with patch('sys.argv', mflux_generate_minimal_argv + ['--config-from-metadata', metadata_file.as_posix()]):  # fmt: off
+        args = mflux_generate_parser.parse_args()
+        assert args.model == "some-lab/some-model"
+        assert args.base_model == "dev"
+    # test value read from flag, overrides value from file
+    with patch('sys.argv', mflux_generate_minimal_argv + ['--base-model', 'schnell', '--config-from-metadata', metadata_file.as_posix()]):  # fmt: off
+        args = mflux_generate_parser.parse_args()
+        assert args.model == "some-lab/some-model"
+        # override metadata base model with CLI --base-model
+        assert args.base_model == "schnell"
+
+
 def test_prompt_arg(mflux_generate_parser, mflux_generate_minimal_argv, base_metadata_dict, temp_dir):
     metadata_file = temp_dir / "prompt.json"
     file_prompt = "origin of the universe"

tests/test_model_config.py

@@ -0,0 +1,68 @@
+import pytest
+
+from mflux.config.model_config import InvalidBaseModel, ModelConfig, ModelConfigError, ModelLookup
+
+
+def test_from_alias_function_redirect():
+    # backwards compatibility for when user follows older docs
+    # but is using a newer mflux version >= 0.5
+    assert ModelConfig.from_alias == ModelLookup.from_name
+
+
+def test_model_config_class_members_and_alias():
+    # these FLUX1_* class members and the alias attribute
+    # existed as members of ModelConfig when it was an Enum
+    # keep them around for backwards compatibility
+    assert ModelConfig.FLUX1_DEV.alias == "dev"
+    assert ModelConfig.FLUX1_SCHNELL.alias == "schnell"
+
+
+def test_bfl_dev():
+    model_attrs = ModelLookup.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_schnell():
+    model_attrs = ModelLookup.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_community_dev_implicit_base_model():
+    model_attrs = ModelLookup.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")
+    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")
+    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")
+    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")
+
+
+def test_invalid_base_model_error():
+    assert pytest.raises(
+        InvalidBaseModel,
+        ModelLookup.from_name,
+        "acme-lab/some-model-who-knows-what-its-based-on",
+        base_model="something_unknown",
+    )