Refactor: Remove ConfigControlnet

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

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

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
+            # For image-to-image: higher strength means we skip more steps.
-            # the higher the strength number, the more time steps we skip
             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:
+    def controlnet_strength(self) -> float | None:
-        if isinstance(self.config, ConfigControlnet):
+        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:

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:

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

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)

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)

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

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)
 

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,