Merge pull request #124 from filipstrand/refactor-and-add-callback-mechanism

Refactor and add callback mechanism

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/callbacks/callback.py

@@ -0,0 +1,43 @@
+from typing import Protocol
+
+import mlx.core as mx
+import PIL.Image
+import tqdm
+
+from mflux.config.runtime_config import RuntimeConfig
+
+
+class BeforeLoopCallback(Protocol):
+    def call_before_loop(
+        self,
+        seed: int,
+        prompt: str,
+        canny_image: PIL.Image.Image | None = None,
+    ) -> None:  # fmt: off
+        ...
+
+
+class InLoopCallback(Protocol):
+    def call_in_loop(
+        self,
+        seed: int,
+        prompt: str,
+        step: int,
+        latents: mx.array,
+        config: RuntimeConfig,
+        time_steps: tqdm
+    ) -> None:  # fmt: off
+        ...
+
+
+class InterruptCallback(Protocol):
+    def call_interrupt(
+        self,
+        seed: int,
+        prompt: str,
+        step: int,
+        latents: mx.array,
+        config: RuntimeConfig,
+        time_steps: tqdm
+    ) -> None:  # fmt: off
+        ...

src/mflux/callbacks/callback_registry.py

@@ -0,0 +1,31 @@
+from mflux.callbacks.callback import BeforeLoopCallback, InLoopCallback, InterruptCallback
+
+
+class CallbackRegistry:
+    in_loop = []
+    before_loop = []
+    interrupt = []
+
+    @staticmethod
+    def register_in_loop(callback: InLoopCallback) -> None:
+        CallbackRegistry.in_loop.append(callback)
+
+    @staticmethod
+    def register_before_loop(callback: BeforeLoopCallback) -> None:
+        CallbackRegistry.before_loop.append(callback)
+
+    @staticmethod
+    def register_interrupt(callback: InterruptCallback) -> None:
+        CallbackRegistry.interrupt.append(callback)
+
+    @staticmethod
+    def in_loop_callbacks() -> list[InLoopCallback]:
+        return CallbackRegistry.in_loop
+
+    @staticmethod
+    def before_loop_callbacks() -> list[BeforeLoopCallback]:
+        return CallbackRegistry.before_loop
+
+    @staticmethod
+    def interrupt_callbacks() -> list[InterruptCallback]:
+        return CallbackRegistry.interrupt

src/mflux/callbacks/callbacks.py

@@ -0,0 +1,59 @@
+import mlx.core as mx
+import PIL.Image
+import tqdm
+
+from mflux.callbacks.callback_registry import CallbackRegistry
+from mflux.config.runtime_config import RuntimeConfig
+
+
+class Callbacks:
+    @staticmethod
+    def before_loop(
+        seed: int,
+        prompt: str,
+        canny_image: PIL.Image.Image | None = None,
+    ):  # fmt: off
+        for subscriber in CallbackRegistry.before_loop_callbacks():
+            subscriber.call_before_loop(
+                seed=seed,
+                prompt=prompt,
+                canny_image=canny_image,
+            )
+
+    @staticmethod
+    def in_loop(
+        seed: int,
+        prompt: str,
+        step: int,
+        latents: mx.array,
+        config: RuntimeConfig,
+        time_steps: tqdm
+    ):  # fmt: off
+        for subscriber in CallbackRegistry.in_loop_callbacks():
+            subscriber.call_in_loop(
+                seed=seed,
+                prompt=prompt,
+                step=step,
+                latents=latents,
+                config=config,
+                time_steps=time_steps,
+            )
+
+    @staticmethod
+    def interruption(
+            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(
+                seed=seed,
+                prompt=prompt,
+                step=step,
+                latents=latents,
+                config=config,
+                time_steps=time_steps,
+            )

src/mflux/callbacks/instances/canny_saver.py

@@ -0,0 +1,24 @@
+import os
+from pathlib import Path
+
+import PIL.Image
+
+from mflux import ImageUtil
+from mflux.callbacks.callback import BeforeLoopCallback
+
+
+class CannyImageSaver(BeforeLoopCallback):
+    def __init__(self, path: str):
+        self.path = Path(path)
+
+    def call_before_loop(
+        self,
+        seed: int,
+        prompt: str,
+        canny_image: PIL.Image.Image | None = None,
+    ) -> None:  # fmt: off
+        base, ext = os.path.splitext(self.path)
+        ImageUtil.save_image(
+            image=canny_image,
+            path=f"{base}_controlnet_canny{ext}"
+        )  # fmt: off

src/mflux/callbacks/instances/stepwise_handler.py

@@ -0,0 +1,70 @@
+from pathlib import Path
+
+import mlx.core as mx
+import tqdm
+
+from mflux import StopImageGenerationException
+from mflux.callbacks.callback import InLoopCallback, InterruptCallback
+from mflux.config.runtime_config import RuntimeConfig
+from mflux.post_processing.array_util import ArrayUtil
+from mflux.post_processing.image_util import ImageUtil
+
+
+class StepwiseHandler(InLoopCallback, InterruptCallback):
+    def __init__(
+        self,
+        flux,
+        output_dir: str,
+    ):
+        self.flux = flux
+        self.output_dir = Path(output_dir)
+        self.step_wise_images = []
+
+        if self.output_dir:
+            self.output_dir.mkdir(parents=True, exist_ok=True)
+
+    def call_in_loop(
+        self,
+        seed: int,
+        prompt: str,
+        step: int,
+        latents: mx.array,
+        config: RuntimeConfig,
+        time_steps: tqdm
+    ) -> None:  # fmt: off
+        unpack_latents = ArrayUtil.unpack_latents(latents=latents, height=config.height, width=config.width)
+        stepwise_decoded = self.flux.vae.decode(unpack_latents)
+        stepwise_img = ImageUtil.to_image(
+            decoded_latents=stepwise_decoded,
+            config=config,
+            seed=seed,
+            prompt=prompt,
+            quantization=self.flux.bits,
+            lora_paths=self.flux.lora_paths,
+            lora_scales=self.flux.lora_scales,
+            generation_time=time_steps.format_dict["elapsed"],
+        )
+        self.step_wise_images.append(stepwise_img)
+
+        stepwise_img.save(
+            path=self.output_dir / f"seed_{seed}_step{step}of{len(time_steps)}.png",
+            export_json_metadata=False,
+        )
+        self._save_composite(seed=seed)
+
+    def call_interrupt(
+        self,
+        seed: int,
+        prompt: str,
+        step: int,
+        latents: mx.array,
+        config: RuntimeConfig,
+        time_steps: tqdm
+    ) -> None:  # fmt: off
+        self._save_composite(seed=seed)
+        raise StopImageGenerationException(f"Stopping image generation at step {step + 1}/{len(time_steps)}")
+
+    def _save_composite(self, seed: int) -> None:
+        if self.step_wise_images:
+            composite_img = ImageUtil.to_composite_image(self.step_wise_images)
+            composite_img.save(self.output_dir / f"seed_{seed}_composite.png")

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,14 +3,18 @@ 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__)
 
 
 class RuntimeConfig:
