diff --git a/LICENSE b/LICENSE index 3a493f8..d1f1eea 100644 --- a/LICENSE +++ b/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 diff --git a/README.md b/README.md index b2c85dd..19f783e 100644 --- a/README.md +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index a613a92..315c650 100644 --- a/pyproject.toml +++ b/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", diff --git a/src/mflux/__init__.py b/src/mflux/__init__.py index 67b368f..b6cdb25 100644 --- a/src/mflux/__init__.py +++ b/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", ] diff --git a/src/mflux/callbacks/callbacks.py b/src/mflux/callbacks/callbacks.py index f3e2499..220b7b9 100644 --- a/src/mflux/callbacks/callbacks.py +++ b/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( diff --git a/src/mflux/config/model_config.py b/src/mflux/config/model_config.py index 46e30e6..8daa6f3 100644 --- a/src/mflux/config/model_config.py +++ b/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" diff --git a/src/mflux/config/runtime_config.py b/src/mflux/config/runtime_config.py index 096338d..2ae1a8c 100644 --- a/src/mflux/config/runtime_config.py +++ b/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 diff --git a/src/mflux/controlnet/flux_controlnet.py b/src/mflux/controlnet/flux_controlnet.py index 5d38a11..78a71e9 100644 --- a/src/mflux/controlnet/flux_controlnet.py +++ b/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( diff --git a/src/mflux/dreambooth/dreambooth_initializer.py b/src/mflux/dreambooth/dreambooth_initializer.py index bd668af..0c85e90 100644 --- a/src/mflux/dreambooth/dreambooth_initializer.py +++ b/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, diff --git a/src/mflux/error/error.py b/src/mflux/error/error.py new file mode 100644 index 0000000..886a2b5 --- /dev/null +++ b/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.""" diff --git a/src/mflux/flux/flux.py b/src/mflux/flux/flux.py index 178b68e..31c4f6b 100644 --- a/src/mflux/flux/flux.py +++ b/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, ) diff --git a/src/mflux/flux/flux_initializer.py b/src/mflux/flux/flux_initializer.py index 2e57ab1..1997486 100644 --- a/src/mflux/flux/flux_initializer.py +++ b/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, diff --git a/src/mflux/generate.py b/src/mflux/generate.py index ea95a02..5666c80 100644 --- a/src/mflux/generate.py +++ b/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, diff --git a/src/mflux/generate_controlnet.py b/src/mflux/generate_controlnet.py index d467b4a..bd5deab 100644 --- a/src/mflux/generate_controlnet.py +++ b/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, diff --git a/src/mflux/models/text_encoder/prompt_encoder.py b/src/mflux/models/text_encoder/prompt_encoder.py new file mode 100644 index 0000000..48fa189 --- /dev/null +++ b/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 diff --git a/src/mflux/save.py b/src/mflux/save.py index e0c04dd..0dfe954 100644 --- a/src/mflux/save.py +++ b/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, diff --git a/src/mflux/ui/cli/parsers.py b/src/mflux/ui/cli/parsers.py index 92a1bb9..8135ac2 100644 --- a/src/mflux/ui/cli/parsers.py +++ b/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 diff --git a/src/mflux/ui/defaults.py b/src/mflux/ui/defaults.py index 7ed4bd7..31a307f 100644 --- a/src/mflux/ui/defaults.py +++ b/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] diff --git a/src/mflux/weights/quantization_util.py b/src/mflux/weights/quantization_util.py index 3592f69..0afa685 100644 --- a/src/mflux/weights/quantization_util.py +++ b/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) diff --git a/tests/dreambooth/test_train_and_load_weights.py b/tests/dreambooth/test_train_and_load_weights.py index 10928eb..58163e9 100644 --- a/tests/dreambooth/test_train_and_load_weights.py +++ b/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], diff --git a/tests/image_generation/test_generate_image.py b/tests/image_generation/test_generate_image.py index 9f18b06..d45fca0 100644 --- a/tests/image_generation/test_generate_image.py +++ b/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, diff --git a/tests/image_generation/test_generate_image_controlnet.py b/tests/image_generation/test_generate_image_controlnet.py index cfc52d9..af9131f 100644 --- a/tests/image_generation/test_generate_image_controlnet.py +++ b/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", diff --git a/tests/model_config/__init__.py b/tests/model_config/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_model_config.py b/tests/model_config/test_model_config.py similarity index 68% rename from tests/test_model_config.py rename to tests/model_config/test_model_config.py index 783b005..99cf540 100644 --- a/tests/test_model_config.py +++ b/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", ) diff --git a/tests/model_saving/test_model_saving.py b/tests/model_saving/test_model_saving.py index 9d2fc08..82a0e7a 100644 --- a/tests/model_saving/test_model_saving.py +++ b/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, ) diff --git a/tests/model_saving/test_model_saving_lora.py b/tests/model_saving/test_model_saving_lora.py index 2e62ee3..acf3a88 100644 --- a/tests/model_saving/test_model_saving_lora.py +++ b/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], diff --git a/tests/resources/reference_controlnet_dev.png b/tests/resources/reference_controlnet_dev.png index 8afedee..1777975 100644 Binary files a/tests/resources/reference_controlnet_dev.png and b/tests/resources/reference_controlnet_dev.png differ diff --git a/tests/resources/reference_controlnet_dev_lora.png b/tests/resources/reference_controlnet_dev_lora.png index 1752c70..ab9912f 100644 Binary files a/tests/resources/reference_controlnet_dev_lora.png and b/tests/resources/reference_controlnet_dev_lora.png differ diff --git a/tests/resources/reference_controlnet_schnell.png b/tests/resources/reference_controlnet_schnell.png index 5963b98..6716c44 100644 Binary files a/tests/resources/reference_controlnet_schnell.png and b/tests/resources/reference_controlnet_schnell.png differ diff --git a/tests/resources/reference_dev.png b/tests/resources/reference_dev.png index a0e0cfc..886f9cf 100644 Binary files a/tests/resources/reference_dev.png and b/tests/resources/reference_dev.png differ diff --git a/tests/resources/reference_dev_image_to_image_result.png b/tests/resources/reference_dev_image_to_image_result.png index 0ba8050..3e3322e 100644 Binary files a/tests/resources/reference_dev_image_to_image_result.png and b/tests/resources/reference_dev_image_to_image_result.png differ diff --git a/tests/resources/reference_dev_lora.png b/tests/resources/reference_dev_lora.png index 0f38df0..6b6abad 100644 Binary files a/tests/resources/reference_dev_lora.png and b/tests/resources/reference_dev_lora.png differ diff --git a/tests/resources/reference_dev_lora_multiple.png b/tests/resources/reference_dev_lora_multiple.png index 9f63d0f..932e112 100644 Binary files a/tests/resources/reference_dev_lora_multiple.png and b/tests/resources/reference_dev_lora_multiple.png differ diff --git a/tests/resources/reference_schnell.png b/tests/resources/reference_schnell.png index fd345fe..f8656c5 100644 Binary files a/tests/resources/reference_schnell.png and b/tests/resources/reference_schnell.png differ