From 344ac90c74041ed1e5bc9b3f2790089d9c98d9bb Mon Sep 17 00:00:00 2001 From: Anthony Wu Date: Fri, 17 Jan 2025 06:27:33 -0800 Subject: [PATCH] comprehensive ModelConfig refactor to support compatible HuggingFace dev/schnell models --- README.md | 2 +- pyproject.toml | 2 +- src/mflux/__init__.py | 3 +- src/mflux/config/model_config.py | 115 ++++++++++++++---- src/mflux/flux/flux.py | 9 +- src/mflux/generate.py | 4 +- src/mflux/generate_controlnet.py | 4 +- .../models/transformer/time_text_embed.py | 2 +- src/mflux/post_processing/generated_image.py | 7 +- src/mflux/save.py | 4 +- src/mflux/ui/cli/parsers.py | 22 +++- tests/arg_parser/test_cli_argparser.py | 21 ++++ tests/test_model_config.py | 68 +++++++++++ 13 files changed, 224 insertions(+), 39 deletions(-) create mode 100644 tests/test_model_config.py diff --git a/README.md b/README.md index 78cea9c..130e8af 100644 --- a/README.md +++ b/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 ) diff --git a/pyproject.toml b/pyproject.toml index 382168a..a613a92 100644 --- a/pyproject.toml +++ b/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"] diff --git a/src/mflux/__init__.py b/src/mflux/__init__.py index 553f634..83273fc 100644 --- a/src/mflux/__init__.py +++ b/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", ] diff --git a/src/mflux/config/model_config.py b/src/mflux/config/model_config.py index e67a581..c143567 100644 --- a/src/mflux/config/model_config.py +++ b/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"] diff --git a/src/mflux/flux/flux.py b/src/mflux/flux/flux.py index 08e0de5..f04b7b2 100644 --- a/src/mflux/flux/flux.py +++ b/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) diff --git a/src/mflux/generate.py b/src/mflux/generate.py index 236b6b0..878f78e 100644 --- a/src/mflux/generate.py +++ b/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, diff --git a/src/mflux/generate_controlnet.py b/src/mflux/generate_controlnet.py index acd7443..2d116d6 100644 --- a/src/mflux/generate_controlnet.py +++ b/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, diff --git a/src/mflux/models/transformer/time_text_embed.py b/src/mflux/models/transformer/time_text_embed.py index 3492a83..eb6f24f 100644 --- a/src/mflux/models/transformer/time_text_embed.py +++ b/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__( diff --git a/src/mflux/post_processing/generated_image.py b/src/mflux/post_processing/generated_image.py index 893da42..b8ef65b 100644 --- a/src/mflux/post_processing/generated_image.py +++ b/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), diff --git a/src/mflux/save.py b/src/mflux/save.py index 21c0b57..e0c04dd 100644 --- a/src/mflux/save.py +++ b/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, diff --git a/src/mflux/ui/cli/parsers.py b/src/mflux/ui/cli/parsers.py index 215a4ed..7710e60 100644 --- a/src/mflux/ui/cli/parsers.py +++ b/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) diff --git a/tests/arg_parser/test_cli_argparser.py b/tests/arg_parser/test_cli_argparser.py index 0a043e1..1558a3b 100644 --- a/tests/arg_parser/test_cli_argparser.py +++ b/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" diff --git a/tests/test_model_config.py b/tests/test_model_config.py new file mode 100644 index 0000000..636b093 --- /dev/null +++ b/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", + )