-    def __init__(self, config: Config | ConfigControlnet, model_config):
+    def __init__(
+        self,
+        config: Config,
+        model_config: ModelConfig,
+    ):
         self.config = config
         self.model_config = model_config
         self.sigmas = self._create_sigmas(config, model_config)
@@ -50,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/controlnet_util.py

@@ -1,16 +1,37 @@
 import logging
-import os
 
 import cv2
+import mlx.core as mx
 import numpy as np
 import PIL.Image
 
+from mflux.models.vae.vae import VAE
+from mflux.post_processing.array_util import ArrayUtil
+
 log = logging.getLogger(__name__)
 
 
 class ControlnetUtil:
     @staticmethod
-    def preprocess_canny(img: PIL.Image) -> PIL.Image:
+    def encode_image(
+        vae: VAE,
+        height: int,
+        width: int,
+        controlnet_image_path: str,
+    ) -> (mx.array, PIL.Image):
+        from mflux import ImageUtil
+
+        control_image = ImageUtil.load_image(controlnet_image_path)
+        control_image = ControlnetUtil._scale_image(height=height, width=width, img=control_image)
+        control_image = ControlnetUtil._preprocess_canny(control_image)
+        controlnet_cond = ImageUtil.to_array(control_image)
+        controlnet_cond = vae.encode(controlnet_cond)
+        controlnet_cond = (controlnet_cond / vae.scaling_factor) + vae.shift_factor
+        controlnet_cond = ArrayUtil.pack_latents(latents=controlnet_cond, height=height, width=width)
+        return controlnet_cond, control_image
+
+    @staticmethod
+    def _preprocess_canny(img: PIL.Image) -> PIL.Image:
         image_to_canny = np.array(img)
         image_to_canny = cv2.Canny(image_to_canny, 100, 200)
         image_to_canny = np.array(image_to_canny[:, :, None])
@@ -18,16 +39,8 @@ class ControlnetUtil:
         return PIL.Image.fromarray(image_to_canny)
 
     @staticmethod
-    def scale_image(height: int, width: int, img: PIL.Image) -> PIL.Image:
+    def _scale_image(height: int, width: int, img: PIL.Image) -> PIL.Image:
         if height != img.height or width != img.width:
             log.warning(f"Control image has different dimensions than the model. Resizing to {width}x{height}")
             img = img.resize((width, height), PIL.Image.LANCZOS)
         return img
-
-    @staticmethod
-    def save_canny_image(control_image: PIL.Image, path: str):
-        from mflux import ImageUtil
-
-        base, ext = os.path.splitext(path)
-        new_filename = f"{base}_controlnet_canny{ext}"
-        ImageUtil.save_image(control_image, new_filename)

src/mflux/controlnet/flux_controlnet.py

@@ -1,15 +1,14 @@
-from pathlib import Path
-
 import mlx.core as mx
