Let before-loop callback accept latents and let StepwiseHandler include initial latent

src/mflux/callbacks/callback.py

@@ -12,6 +12,8 @@ class BeforeLoopCallback(Protocol):
         self,
         seed: int,
         prompt: str,
+        latents: mx.array,
+        config: RuntimeConfig,
         canny_image: PIL.Image.Image | None = None,
     ) -> None:  # fmt: off
         ...
@@ -20,9 +22,9 @@ class BeforeLoopCallback(Protocol):
 class InLoopCallback(Protocol):
     def call_in_loop(
         self,
+        t: int,
         seed: int,
         prompt: str,
-        step: int,
         latents: mx.array,
         config: RuntimeConfig,
         time_steps: tqdm
@@ -33,9 +35,9 @@ class InLoopCallback(Protocol):
 class InterruptCallback(Protocol):
     def call_interrupt(
         self,
+        t: int,
         seed: int,
         prompt: str,
-        step: int,
         latents: mx.array,
         config: RuntimeConfig,
         time_steps: tqdm

src/mflux/callbacks/callbacks.py

@@ -11,29 +11,33 @@ class Callbacks:
     def before_loop(
         seed: int,
         prompt: str,
+        latents: mx.array,
+        config: RuntimeConfig,
         canny_image: PIL.Image.Image | None = None,
     ):  # fmt: off
         for subscriber in CallbackRegistry.before_loop_callbacks():
             subscriber.call_before_loop(
                 seed=seed,
                 prompt=prompt,
+                latents=latents,
+                config=config,
                 canny_image=canny_image,
             )
 
     @staticmethod
     def in_loop(
+        t: int,
         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(
+                t=t,
                 seed=seed,
                 prompt=prompt,
-                step=step,
                 latents=latents,
                 config=config,
                 time_steps=time_steps,
@@ -41,18 +45,18 @@ class Callbacks:
 
     @staticmethod
     def interruption(
+        t: int,
         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(
+                t=t,
                 seed=seed,
                 prompt=prompt,
-                step=step,
                 latents=latents,
                 config=config,
                 time_steps=time_steps,

src/mflux/callbacks/instances/canny_saver.py

@@ -1,10 +1,12 @@
 import os
 from pathlib import Path
 
+import mlx.core as mx
 import PIL.Image
 
 from mflux import ImageUtil
 from mflux.callbacks.callback import BeforeLoopCallback
+from mflux.config.runtime_config import RuntimeConfig
 
 
 class CannyImageSaver(BeforeLoopCallback):
@@ -15,8 +17,10 @@ class CannyImageSaver(BeforeLoopCallback):
         self,
         seed: int,
         prompt: str,
+        latents: mx.array,
+        config: RuntimeConfig,
         canny_image: PIL.Image.Image | None = None,
-    ) -> None:  # fmt: off
+    ) -> None:
         base, ext = os.path.splitext(self.path)
         ImageUtil.save_image(
             image=canny_image,

src/mflux/callbacks/instances/stepwise_handler.py

@@ -1,16 +1,17 @@
 from pathlib import Path
 
 import mlx.core as mx
+import PIL.Image
 import tqdm
 
 from mflux import StopImageGenerationException
-from mflux.callbacks.callback import InLoopCallback, InterruptCallback
+from mflux.callbacks.callback import BeforeLoopCallback, 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):
+class StepwiseHandler(BeforeLoopCallback, InLoopCallback, InterruptCallback):
     def __init__(
         self,
         flux,
@@ -23,17 +24,65 @@ class StepwiseHandler(InLoopCallback, InterruptCallback):
         if self.output_dir:
             self.output_dir.mkdir(parents=True, exist_ok=True)
 
-    def call_in_loop(
+    def call_before_loop(
         self,
         seed: int,
         prompt: str,
+        latents: mx.array,
+        config: RuntimeConfig,
+        canny_image: PIL.Image.Image | None = None,
+    ) -> None:  # fmt: off
+        self._save_image(
+            step=config.init_time_step,
+            seed=seed,
+            prompt=prompt,
+            latents=latents,
+            config=config,
+            time_steps=None,
+        )
+
+    def call_in_loop(
+        self,
+        t: int,
+        seed: int,
+        prompt: str,
+        latents: mx.array,
+        config: RuntimeConfig,
+        time_steps: tqdm
+    ) -> None:  # fmt: off
+        self._save_image(
+            step=t + 1,
+            seed=seed,
+            prompt=prompt,
+            latents=latents,
+            config=config,
+            time_steps=time_steps,
+        )
+
+    def call_interrupt(
+        self,
+        t: int,
+        seed: int,
+        prompt: str,
+        latents: mx.array,
+        config: RuntimeConfig,
+        time_steps: tqdm
+    ) -> None:  # fmt: off
+        self._save_composite(seed=seed)
+        raise StopImageGenerationException(f"Stopping image generation at step {t + 1}/{len(time_steps)}")
+
+    def _save_image(
+        self,
         step: int,
+        seed: int,
+        prompt: str,
         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)
+        generation_time = time_steps.format_dict["elapsed"] if time_steps is not None else 0
         stepwise_img = ImageUtil.to_image(
             decoded_latents=stepwise_decoded,
             config=config,
@@ -42,28 +91,15 @@ class StepwiseHandler(InLoopCallback, InterruptCallback):
             quantization=self.flux.bits,
             lora_paths=self.flux.lora_paths,
             lora_scales=self.flux.lora_scales,
-            generation_time=time_steps.format_dict["elapsed"],
+            generation_time=generation_time,
         )
-        self.step_wise_images.append(stepwise_img)
-
         stepwise_img.save(
-            path=self.output_dir / f"seed_{seed}_step{step}of{len(time_steps)}.png",
+            path=self.output_dir / f"seed_{seed}_step{step}of{config.num_inference_steps}.png",
             export_json_metadata=False,
         )
+        self.step_wise_images.append(stepwise_img)
         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)

src/mflux/controlnet/flux_controlnet.py

@@ -87,10 +87,12 @@ class Flux1Controlnet(nn.Module):
         Callbacks.before_loop(
             seed=seed,
             prompt=prompt,
+            latents=latents,
+            config=config,
             canny_image=canny_image
         )  # fmt: off
 
-        for gen_step, t in enumerate(time_steps, 1):
+        for t in time_steps:
             try:
                 # 4.t Compute controlnet samples
                 controlnet_block_samples, controlnet_single_block_samples = self.transformer_controlnet(
@@ -119,9 +121,9 @@ class Flux1Controlnet(nn.Module):
 
                 # (Optional) Call subscribes at end of loop
                 Callbacks.in_loop(
+                    t=t,
                     seed=seed,
                     prompt=prompt,
-                    step=gen_step,
                     latents=latents,
                     config=config,
                     time_steps=time_steps,
@@ -132,9 +134,9 @@ class Flux1Controlnet(nn.Module):
 
             except KeyboardInterrupt:  # noqa: PERF203
                 Callbacks.interruption(
+                    t=t,
                     seed=seed,
                     prompt=prompt,
-                    step=gen_step,
                     latents=latents,
                     config=config,
                     time_steps=time_steps,

src/mflux/flux/flux.py

@@ -79,10 +79,12 @@ class Flux1(nn.Module):
         # (Optional) Call subscribers for beginning of loop
         Callbacks.before_loop(
             seed=seed,
-            prompt=prompt
+            prompt=prompt,
+            latents=latents,
+            config=config,
         )  # fmt: off
 
-        for gen_step, t in enumerate(time_steps, 1):
+        for t in time_steps:
             try:
                 # 3.t Predict the noise
                 noise = self.transformer(
@@ -99,9 +101,9 @@ class Flux1(nn.Module):
 
                 # (Optional) Call subscribes at end of loop
                 Callbacks.in_loop(
+                    t=t,
                     seed=seed,
                     prompt=prompt,
-                    step=gen_step,
                     latents=latents,
                     config=config,
                     time_steps=time_steps,
@@ -112,9 +114,9 @@ class Flux1(nn.Module):
 
             except KeyboardInterrupt:  # noqa: PERF203
                 Callbacks.interruption(
+                    t=t,
                     seed=seed,
                     prompt=prompt,
-                    step=gen_step,
                     latents=latents,
                     config=config,
                     time_steps=time_steps,

src/mflux/generate.py

@@ -26,6 +26,7 @@ def main():
     # 2. Register the optional callbacks
     if args.stepwise_image_output_dir:
         handler = StepwiseHandler(flux=flux, output_dir=args.stepwise_image_output_dir)
+        CallbackRegistry.register_before_loop(handler)
         CallbackRegistry.register_in_loop(handler)
         CallbackRegistry.register_interrupt(handler)
 

src/mflux/generate_controlnet.py

@@ -29,6 +29,7 @@ def main():
         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_before_loop(handler)
         CallbackRegistry.register_in_loop(handler)
         CallbackRegistry.register_interrupt(handler)