Implement general callback mechanism

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/controlnet/controlnet_util.py

@@ -1,12 +1,10 @@
 import logging
-import os
 
 import cv2
 import mlx.core as mx
 import numpy as np
 import PIL.Image
 
-from mflux.config.runtime_config import RuntimeConfig
 from mflux.models.vae.vae import VAE
 from mflux.post_processing.array_util import ArrayUtil
 
@@ -17,29 +15,20 @@ class ControlnetUtil:
     @staticmethod
     def encode_image(
         vae: VAE,
-        config: RuntimeConfig,
+        height: int,
+        width: int,
         controlnet_image_path: str,
-        controlnet_save_canny: bool,
+    ) -> (mx.array, PIL.Image):
-        output: str,
-    ) -> mx.array:
         from mflux import ImageUtil
 
         control_image = ImageUtil.load_image(controlnet_image_path)
-        control_image = ControlnetUtil._scale_image(config.height, config.width, control_image)
+        control_image = ControlnetUtil._scale_image(height=height, width=width, img=control_image)
         control_image = ControlnetUtil._preprocess_canny(control_image)
-
-        if controlnet_save_canny:
-            base, ext = os.path.splitext(output)
-            ImageUtil.save_image(
-                image=control_image,
-                path=f"{base}_controlnet_canny{ext}"
-            )  # fmt: off
-
         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=config.height, width=config.width)
+        controlnet_cond = ArrayUtil.pack_latents(latents=controlnet_cond, height=height, width=width)
-        return controlnet_cond
+        return controlnet_cond, control_image
 
     @staticmethod
     def _preprocess_canny(img: PIL.Image) -> PIL.Image:

src/mflux/controlnet/flux_controlnet.py

@@ -1,15 +1,13 @@
-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
 from mflux.config.runtime_config import RuntimeConfig
 from mflux.controlnet.controlnet_util import ControlnetUtil
 from mflux.controlnet.transformer_controlnet import TransformerControlnet
-from mflux.error.exceptions import StopImageGenerationException
 from mflux.flux.flux_initializer import FluxInitializer
 from mflux.latent_creator.latent_creator import LatentCreator
 from mflux.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder
@@ -19,7 +17,6 @@ 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.weights.model_saver import ModelSaver
 
 
@@ -53,49 +50,44 @@ class Flux1Controlnet(nn.Module):
         self,
         seed: int,
         prompt: str,
-        output: str,
         controlnet_image_path: str,
-        controlnet_save_canny: bool = False,
         config: Config = Config(),
-        stepwise_output_dir: Path = None,
     ) -> GeneratedImage:
-        # Convert the user config to a runtime config with derived 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.num_inference_steps))
-        stepwise_handler = StepwiseHandler(
-            flux=self,
-            config=config,
-            seed=seed,
-            prompt=prompt,
-            time_steps=time_steps,
-            output_dir=stepwise_output_dir,
-        )
 
-        # 0. Encode the controlnet reference image
+        # 1. Encode the controlnet reference image
-        controlnet_condition = ControlnetUtil.encode_image(
+        controlnet_condition, canny_image = ControlnetUtil.encode_image(
             vae=self.vae,
-            config=config,
+            height=config.height,
+            width=config.width,
             controlnet_image_path=controlnet_image_path,
-            controlnet_save_canny=controlnet_save_canny,
-            output=output,
         )
 
-        # 1. Create the initial latents
+        # 2. Create the initial latents
         latents = LatentCreator.create(
             seed=seed,
             height=config.height,
             width=config.width
         )  # fmt: off
 
-        # 2. Embed the prompt
+        # 3. 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,
+            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,
@@ -105,7 +97,7 @@ class Flux1Controlnet(nn.Module):
                     controlnet_condition=controlnet_condition,
                 )
 
-                # 4.t Predict the noise
+                # 5.t Predict the noise
                 noise = self.transformer(
                     t=t,
                     config=config,
@@ -116,21 +108,34 @@ class Flux1Controlnet(nn.Module):
                     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(

src/mflux/flux/flux.py

@@ -1,16 +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
@@ -18,7 +15,6 @@ 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.weights.model_saver import ModelSaver
 
 
@@ -51,33 +47,36 @@ class Flux1(nn.Module):
         seed: int,
         prompt: str,
         config: Config = Config(),
-        stepwise_output_dir: Path = None,
     ) -> GeneratedImage:
-        # Convert the user config to a runtime config with derived 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=seed,
-            vae=self.vae,
+            height=config.height,
-            runtime_conf=config,
+            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
@@ -93,17 +92,30 @@ 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(
@@ -119,15 +131,6 @@ class Flux1(nn.Module):
             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/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()
@@ -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:
-            # 2. 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,
                 ),
             )
-            # 3. 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.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()
@@ -22,16 +24,21 @@ 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:
-            # 2. 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,
-                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,
@@ -41,8 +48,8 @@ def main():
                 ),
             )
 
-            # 3. 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/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,
-                config=self.config,
-                seed=self.seed,
-                prompt=self.prompt,
-                quantization=self.flux.bits,
-                lora_paths=self.flux.lora_paths,
-                lora_scales=self.flux.lora_scales,
-                generation_time=self.time_steps.format_dict["elapsed"],
-            )
-            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

@@ -40,9 +40,7 @@ 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(
                     num_inference_steps=steps,
                     height=768,