+from mlx import nn
 from tqdm import tqdm
 
-from mflux.config.config import ConfigControlnet
+from mflux.callbacks.callbacks import Callbacks
+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
 from mflux.controlnet.transformer_controlnet import TransformerControlnet
-from mflux.controlnet.weight_handler_controlnet import WeightHandlerControlnet
+from mflux.flux.flux_initializer import FluxInitializer
-from mflux.error.exceptions import StopImageGenerationException
 from mflux.latent_creator.latent_creator import LatentCreator
 from mflux.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder
 from mflux.models.text_encoder.t5_encoder.t5_encoder import T5Encoder
@@ -18,17 +17,16 @@ from mflux.models.vae.vae import VAE
 from mflux.post_processing.array_util import ArrayUtil
 from mflux.post_processing.generated_image import GeneratedImage
 from mflux.post_processing.image_util import ImageUtil
-from mflux.post_processing.stepwise_handler import StepwiseHandler
-from mflux.tokenizer.clip_tokenizer import TokenizerCLIP
-from mflux.tokenizer.t5_tokenizer import TokenizerT5
-from mflux.tokenizer.tokenizer_handler import TokenizerHandler
 from mflux.weights.model_saver import ModelSaver
-from mflux.weights.weight_handler import WeightHandler
-from mflux.weights.weight_handler_lora import WeightHandlerLoRA
-from mflux.weights.weight_util import WeightUtil
 
 
-class Flux1Controlnet:
+class Flux1Controlnet(nn.Module):
+    vae: VAE
+    transformer: Transformer
+    transformer_controlnet: TransformerControlnet
+    t5_text_encoder: T5Encoder
+    clip_text_encoder: CLIPEncoder
+
     def __init__(
         self,
         model_config: ModelConfig,
@@ -38,84 +36,58 @@ class Flux1Controlnet:
         lora_scales: list[float] | None = None,
         controlnet_path: str | None = None,
     ):
-        self.lora_paths = lora_paths
+        super().__init__()
-        self.lora_scales = lora_scales
+        FluxInitializer.init_controlnet(
-        self.model_config = model_config
+            flux_model=self,
-
+            model_config=model_config,
-        # Load and initialize the tokenizers from disk, huggingface cache, or download from huggingface
+            quantize=quantize,
-        tokenizers = TokenizerHandler(model_config.model_name, self.model_config.max_sequence_length, local_path)
+            local_path=local_path,
-        self.t5_tokenizer = TokenizerT5(tokenizers.t5, max_length=self.model_config.max_sequence_length)
+            lora_paths=lora_paths,
-        self.clip_tokenizer = TokenizerCLIP(tokenizers.clip)
+            lora_scales=lora_scales,
-
-        # Load the weights
-        weights = WeightHandler.load_regular_weights(repo_id=model_config.model_name, local_path=local_path)
-
-        # Initialize the models
-        self.vae = VAE()
-        self.transformer = Transformer(model_config, num_transformer_blocks=weights.num_transformer_blocks(), num_single_transformer_blocks=weights.num_single_transformer_blocks())  # fmt: off
-        self.t5_text_encoder = T5Encoder()
-        self.clip_text_encoder = CLIPEncoder()
-
-        # Set the weights and quantize the model
-        self.bits = WeightUtil.set_weights_and_quantize(
-            quantize_arg=quantize,
-            weights=weights,
-            vae=self.vae,
-            transformer=self.transformer,
-            t5_text_encoder=self.t5_text_encoder,
-            clip_text_encoder=self.clip_text_encoder,
-        )
-
-        # Set LoRA weights
-        lora_weights = WeightHandlerLoRA.load_lora_weights(transformer=self.transformer, lora_files=lora_paths, lora_scales=lora_scales)  # fmt:off
-        WeightHandlerLoRA.set_lora_weights(transformer=self.transformer, loras=lora_weights)
-
-        # Set Controlnet weights
-        weights_controlnet = WeightHandlerControlnet.load_controlnet_transformer()
-        self.transformer_controlnet = TransformerControlnet(model_config=model_config, num_transformer_blocks=weights_controlnet.num_transformer_blocks(), num_single_transformer_blocks=weights_controlnet.num_single_transformer_blocks())  # fmt:off
-        WeightUtil.set_controlnet_weights_and_quantize(
-            quantize_arg=quantize,
-            weights=weights_controlnet,
-            transformer_controlnet=self.transformer_controlnet,
         )
 
     def generate_image(
         self,
         seed: int,
         prompt: str,
-        output: str,
         controlnet_image_path: str,
-        controlnet_save_canny: bool = False,
+        config: Config = Config(),
-        config: ConfigControlnet = ConfigControlnet(),
+    ) -> GeneratedImage:
-        stepwise_output_dir: Path = None,
+        # 0. Create a new runtime config based on the model type and input parameters
-    ) -> GeneratedImage:  # fmt: off
-        # Create a new runtime config based on the model type and input parameters
         config = RuntimeConfig(config, self.model_config)
         time_steps = tqdm(range(config.num_inference_steps))
