diff --git a/src/mflux/__init__.py b/src/mflux/__init__.py index 83273fc..67b368f 100644 --- a/src/mflux/__init__.py +++ b/src/mflux/__init__.py @@ -1,4 +1,4 @@ -from mflux.config.config import Config, ConfigControlnet +from mflux.config.config import Config from mflux.config.model_config import ModelConfig, ModelLookup from mflux.controlnet.flux_controlnet import Flux1Controlnet from mflux.error.exceptions import StopImageGenerationException @@ -9,7 +9,6 @@ __all__ = [ "Flux1", "Flux1Controlnet", "Config", - "ConfigControlnet", "ModelConfig", "ModelLookup", "ImageUtil", diff --git a/src/mflux/config/config.py b/src/mflux/config/config.py index d2b0584..c77147c 100644 --- a/src/mflux/config/config.py +++ b/src/mflux/config/config.py @@ -17,6 +17,7 @@ class Config: guidance: float = 4.0, init_image_path: Path | None = None, init_image_strength: float | None = None, + controlnet_strength: float | None = None, ): if width % 16 != 0 or height % 16 != 0: log.warning("Width and height should be multiples of 16. Rounding down.") @@ -26,16 +27,4 @@ class Config: self.guidance = guidance self.init_image_path = init_image_path self.init_image_strength = init_image_strength - - -class ConfigControlnet(Config): - def __init__( - self, - num_inference_steps: int = 4, - width: int = 1024, - height: int = 1024, - guidance: float = 4.0, - controlnet_strength: float = 1.0, - ): - super().__init__(num_inference_steps=num_inference_steps, width=width, height=height, guidance=guidance) self.controlnet_strength = controlnet_strength diff --git a/src/mflux/config/runtime_config.py b/src/mflux/config/runtime_config.py index 3c0cbc1..096338d 100644 --- a/src/mflux/config/runtime_config.py +++ b/src/mflux/config/runtime_config.py @@ -3,7 +3,7 @@ import logging import mlx.core as mx import numpy as np -from mflux.config.config import Config, ConfigControlnet +from mflux.config.config import Config from mflux.config.model_config import ModelConfig logger = logging.getLogger(__name__) @@ -12,7 +12,7 @@ logger = logging.getLogger(__name__) class RuntimeConfig: def __init__( self, - config: Config | ConfigControlnet, + config: Config, model_config: ModelConfig, ): self.config = config @@ -54,23 +54,20 @@ class RuntimeConfig: @property def init_time_step(self) -> int: if self.config.init_image_path is None: - # text to image, always begin at time step 0 + # For text-to-image, always begin at time step 0. return 0 else: - # we skip to the time step as informed by the init_image_strength - # the higher the strength number, the more time steps we skip + # For image-to-image: higher strength means we skip more steps. strength = max(0.0, min(1.0, self.config.init_image_strength)) - # if the strength is too small to even influence the image - # help the user round up so the init_image has influence at step 1 t = max(1, int(self.num_inference_steps * strength)) return t @property - def controlnet_strength(self) -> float: - if isinstance(self.config, ConfigControlnet): + def controlnet_strength(self) -> float | None: + if self.config.controlnet_strength is not None: return self.config.controlnet_strength - else: - raise NotImplementedError("Controlnet conditioning scale is only available for ConfigControlnet") + + return None @staticmethod def _create_sigmas(config, model) -> mx.array: diff --git a/src/mflux/controlnet/flux_controlnet.py b/src/mflux/controlnet/flux_controlnet.py index 2642739..41c9b03 100644 --- a/src/mflux/controlnet/flux_controlnet.py +++ b/src/mflux/controlnet/flux_controlnet.py @@ -4,7 +4,7 @@ import mlx.core as mx from mlx import nn from tqdm import tqdm -from mflux.config.config import ConfigControlnet +from mflux.config.config import Config from mflux.config.model_config import ModelConfig from mflux.config.runtime_config import RuntimeConfig from mflux.controlnet.controlnet_util import ControlnetUtil @@ -56,10 +56,10 @@ class Flux1Controlnet(nn.Module): output: str, controlnet_image_path: str, controlnet_save_canny: bool = False, - config: ConfigControlnet = ConfigControlnet(), + config: Config = Config(), stepwise_output_dir: Path = None, ) -> GeneratedImage: - # Create a new runtime config based on the model type and input parameters + # Convert the user config to a runtime config with derived parameters. config = RuntimeConfig(config, self.model_config) time_steps = tqdm(range(config.num_inference_steps)) stepwise_handler = StepwiseHandler( @@ -135,14 +135,14 @@ class Flux1Controlnet(nn.Module): decoded = self.vae.decode(latents) return ImageUtil.to_image( decoded_latents=decoded, + config=config, seed=seed, prompt=prompt, quantization=self.bits, - generation_time=time_steps.format_dict["elapsed"], lora_paths=self.lora_paths, lora_scales=self.lora_scales, - config=config, controlnet_image_path=controlnet_image_path, + generation_time=time_steps.format_dict["elapsed"], ) def save_model(self, base_path: str) -> None: diff --git a/src/mflux/flux/flux.py b/src/mflux/flux/flux.py index 5dc6090..15ed094 100644 --- a/src/mflux/flux/flux.py +++ b/src/mflux/flux/flux.py @@ -53,7 +53,7 @@ class Flux1(nn.Module): config: Config = Config(), stepwise_output_dir: Path = None, ) -> GeneratedImage: - # Create a new runtime config based on the model type and input parameters + # Convert the user config to a runtime config with derived parameters. config = RuntimeConfig(config, self.model_config) time_steps = tqdm(range(config.init_time_step, config.num_inference_steps)) stepwise_handler = StepwiseHandler( @@ -108,15 +108,15 @@ class Flux1(nn.Module): decoded = self.vae.decode(latents) return ImageUtil.to_image( decoded_latents=decoded, + config=config, seed=seed, prompt=prompt, quantization=self.bits, - generation_time=time_steps.format_dict["elapsed"], lora_paths=self.lora_paths, lora_scales=self.lora_scales, init_image_path=config.init_image_path, init_image_strength=config.init_image_strength, - config=config, + generation_time=time_steps.format_dict["elapsed"], ) @staticmethod diff --git a/src/mflux/generate.py b/src/mflux/generate.py index 653380e..72f5a49 100644 --- a/src/mflux/generate.py +++ b/src/mflux/generate.py @@ -14,7 +14,7 @@ def main(): parser.add_output_arguments() args = parser.parse_args() - # Load the model + # 1. Load the model flux = Flux1( model_config=ModelLookup.from_name(model_name=args.model, base_model=args.base_model), quantize=args.quantize, @@ -25,7 +25,7 @@ def main(): try: for seed_value in args.seed: - # Generate an image for each seed value + # 2. Generate an image for each seed value image = flux.generate_image( seed=seed_value, prompt=args.prompt, @@ -39,7 +39,7 @@ def main(): init_image_strength=args.init_image_strength, ), ) - # Save the image + # 3. Save the image image.save(path=args.output.format(seed=seed_value), export_json_metadata=args.metadata) except StopImageGenerationException as stop_exc: print(stop_exc) diff --git a/src/mflux/generate_controlnet.py b/src/mflux/generate_controlnet.py index 03e3b83..ccbf394 100644 --- a/src/mflux/generate_controlnet.py +++ b/src/mflux/generate_controlnet.py @@ -1,6 +1,6 @@ from pathlib import Path -from mflux import ConfigControlnet, Flux1Controlnet, ModelLookup, StopImageGenerationException +from mflux import Config, Flux1Controlnet, ModelLookup, StopImageGenerationException from mflux.ui.cli.parsers import CommandLineParser @@ -13,7 +13,7 @@ def main(): parser.add_output_arguments() args = parser.parse_args() - # Load the model + # 1. Load the model flux = Flux1Controlnet( model_config=ModelLookup.from_name(model_name=args.model, base_model=args.base_model), quantize=args.quantize, @@ -24,7 +24,7 @@ def main(): try: for seed_value in args.seed: - # Generate an image for each seed value + # 2. Generate an image for each seed value image = flux.generate_image( seed=seed_value, prompt=args.prompt, @@ -32,7 +32,7 @@ def main(): controlnet_image_path=args.controlnet_image_path, controlnet_save_canny=args.controlnet_save_canny, stepwise_output_dir=Path(args.stepwise_image_output_dir) if args.stepwise_image_output_dir else None, - config=ConfigControlnet( + config=Config( num_inference_steps=args.steps, height=args.height, width=args.width, @@ -41,7 +41,7 @@ def main(): ), ) - # Save the image + # 3. Save the image image.save(path=args.output.format(seed=seed_value), export_json_metadata=args.metadata) except StopImageGenerationException as stop_exc: print(stop_exc) diff --git a/src/mflux/post_processing/image_util.py b/src/mflux/post_processing/image_util.py index c7ccca4..a21da15 100644 --- a/src/mflux/post_processing/image_util.py +++ b/src/mflux/post_processing/image_util.py @@ -8,25 +8,23 @@ import numpy as np import piexif import PIL.Image -from mflux.config.config import ConfigControlnet +from mflux.config.runtime_config import RuntimeConfig from mflux.post_processing.generated_image import GeneratedImage log = logging.getLogger(__name__) -RuntimeConfig: t.TypeAlias = "mflux.config.runtime_config.RuntimeConfig" # noqa: F821 - class ImageUtil: @staticmethod def to_image( decoded_latents: mx.array, + config: RuntimeConfig, seed: int, prompt: str, quantization: int, generation_time: float, lora_paths: list[str], lora_scales: list[float], - config: RuntimeConfig, controlnet_image_path: str | None = None, init_image_path: str | None = None, init_image_strength: float | None = None, @@ -49,7 +47,7 @@ class ImageUtil: init_image_path=init_image_path, init_image_strength=init_image_strength, controlnet_image_path=controlnet_image_path, - controlnet_strength=config.controlnet_strength if isinstance(config.config, ConfigControlnet) else None, + controlnet_strength=config.controlnet_strength, ) @staticmethod diff --git a/src/mflux/post_processing/stepwise_handler.py b/src/mflux/post_processing/stepwise_handler.py index 60332bb..59dcc25 100644 --- a/src/mflux/post_processing/stepwise_handler.py +++ b/src/mflux/post_processing/stepwise_handler.py @@ -40,13 +40,13 @@ class StepwiseHandler: stepwise_decoded = self.flux.vae.decode(unpack_latents) stepwise_img = ImageUtil.to_image( decoded_latents=stepwise_decoded, + config=self.config, seed=self.seed, prompt=self.prompt, quantization=self.flux.bits, - generation_time=self.time_steps.format_dict["elapsed"], lora_paths=self.flux.lora_paths, lora_scales=self.flux.lora_scales, - config=self.config, + generation_time=self.time_steps.format_dict["elapsed"], ) self.step_wise_images.append(stepwise_img) diff --git a/tests/image_generation/helpers/image_generation_controlnet_test_helper.py b/tests/image_generation/helpers/image_generation_controlnet_test_helper.py index 9a47cd3..e2794a4 100644 --- a/tests/image_generation/helpers/image_generation_controlnet_test_helper.py +++ b/tests/image_generation/helpers/image_generation_controlnet_test_helper.py @@ -3,7 +3,7 @@ import os import numpy as np from PIL import Image -from mflux import ConfigControlnet, Flux1Controlnet, ModelConfig +from mflux import Config, Flux1Controlnet, ModelConfig from tests.image_generation.helpers.image_generation_test_helper import ImageGeneratorTestHelper @@ -43,7 +43,7 @@ class ImageGeneratorControlnetTestHelper: output=str(output_image_path), controlnet_image_path=controlnet_image_path, controlnet_save_canny=False, - config=ConfigControlnet( + config=Config( num_inference_steps=steps, height=768, width=493,