Add stepwise handler

src/mflux/__init__.py

@@ -3,7 +3,7 @@ from mflux.config.config import ConfigControlnet
 from mflux.config.model_config import ModelConfig
 from mflux.controlnet.flux_controlnet import Flux1Controlnet
 from mflux.flux.flux import Flux1
-from mflux.exceptions import StopImageGenerationException
+from mflux.error.exceptions import StopImageGenerationException
 from mflux.post_processing.image_util import ImageUtil
 
 __all__ = [

src/mflux/controlnet/flux_controlnet.py

@@ -10,12 +10,14 @@ 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.exceptions import StopImageGenerationException
+from mflux.error.exceptions import StopImageGenerationException
 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.post_processing.array_util import ArrayUtil
 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
@@ -112,11 +114,18 @@ class Flux1Controlnet:
             controlnet_save_canny: bool = False,
             config: ConfigControlnet = ConfigControlnet(),
             stepwise_output_dir: Path = None,
-            stepwise_composite_only=False
     ) -> "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,
+            config=config,
+            seed=seed,
+            prompt=prompt,
+            time_steps=time_steps,
+            output_dir=stepwise_output_dir,
+        )
 
         # Embedd the controlnet reference image
         control_image = ImageUtil.load_image(controlnet_image_path)
@@ -127,7 +136,7 @@ class Flux1Controlnet:
         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 = Flux1Controlnet._pack_latents(controlnet_cond, config.height, config.width)
+        controlnet_cond = ArrayUtil.pack_latents(controlnet_cond, config.height, config.width)
 
         # 1. Create the initial latents
         latents = mx.random.normal(
@@ -135,15 +144,12 @@ class Flux1Controlnet:
             key=mx.random.key(seed)
         )  # fmt: off
 
-        # 2. Embedd the prompt
+        # 2. Embed the prompt
         t5_tokens = self.t5_tokenizer.tokenize(prompt)
         clip_tokens = self.clip_tokenizer.tokenize(prompt)
         prompt_embeds = self.t5_text_encoder.forward(t5_tokens)
         pooled_prompt_embeds = self.clip_text_encoder.forward(clip_tokens)
 
-        step_wise_images = []
-        if stepwise_output_dir:
-            stepwise_output_dir.mkdir(parents=True, exist_ok=True)
         for t in time_steps:
             try:
                 # Compute controlnet samples
@@ -171,41 +177,18 @@ class Flux1Controlnet:
                 dt = config.sigmas[t + 1] - config.sigmas[t]
                 latents += noise * dt
 
+                # Handle stepwise output if enabled
+                stepwise_handler.process_step(t, latents)
+
                 # Evaluate to enable progress tracking
                 mx.eval(latents)
 
-                if stepwise_output_dir:
+            except KeyboardInterrupt:
-                    stepwise_decoded = self.vae.decode(
+                stepwise_handler.handle_interruption()
-                        Flux1Controlnet._unpack_latents(latents, config.height, config.width)
+                raise StopImageGenerationException(f"Stopping image generation at step {t + 1}/{len(time_steps)}")
-                    )
-                    # performance todo: Pillow mostly uses CPU,
-                    # can try to improve image generation performance by offloading
-                    # stepwise image processing to the CPU via threading
-                    stepwise_img = ImageUtil.to_image(
-                        decoded_latents=stepwise_decoded,
-                        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,
-                    )
-                    step_wise_images.append(stepwise_img)
-                    if not stepwise_composite_only:
-                        stepwise_img.save(
-                            path=stepwise_output_dir / f"seed_{seed}_step{t+1}of{len(time_steps)}.png",
-                            export_json_metadata=False,
-                        )
-            except KeyboardInterrupt:  # noqa: PERF203
-                raise StopImageGenerationException(f"Stopping image generation at step {t+1}/{len(time_steps)}")
-            finally:
-                if step_wise_images and stepwise_output_dir:
-                    composite_img = ImageUtil.to_composite_image(step_wise_images)
-                    composite_img.save(stepwise_output_dir / f"seed_{seed}_composite.png")
 
         # 5. Decode the latent array and return the image
