From 3cecfae75b6f9f8e0b0a2c0aaa635b3c771f8a33 Mon Sep 17 00:00:00 2001 From: filipstrand Date: Wed, 12 Feb 2025 22:27:12 +0100 Subject: [PATCH] Refactor model config --- src/mflux/__init__.py | 3 +- src/mflux/callbacks/callbacks.py | 12 +- src/mflux/config/model_config.py | 144 +++++++++--------- src/mflux/config/runtime_config.py | 4 +- src/mflux/controlnet/flux_controlnet.py | 2 +- .../dreambooth/dreambooth_initializer.py | 4 +- src/mflux/error/error.py | 6 + src/mflux/flux/flux.py | 6 +- src/mflux/flux/flux_initializer.py | 5 +- src/mflux/generate.py | 4 +- src/mflux/generate_controlnet.py | 4 +- src/mflux/save.py | 4 +- src/mflux/ui/cli/parsers.py | 6 +- .../dreambooth/test_train_and_load_weights.py | 2 +- tests/image_generation/test_generate_image.py | 10 +- .../test_generate_image_controlnet.py | 6 +- tests/model_config/__init__.py | 0 tests/{ => model_config}/test_model_config.py | 27 ++-- tests/model_saving/test_model_saving.py | 4 +- tests/model_saving/test_model_saving_lora.py | 6 +- 20 files changed, 130 insertions(+), 129 deletions(-) create mode 100644 src/mflux/error/error.py create mode 100644 tests/model_config/__init__.py rename tests/{ => model_config}/test_model_config.py (68%) 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 043515d..78a71e9 100644 --- a/src/mflux/controlnet/flux_controlnet.py +++ b/src/mflux/controlnet/flux_controlnet.py @@ -52,7 +52,7 @@ class Flux1Controlnet(nn.Module): seed: int, prompt: str, controlnet_image_path: str, - config: Config = Config(), + config: Config, ) -> GeneratedImage: # 0. Create a new runtime config based on the model type and input parameters config = RuntimeConfig(config, self.model_config) 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 c5befb2..31c4f6b 100644 --- a/src/mflux/flux/flux.py +++ b/src/mflux/flux/flux.py @@ -4,7 +4,7 @@ from tqdm import tqdm from mflux.callbacks.callbacks import Callbacks from mflux.config.config import Config -from mflux.config.model_config import ModelConfig, ModelLookup +from mflux.config.model_config import ModelConfig from mflux.config.runtime_config import RuntimeConfig from mflux.flux.flux_initializer import FluxInitializer from mflux.latent_creator.latent_creator import Img2Img, LatentCreator @@ -47,7 +47,7 @@ class Flux1(nn.Module): self, seed: int, prompt: str, - config: Config = Config(), + config: Config, ) -> GeneratedImage: # 0. Create a new runtime config based on the model type and input parameters config = RuntimeConfig(config, self.model_config) @@ -139,7 +139,7 @@ class Flux1(nn.Module): @staticmethod def from_name(model_name: str, quantize: int | None = None) -> "Flux1": return Flux1( - model_config=ModelLookup.from_name(model_name=model_name, base_model=None), + model_config=ModelConfig.from_name(model_name=model_name, base_model=None), quantize=quantize, ) diff --git a/src/mflux/flux/flux_initializer.py b/src/mflux/flux/flux_initializer.py index 75a7c6e..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,7 +17,7 @@ class FluxInitializer: @staticmethod def init( flux_model, - model_config, + model_config: ModelConfig, quantize: int | None, local_path: str | None, lora_paths: list[str] | None, @@ -82,7 +83,7 @@ class FluxInitializer: @staticmethod def init_controlnet( flux_model, - model_config, + model_config: ModelConfig, quantize: int | None, local_path: str | None, lora_paths: list[str] | None, 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/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/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],