monster/Qwen-Image-Layered-MRP-MLX
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refactor model config

Browse Source
This commit is contained in:
filipstrand 2025-02-12 22:27:12 +01:00
parent f21710eb7e
commit 3cecfae75b
20 changed files with 130 additions and 129 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
src/mflux/__init__.py
View File

@ -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",
]

12
src/mflux/callbacks/callbacks.py
View File

@ -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(

144
src/mflux/config/model_config.py
View File

@ -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"

4
src/mflux/config/runtime_config.py
View File

@ -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

2
src/mflux/controlnet/flux_controlnet.py
View File

@ -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)

4
src/mflux/dreambooth/dreambooth_initializer.py
View File

@ -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,

6
src/mflux/error/error.py Normal file
View File

@ -0,0 +1,6 @@
class ModelConfigError(ValueError):
"""User error in model config."""
class InvalidBaseModel(ModelConfigError):
"""Invalid base model, cannot infer model properties."""

6
src/mflux/flux/flux.py
View File

@ -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,
)

5
src/mflux/flux/flux_initializer.py
View File

@ -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,

4
src/mflux/generate.py
View File

@ -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,

4
src/mflux/generate_controlnet.py
View File

@ -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,

4
src/mflux/save.py
View File

@ -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,

6
src/mflux/ui/cli/parsers.py
View File

@ -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

2
tests/dreambooth/test_train_and_load_weights.py
View File

@ -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],

10
tests/image_generation/test_generate_image.py
View File

@ -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,

6
tests/image_generation/test_generate_image_controlnet.py
View File

@ -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",

0
tests/model_config/__init__.py Normal file
View File

27
tests/test_model_config.py → tests/model_config/test_model_config.py
View File

@ -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",
)

4
tests/model_saving/test_model_saving.py
View File

@ -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,
)

6
tests/model_saving/test_model_saving_lora.py
View File

@ -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],