-        latents = Flux1Controlnet._unpack_latents(latents, config.height, config.width)
+        latents = ArrayUtil.unpack_latents(latents, config.height, config.width)
         decoded = self.vae.decode(latents)
         return ImageUtil.to_image(
             decoded_latents=decoded,
@@ -219,20 +202,6 @@ class Flux1Controlnet:
             controlnet_image_path=controlnet_image_path,
         )
 
-    @staticmethod
-    def _unpack_latents(latents: mx.array, height: int, width: int) -> mx.array:
-        latents = mx.reshape(latents, (1, height // 16, width // 16, 16, 2, 2))
-        latents = mx.transpose(latents, (0, 3, 1, 4, 2, 5))
-        latents = mx.reshape(latents, (1, 16, height // 16 * 2, width // 16 * 2))
-        return latents
-
-    @staticmethod
-    def _pack_latents(latents: mx.array, height: int, width: int) -> mx.array:
-        latents = mx.reshape(latents, (1, 16, height // 16, 2, width // 16, 2))
-        latents = mx.transpose(latents, (0, 2, 4, 1, 3, 5))
-        latents = mx.reshape(latents, (1, (width // 16) * (height // 16), 64))
-        return latents
-
     def _set_model_weights(self, weights):
         self.vae.update(weights.vae)
         self.transformer.update(weights.transformer)

src/mflux/flux/flux.py

@@ -1,16 +1,19 @@
-import mlx.core as mx
 from pathlib import Path
+
+import mlx.core as mx
 from mlx import nn
 from tqdm import tqdm
 
 from mflux.config.config import Config
 from mflux.config.model_config import ModelConfig
 from mflux.config.runtime_config import RuntimeConfig
-from mflux.exceptions import StopImageGenerationException
+from mflux.error.exceptions import StopImageGenerationException
+from mflux.post_processing.stepwise_handler import StepwiseHandler
 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.post_processing.array_util import ArrayUtil
 from mflux.post_processing.generated_image import GeneratedImage
 from mflux.post_processing.image_util import ImageUtil
 from mflux.tokenizer.clip_tokenizer import TokenizerCLIP
@@ -71,10 +74,24 @@ class Flux1:
         if weights.quantization_level is not None:
             self._set_model_weights(weights)
 
-    def generate_image(self, seed: int, prompt: str, config: Config = Config(), stepwise_output_dir: Path = None, stepwise_composite_only=False) -> GeneratedImage:  # fmt: off
+    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
         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,
+        )
 
         # 1. Create the initial latents
         latents = mx.random.normal(
@@ -82,15 +99,12 @@ class Flux1:
             key=mx.random.key(seed)
         )  # fmt: off
 
-        # 2. Embedd the prompt
+        # 2. Embed the prompt
         t5_tokens = self.t5_tokenizer.tokenize(prompt)
         clip_tokens = self.clip_tokenizer.tokenize(prompt)
         prompt_embeds = self.t5_text_encoder.forward(t5_tokens)
         pooled_prompt_embeds = self.clip_text_encoder.forward(clip_tokens)
 
-        step_wise_images = []
-        if stepwise_output_dir:
-            stepwise_output_dir.mkdir(parents=True, exist_ok=True)
         for t in time_steps:
             try:
                 # 3.t Predict the noise
@@ -106,39 +120,18 @@ class Flux1:
                 dt = config.sigmas[t + 1] - config.sigmas[t]
                 latents += noise * dt
 
+                # Handle stepwise output if enabled
+                stepwise_handler.process_step(t, latents)
+
                 # Evaluate to enable progress tracking
                 mx.eval(latents)
 
-                if stepwise_output_dir:
+            except KeyboardInterrupt:
-                    stepwise_decoded = self.vae.decode(Flux1._unpack_latents(latents, config.height, config.width))
+                stepwise_handler.handle_interruption()
-                    # performance todo: Pillow mostly uses CPU,
+                raise StopImageGenerationException(f"Stopping image generation at step {t + 1}/{len(time_steps)}")
-                    # can try to improve image generation performance by offloading
-                    # stepwise image processing to the CPU via threading
-                    stepwise_img = ImageUtil.to_image(
-                        decoded_latents=stepwise_decoded,
-                        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,
-                    )
-                    step_wise_images.append(stepwise_img)
-                    if not stepwise_composite_only:
-                        stepwise_img.save(
-                            path=stepwise_output_dir / f"seed_{seed}_step{t+1}of{len(time_steps)}.png",
-                            export_json_metadata=False,
-                        )
-            except KeyboardInterrupt:  # noqa: PERF203
-                raise StopImageGenerationException(f"Stopping image generation at step {t+1}/{len(time_steps)}")
-            finally:
-                if step_wise_images and stepwise_output_dir:
-                    composite_img = ImageUtil.to_composite_image(step_wise_images)
-                    composite_img.save(stepwise_output_dir / f"seed_{seed}_composite.png")
 
         # 5. Decode the latent array and return the image
-        latents = Flux1._unpack_latents(latents, config.height, config.width)
+        latents = ArrayUtil.unpack_latents(latents, config.height, config.width)
         decoded = self.vae.decode(latents)
         return ImageUtil.to_image(
             decoded_latents=decoded,
@@ -151,13 +144,6 @@ class Flux1:
             config=config,
         )
 
-    @staticmethod
-    def _unpack_latents(latents: mx.array, height: int, width: int) -> mx.array:
-        latents = mx.reshape(latents, (1, height // 16, width // 16, 16, 2, 2))
-        latents = mx.transpose(latents, (0, 3, 1, 4, 2, 5))
-        latents = mx.reshape(latents, (1, 16, height // 16 * 2, width // 16 * 2))
-        return latents
-
     @staticmethod
     def from_alias(alias: str, quantize: int | None = None) -> "Flux1":
         return Flux1(

src/mflux/generate.py

@@ -46,13 +46,13 @@ def main():
         image = flux.generate_image(
             seed=int(time.time()) if args.seed is None else args.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,
                 width=args.width,
                 guidance=args.guidance,
             ),
-            stepwise_output_dir=Path(args.stepwise_image_output_dir) if args.stepwise_image_output_dir else None,
         )
 
         # Save the image

src/mflux/generate_controlnet.py

@@ -52,6 +52,7 @@ def main():
             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=ConfigControlnet(
                 num_inference_steps=args.steps,
                 height=args.height,
@@ -59,7 +60,6 @@ def main():
                 guidance=args.guidance,
                 controlnet_strength=args.controlnet_strength,
             ),
-            stepwise_output_dir=Path(args.stepwise_image_output_dir) if args.stepwise_image_output_dir else None,
         )
 
         # Save the image

src/mflux/post_processing/array_util.py

@@ -0,0 +1,17 @@
+import mlx.core as mx
+
+
+class ArrayUtil:
+    @staticmethod
+    def unpack_latents(latents: mx.array, height: int, width: int) -> mx.array:
+        latents = mx.reshape(latents, (1, height // 16, width // 16, 16, 2, 2))
+        latents = mx.transpose(latents, (0, 3, 1, 4, 2, 5))
+        latents = mx.reshape(latents, (1, 16, height // 16 * 2, width // 16 * 2))
+        return latents
+
+    @staticmethod
+    def pack_latents(latents: mx.array, height: int, width: int) -> mx.array:
+        latents = mx.reshape(latents, (1, 16, height // 16, 2, width // 16, 2))
+        latents = mx.transpose(latents, (0, 2, 4, 1, 3, 5))
+        latents = mx.reshape(latents, (1, (width // 16) * (height // 16), 64))
+        return latents

src/mflux/post_processing/stepwise_handler.py

@@ -0,0 +1,55 @@
+from pathlib import Path
+
+import mlx.core as mx
+
+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,
+        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 process_step(self, step: int, latents: mx.array):
+        if self.output_dir:
+            unpack_latents = ArrayUtil.unpack_latents(latents, self.config.height, self.config.width)
+            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{step + 1}of{len(self.time_steps)}.png",
+                export_json_metadata=False,
+            )
+
+    def handle_interruption(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")