-        stepwise_handler = StepwiseHandler(
+
-            flux=self,
+        # 1. Encode the controlnet reference image
-            config=config,
+        controlnet_condition, canny_image = ControlnetUtil.encode_image(
-            seed=seed,
+            vae=self.vae,
-            prompt=prompt,
+            height=config.height,
-            time_steps=time_steps,
+            width=config.width,
-            output_dir=stepwise_output_dir,
+            controlnet_image_path=controlnet_image_path,
         )
 
-        # 0. Embed the controlnet reference image
+        # 2. Create the initial latents
-        controlnet_condition = self._embed_image(config, controlnet_image_path, controlnet_save_canny, output)
+        latents = LatentCreator.create(
+            seed=seed,
+            height=config.height,
+            width=config.width
+        )  # fmt: off
 
-        # 1. Create the initial latents
+        # 3. Encode the prompt
-        latents = LatentCreator.create(seed=seed, height=config.height, width=config.width)
-
-        # 2. Embed 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)
 
+        # (Optional) Call subscribers for beginning of loop
+        Callbacks.before_loop(
+            seed=seed,
+            prompt=prompt,
+            canny_image=canny_image
+        )  # fmt: off
+
         for gen_step, t in enumerate(time_steps, 1):
             try:
-                # 3.t Compute controlnet samples
+                # 4.t Compute controlnet samples
                 controlnet_block_samples, controlnet_single_block_samples = self.transformer_controlnet(
                     t=t,
                     config=config,
@@ -125,7 +97,7 @@ class Flux1Controlnet:
                     controlnet_condition=controlnet_condition,
                 )
 
-                # 4.t Predict the noise
+                # 5.t Predict the noise
                 noise = self.transformer(
                     t=t,
                     config=config,
@@ -136,55 +108,48 @@ class Flux1Controlnet:
                     controlnet_single_block_samples=controlnet_single_block_samples,
                 )
 
-                # 5.t Take one denoise step
+                # 6.t Take one denoise step
                 dt = config.sigmas[t + 1] - config.sigmas[t]
                 latents += noise * dt
 
-                # Handle stepwise output if enabled
+                # (Optional) Call subscribes at end of loop
-                stepwise_handler.process_step(gen_step, latents)
+                Callbacks.in_loop(
+                    seed=seed,
+                    prompt=prompt,
+                    step=gen_step,
+                    latents=latents,
+                    config=config,
+                    time_steps=time_steps,
+                )  # fmt: off
 
-                # Evaluate to enable progress tracking
+                # (Optional) Evaluate to enable progress tracking
                 mx.eval(latents)
 
             except KeyboardInterrupt:  # noqa: PERF203
-                stepwise_handler.handle_interruption()
+                Callbacks.interruption(
-                raise StopImageGenerationException(f"Stopping image generation at step {t + 1}/{len(time_steps)}")
+                    seed=seed,
+                    prompt=prompt,
+                    step=gen_step,
+                    latents=latents,
+                    config=config,
+                    time_steps=time_steps,
+                )
 
-        # 5. Decode the latent array and return the image
+        # 7. Decode the latent array and return the image
         latents = ArrayUtil.unpack_latents(latents=latents, height=config.height, width=config.width)
         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 _embed_image(
-        self,
-        config: RuntimeConfig,
-        controlnet_image_path: str,
-        controlnet_save_canny: bool,
-        output: str,
-    ):
-        control_image = ImageUtil.load_image(controlnet_image_path)
-        control_image = ControlnetUtil.scale_image(config.height, config.width, control_image)
-        control_image = ControlnetUtil.preprocess_canny(control_image)
-
-        if controlnet_save_canny:
-            ControlnetUtil.save_canny_image(control_image, output)
-
-        controlnet_cond = ImageUtil.to_array(control_image)
-        controlnet_cond = self.vae.encode(controlnet_cond)
-        controlnet_cond = (controlnet_cond / self.vae.scaling_factor) + self.vae.shift_factor
-        controlnet_cond = ArrayUtil.pack_latents(latents=controlnet_cond, height=config.height, width=config.width)
-        return controlnet_cond
-
     def save_model(self, base_path: str) -> None:
         ModelSaver.save_model(self, self.bits, base_path)
         ModelSaver.save_weights(base_path, self.bits, self.transformer_controlnet, "transformer_controlnet")

src/mflux/flux/flux.py

@@ -1,15 +1,13 @@
-import warnings
-from pathlib import Path
-
 import mlx.core as mx
 from mlx import nn
 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.runtime_config import RuntimeConfig
-from mflux.error.exceptions import StopImageGenerationException
+from mflux.flux.flux_initializer import FluxInitializer
-from mflux.latent_creator.latent_creator import LatentCreator
+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.t5_encoder.t5_encoder import T5Encoder
 from mflux.models.transformer.transformer import Transformer
@@ -17,17 +15,15 @@ from mflux.models.vae.vae import VAE
 from mflux.post_processing.array_util import ArrayUtil
 from mflux.post_processing.generated_image import GeneratedImage
 from mflux.post_processing.image_util import ImageUtil
-from mflux.post_processing.stepwise_handler import StepwiseHandler
-from mflux.tokenizer.clip_tokenizer import TokenizerCLIP
-from mflux.tokenizer.t5_tokenizer import TokenizerT5
-from mflux.tokenizer.tokenizer_handler import TokenizerHandler
 from mflux.weights.model_saver import ModelSaver
-from mflux.weights.weight_handler import WeightHandler
-from mflux.weights.weight_handler_lora import WeightHandlerLoRA
-from mflux.weights.weight_util import WeightUtil
 
 
 class Flux1(nn.Module):
+    vae: VAE
+    transformer: Transformer
+    t5_text_encoder: T5Encoder
+    clip_text_encoder: CLIPEncoder
+
     def __init__(
         self,
         model_config: ModelConfig,
@@ -37,66 +33,50 @@ class Flux1(nn.Module):
         lora_scales: list[float] | None = None,
     ):
         super().__init__()
-        self.lora_paths = lora_paths
+        FluxInitializer.init(
-        self.lora_scales = lora_scales
+            flux_model=self,
-        self.model_config = model_config
+            model_config=model_config,
-
+            quantize=quantize,
-        # Load and initialize the tokenizers from disk, huggingface cache, or download from huggingface
+            local_path=local_path,
-        tokenizers = TokenizerHandler(model_config.model_name, self.model_config.max_sequence_length, local_path)
+            lora_paths=lora_paths,
-        self.t5_tokenizer = TokenizerT5(tokenizers.t5, max_length=self.model_config.max_sequence_length)
+            lora_scales=lora_scales,
-        self.clip_tokenizer = TokenizerCLIP(tokenizers.clip)
-
-        # Load the weights
-        weights = WeightHandler.load_regular_weights(repo_id=model_config.model_name, local_path=local_path)
-
-        # Initialize the models
-        self.vae = VAE()
-        self.transformer = Transformer(model_config, num_transformer_blocks=weights.num_transformer_blocks(), num_single_transformer_blocks=weights.num_single_transformer_blocks())  # fmt: off
-        self.t5_text_encoder = T5Encoder()
-        self.clip_text_encoder = CLIPEncoder()
-
-        # Set the weights and quantize the model
-        self.bits = WeightUtil.set_weights_and_quantize(
-            quantize_arg=quantize,
-            weights=weights,
-            vae=self.vae,
-            transformer=self.transformer,
-            t5_text_encoder=self.t5_text_encoder,
-            clip_text_encoder=self.clip_text_encoder,
         )
 
-        # Set LoRA weights
-        lora_weights = WeightHandlerLoRA.load_lora_weights(transformer=self.transformer, lora_files=lora_paths, lora_scales=lora_scales)  # fmt:off
-        WeightHandlerLoRA.set_lora_weights(transformer=self.transformer, loras=lora_weights)
-
     def generate_image(
         self,
         seed: int,
         prompt: str,
         config: Config = Config(),
-        stepwise_output_dir: Path = None,
     ) -> GeneratedImage:
-        # Create a new runtime config based on the model type and input parameters
+        # 0. Create a new runtime config based on the model type and input parameters
         config = RuntimeConfig(config, self.model_config)
         time_steps = tqdm(range(config.init_time_step, config.num_inference_steps))
-        stepwise_handler = StepwiseHandler(
-            flux=self,
-            config=config,
-            seed=seed,
-            prompt=prompt,
-            time_steps=time_steps,
-            output_dir=stepwise_output_dir,
-        )
 
         # 1. Create the initial latents
-        latents = LatentCreator.create_for_txt2img_or_img2img(seed, config, self.vae)
+        latents = LatentCreator.create_for_txt2img_or_img2img(
+            seed=seed,
+            height=config.height,
+            width=config.width,
+            img2img=Img2Img(
+                vae=self.vae,
+                sigmas=config.sigmas,
+                init_time_step=config.init_time_step,
+                init_image_path=config.init_image_path,
+            ),
+        )
 
-        # 2. Embed the prompt
+        # 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)
 
+        # (Optional) Call subscribers for beginning of loop
+        Callbacks.before_loop(
+            seed=seed,
+            prompt=prompt
+        )  # fmt: off
+
         for gen_step, t in enumerate(time_steps, 1):
             try:
                 # 3.t Predict the noise
@@ -112,41 +92,45 @@ class Flux1(nn.Module):
                 dt = config.sigmas[t + 1] - config.sigmas[t]
                 latents += noise * dt
 
-                # Handle stepwise output if enabled
+                # (Optional) Call subscribes at end of loop
-                stepwise_handler.process_step(gen_step, latents)
+                Callbacks.in_loop(
+                    seed=seed,
+                    prompt=prompt,
+                    step=gen_step,
+                    latents=latents,
+                    config=config,
+                    time_steps=time_steps,
+                )  # fmt: off
 
-                # Evaluate to enable progress tracking
+                # (Optional) Evaluate to enable progress tracking
                 mx.eval(latents)
 
             except KeyboardInterrupt:  # noqa: PERF203
-                stepwise_handler.handle_interruption()
+                Callbacks.interruption(
-                raise StopImageGenerationException(f"Stopping image generation at step {t + 1}/{len(time_steps)}")
+                    seed=seed,
+                    prompt=prompt,
+                    step=gen_step,
+                    latents=latents,
+                    config=config,
+                    time_steps=time_steps,
+                )
 
-        # 5. Decode the latent array and return the image
+        # 7. Decode the latent array and return the image
         latents = ArrayUtil.unpack_latents(latents=latents, height=config.height, width=config.width)
         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
-    def from_alias(alias: str, quantize: int | None = None) -> "Flux1":
-        warnings.warn(
-            "from_alias is deprecated and will be removed in a future release. Please use from_name instead.",
-            DeprecationWarning,
-            stacklevel=2,
-        )
-        return Flux1.from_name(model_name=alias, quantize=quantize)
-
     @staticmethod
     def from_name(model_name: str, quantize: int | None = None) -> "Flux1":
         return Flux1(

src/mflux/flux/flux_initializer.py

@@ -0,0 +1,111 @@
+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
+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
+from mflux.tokenizer.clip_tokenizer import TokenizerCLIP
+from mflux.tokenizer.t5_tokenizer import TokenizerT5
+from mflux.tokenizer.tokenizer_handler import TokenizerHandler
+from mflux.weights.weight_handler import WeightHandler
+from mflux.weights.weight_handler_lora import WeightHandlerLoRA
+from mflux.weights.weight_util import WeightUtil
+
+
+class FluxInitializer:
+    @staticmethod
+    def init(
+        flux_model,
+        model_config,
+        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
+        flux_model.lora_paths = lora_paths
+        flux_model.lora_scales = lora_scales
+        flux_model.model_config = model_config
+
+        # 1. Initialize tokenizers
+        tokenizers = TokenizerHandler(
+            repo_id=model_config.model_name,
+            max_t5_length=model_config.max_sequence_length,
+            local_path=local_path,
+        )
+        flux_model.t5_tokenizer = TokenizerT5(
+            tokenizer=tokenizers.t5,
+            max_length=model_config.max_sequence_length
+        )  # fmt: off
+        flux_model.clip_tokenizer = TokenizerCLIP(
+            tokenizer=tokenizers.clip,
+        )
+
+        # 2. Load the regular weights
+        weights = WeightHandler.load_regular_weights(
+            repo_id=model_config.model_name,
+            local_path=local_path
+        )  # fmt: off
+
+        # 3. Initialize all models
+        flux_model.vae = VAE()
+        flux_model.transformer = Transformer(
+            model_config=model_config,
+            num_transformer_blocks=weights.num_transformer_blocks(),
+            num_single_transformer_blocks=weights.num_single_transformer_blocks(),
+        )
+        flux_model.t5_text_encoder = T5Encoder()
+        flux_model.clip_text_encoder = CLIPEncoder()
+
+        # 4. Apply weights and quantize the models
+        flux_model.bits = WeightUtil.set_weights_and_quantize(
+            quantize_arg=quantize,
+            weights=weights,
+            vae=flux_model.vae,
+            transformer=flux_model.transformer,
+            t5_text_encoder=flux_model.t5_text_encoder,
+            clip_text_encoder=flux_model.clip_text_encoder,
+        )
+
+        # 5. Set LoRA weights
+        lora_weights = WeightHandlerLoRA.load_lora_weights(
+            transformer=flux_model.transformer,
+            lora_files=lora_paths,
+            lora_scales=lora_scales,
+        )
+        WeightHandlerLoRA.set_lora_weights(
+            transformer=flux_model.transformer,
+            loras=lora_weights
+        )  # fmt: off
+
+    @staticmethod
+    def init_controlnet(
+        flux_model,
+        model_config,
+        quantize: int | None,
+        local_path: str | None,
+        lora_paths: list[str] | None,
+        lora_scales: list[float] | None,
+    ) -> None:
+        # 1. Start with same init as regular Flux
+        FluxInitializer.init(
+            flux_model=flux_model,
+            model_config=model_config,
+            quantize=quantize,
+            local_path=local_path,
+            lora_paths=lora_paths,
+            lora_scales=lora_scales,
+        )
+
+        # 2. Apply ControlNet-specific initialization
+        weights_controlnet = WeightHandlerControlnet.load_controlnet_transformer()
+        flux_model.transformer_controlnet = TransformerControlnet(
+            model_config=model_config,
+            num_transformer_blocks=weights_controlnet.num_transformer_blocks(),
+            num_single_transformer_blocks=weights_controlnet.num_single_transformer_blocks(),
+        )
+        WeightUtil.set_controlnet_weights_and_quantize(
+            quantize_arg=quantize,
+            weights=weights_controlnet,
+            transformer_controlnet=flux_model.transformer_controlnet,
+        )

src/mflux/generate.py

@@ -1,11 +1,11 @@
-from pathlib import Path
-
 from mflux import Config, Flux1, ModelLookup, StopImageGenerationException
+from mflux.callbacks.callback_registry import CallbackRegistry
+from mflux.callbacks.instances.stepwise_handler import StepwiseHandler
 from mflux.ui.cli.parsers import CommandLineParser
 
 
 def main():
-    # fmt: off
+    # 0. Parse command line arguments
     parser = CommandLineParser(description="Generate an image based on a prompt.")
     parser.add_model_arguments(require_model_arg=False)
     parser.add_lora_arguments()
@@ -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,
@@ -23,13 +23,18 @@ def main():
         lora_scales=args.lora_scales,
     )
 
+    # 2. Register the optional callbacks
+    if args.stepwise_image_output_dir:
+        handler = StepwiseHandler(flux=flux, output_dir=args.stepwise_image_output_dir)
+        CallbackRegistry.register_in_loop(handler)
+        CallbackRegistry.register_interrupt(handler)
+
     try:
-        for seed_value in args.seed:
+        for seed in args.seed:
-            # Generate an image for each seed value
+            # 3. Generate an image for each seed value
             image = flux.generate_image(
-                seed=seed_value,
+                seed=seed,
                 prompt=args.prompt,
-                stepwise_output_dir=Path(args.stepwise_image_output_dir) if args.stepwise_image_output_dir else None,
                 config=Config(
                     num_inference_steps=args.steps,
                     height=args.height,
@@ -39,8 +44,8 @@ def main():
                     init_image_strength=args.init_image_strength,
                 ),
             )
-            # Save the image
+            # 4. Save the image
-            image.save(path=args.output.format(seed=seed_value), export_json_metadata=args.metadata)
+            image.save(path=args.output.format(seed=seed), export_json_metadata=args.metadata)
     except StopImageGenerationException as stop_exc:
         print(stop_exc)
 

src/mflux/generate_controlnet.py

@@ -1,10 +1,12 @@
-from pathlib import Path
+from mflux import Config, Flux1Controlnet, ModelLookup, StopImageGenerationException
-
+from mflux.callbacks.callback_registry import CallbackRegistry
-from mflux import ConfigControlnet, Flux1Controlnet, ModelLookup, StopImageGenerationException
+from mflux.callbacks.instances.canny_saver import CannyImageSaver
+from mflux.callbacks.instances.stepwise_handler import StepwiseHandler
 from mflux.ui.cli.parsers import CommandLineParser
 
 
 def main():
+    # 0. Parse command line arguments
     parser = CommandLineParser(description="Generate an image based on a prompt and a controlnet reference image.")  # fmt: off
     parser.add_model_arguments(require_model_arg=True)
     parser.add_lora_arguments()
@@ -13,7 +15,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,
@@ -22,17 +24,22 @@ def main():
         lora_scales=args.lora_scales,
     )
 
+    # 2. Register the optional callbacks
+    if args.controlnet_save_canny:
+        CallbackRegistry.register_before_loop(CannyImageSaver(path=args.output))
+    if args.stepwise_image_output_dir:
+        handler = StepwiseHandler(flux=flux, output_dir=args.stepwise_image_output_dir)
+        CallbackRegistry.register_in_loop(handler)
+        CallbackRegistry.register_interrupt(handler)
+
     try:
-        for seed_value in args.seed:
+        for seed in args.seed:
-            # Generate an image for each seed value
+            # 3. Generate an image for each seed value
             image = flux.generate_image(
-                seed=seed_value,
+                seed=seed,
                 prompt=args.prompt,
-                output=args.output,
                 controlnet_image_path=args.controlnet_image_path,
-                controlnet_save_canny=args.controlnet_save_canny,
+                config=Config(
-                stepwise_output_dir=Path(args.stepwise_image_output_dir) if args.stepwise_image_output_dir else None,
-                config=ConfigControlnet(
                     num_inference_steps=args.steps,
                     height=args.height,
                     width=args.width,
@@ -41,8 +48,8 @@ def main():
                 ),
             )
 
-            # Save the image
+            # 4. Save the image
-            image.save(path=args.output.format(seed=seed_value), export_json_metadata=args.metadata)
+            image.save(path=args.output.format(seed=seed), export_json_metadata=args.metadata)
     except StopImageGenerationException as stop_exc:
         print(stop_exc)
 

src/mflux/latent_creator/latent_creator.py

@@ -1,11 +1,24 @@
 import mlx.core as mx
-from mlx import nn
 
-from mflux.config.runtime_config import RuntimeConfig
+from mflux.models.vae.vae import VAE
 from mflux.post_processing.array_util import ArrayUtil
 from mflux.post_processing.image_util import ImageUtil
 
 
+class Img2Img:
+    def __init__(
+        self,
+        vae: VAE,
+        sigmas: mx.array,
+        init_time_step: int,
+        init_image_path: int,
+    ):
+        self.vae = vae
+        self.sigmas = sigmas
+        self.init_time_step = init_time_step
+        self.init_image_path = init_image_path
+
+
 class LatentCreator:
     @staticmethod
     def create(
@@ -21,29 +34,41 @@ class LatentCreator:
     @staticmethod
     def create_for_txt2img_or_img2img(
         seed: int,
-        runtime_conf: RuntimeConfig,
+        height: int,
-        vae: nn.Module,
+        width: int,
+        img2img: Img2Img,
     ) -> mx.array:
-        pure_noise = LatentCreator.create(
+        # 0. Determine type of image generation
-            seed=seed,
+        is_text2img = img2img.init_image_path is None
-            height=runtime_conf.height,
-            width=runtime_conf.width,
-        )
 
-        if runtime_conf.config.init_image_path is None:
+        if is_text2img:
-            # Text2Image
+            # 1. Create the pure noise
-            return pure_noise
+            return LatentCreator.create(
-        else:
+                seed=seed,
-            # Image2Image
+                height=height,
-            user_image = ImageUtil.load_image(runtime_conf.config.init_image_path).convert("RGB")
+                width=width,
-            scaled_user_image = ImageUtil.scale_to_dimensions(
-                image=user_image,
-                target_width=runtime_conf.width,
-                target_height=runtime_conf.height,
             )
-            encoded = vae.encode(ImageUtil.to_array(scaled_user_image))
+        else:
-            latents = ArrayUtil.pack_latents(latents=encoded, height=runtime_conf.height, width=runtime_conf.width)
+            # 1. Create the pure noise
-            sigma = runtime_conf.sigmas[runtime_conf.init_time_step]
+            pure_noise = LatentCreator.create(
+                seed=seed,
+                height=height,
+                width=width,
+            )
+
+            # 2. Encode the image
+            scaled_user_image = ImageUtil.scale_to_dimensions(
+                image=ImageUtil.load_image(img2img.init_image_path).convert("RGB"),
+                target_width=width,
+                target_height=height,
+            )
+            encoded = img2img.vae.encode(ImageUtil.to_array(scaled_user_image))
+            latents = ArrayUtil.pack_latents(latents=encoded, height=height, width=width)
+
+            # 3. Find the appropriate sigma value
+            sigma = img2img.sigmas[img2img.init_time_step]
+
+            # 4. Blend the appropriate amount of noise based on linear interpolation
             return LatentCreator.add_noise_by_interpolation(
                 clean=latents,
                 noise=pure_noise,

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

@@ -1,60 +0,0 @@
-from pathlib import Path
-
-import mlx.core as mx
-import tqdm
-
-from mflux.config.runtime_config import RuntimeConfig
-from mflux.post_processing.array_util import ArrayUtil
-from mflux.post_processing.image_util import ImageUtil
-
-
-class StepwiseHandler:
-    def __init__(
-        self,
-        flux,
-        config: RuntimeConfig,
-        seed: int,
-        prompt: str,
-        time_steps: tqdm.std.tqdm,
-        output_dir: Path | None = None,
-    ):
-        self.flux = flux
-        self.config = config
-        self.seed = seed
-        self.prompt = prompt
-        self.output_dir = output_dir
-        self.time_steps = time_steps
-        self.step_wise_images = []
-
-        if self.output_dir:
-            self.output_dir.mkdir(parents=True, exist_ok=True)
-
-    def save_composite(self):
-        if self.step_wise_images:
-            composite_img = ImageUtil.to_composite_image(self.step_wise_images)
-            composite_img.save(self.output_dir / f"seed_{self.seed}_composite.png")
-
-    def process_step(self, gen_step: int, latents: mx.array):
-        if self.output_dir:
-            unpack_latents = ArrayUtil.unpack_latents(latents=latents, height=self.config.height, width=self.config.width)  # fmt: off
-            stepwise_decoded = self.flux.vae.decode(unpack_latents)
-            stepwise_img = ImageUtil.to_image(
-                decoded_latents=stepwise_decoded,
-                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,
-            )
-            self.step_wise_images.append(stepwise_img)
-
-            stepwise_img.save(
-                path=self.output_dir / f"seed_{self.seed}_step{gen_step}of{len(self.time_steps)}.png",
-                export_json_metadata=False,
-            )
-            self.save_composite()
-
-    def handle_interruption(self):
-        self.save_composite()

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
 
 
@@ -40,10 +40,8 @@ class ImageGeneratorControlnetTestHelper:
             image = flux.generate_image(
                 seed=seed,
                 prompt=prompt,
-                output=str(output_image_path),
                 controlnet_image_path=controlnet_image_path,
-                controlnet_save_canny=False,
+                config=Config(
-                config=ConfigControlnet(
                     num_inference_steps=steps,
                     height=768,
                     width=493,