option to save stepwise and composite images, add interruptibility

.gitignore

@@ -14,4 +14,7 @@
 *.safetensors
 *.json
 
-*.egg-info
+*.egg-info
+
+# build/ generated from 'pip install -e .' in developer mode
+build/

pyproject.toml

@@ -57,7 +57,7 @@ respect-gitignore = true
 # Enable Pyflakes (`F`) and a subset of the pycodestyle (`E`) codes by default.
 # Unlike Flake8, Ruff doesn't enable pycodestyle warnings (`W`) or
 # McCabe complexity (`C901`) by default.
-select = ["E4", "E7", "E9", "F"]
+select = ["BLE", "E4", "E7", "E9", "F", "ICN", "LOG", "PERF", "W"]
 ignore = []
 
 # Allow fix for all enabled rules (when `--fix`) is provided.

src/mflux/__init__.py

@@ -3,6 +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.post_processing.image_util import ImageUtil
 
 __all__ = [
@@ -12,4 +13,5 @@ __all__ = [
     "ConfigControlnet",
     "ModelConfig",
     "ImageUtil",
+    "StopImageGenerationException",
 ]

src/mflux/controlnet/flux_controlnet.py

@@ -1,16 +1,16 @@
 import logging
-from typing import TYPE_CHECKING
-
 import mlx.core as mx
 from mlx import nn
+from pathlib import Path
 from tqdm import tqdm
-
+from typing import TYPE_CHECKING
 from mflux.config.config import ConfigControlnet
 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.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
@@ -110,7 +110,9 @@ class Flux1Controlnet:
             output: str,
             controlnet_image_path: str,
             controlnet_save_canny: bool = False,
-            config: ConfigControlnet = ConfigControlnet()
+            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)
@@ -139,34 +141,68 @@ class Flux1Controlnet:
         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:
-            # Compute controlnet samples
-            controlnet_block_samples, controlnet_single_block_samples = self.transformer_controlnet.forward(
-                t=t,
-                prompt_embeds=prompt_embeds,
-                pooled_prompt_embeds=pooled_prompt_embeds,
-                hidden_states=latents,
-                controlnet_cond=controlnet_cond,
-                config=config,
-            )
+            try:
+                # Compute controlnet samples
+                controlnet_block_samples, controlnet_single_block_samples = self.transformer_controlnet.forward(
+                    t=t,
+                    prompt_embeds=prompt_embeds,
+                    pooled_prompt_embeds=pooled_prompt_embeds,
+                    hidden_states=latents,
+                    controlnet_cond=controlnet_cond,
+                    config=config,
+                )
 
-            # 3.t Predict the noise
-            noise = self.transformer.predict(
-                t=t,
-                prompt_embeds=prompt_embeds,
-                pooled_prompt_embeds=pooled_prompt_embeds,
-                hidden_states=latents,
-                config=config,
-                controlnet_block_samples=controlnet_block_samples,
-                controlnet_single_block_samples=controlnet_single_block_samples,
-            )
+                # 3.t Predict the noise
+                noise = self.transformer.predict(
+                    t=t,
+                    prompt_embeds=prompt_embeds,
+                    pooled_prompt_embeds=pooled_prompt_embeds,
+                    hidden_states=latents,
+                    config=config,
+                    controlnet_block_samples=controlnet_block_samples,
+                    controlnet_single_block_samples=controlnet_single_block_samples,
+                )
 
-            # 4.t Take one denoise step
-            dt = config.sigmas[t + 1] - config.sigmas[t]
-            latents += noise * dt
+                # 4.t Take one denoise step
+                dt = config.sigmas[t + 1] - config.sigmas[t]
+                latents += noise * dt
 
-            # Evaluate to enable progress tracking
-            mx.eval(latents)
+                # Evaluate to enable progress tracking
+                mx.eval(latents)
+
+                if stepwise_output_dir:
+                    stepwise_decoded = self.vae.decode(
+                        Flux1Controlnet._unpack_latents(latents, config.height, config.width)
+                    )
+                    # 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)

src/mflux/exceptions.py

@@ -0,0 +1,18 @@
+class MFluxException(Exception):
+    """base class for all custom exceptions in mflux package."""
+
+
+class ImageSavingException(MFluxException):
+    """error ocurred while attempting to save image to storage."""
+
+
+class MetadataEmbedException(MFluxException):
+    """error ocurred while attempting to embed metadata in image"""
+
+
+class MFluxUserException(MFluxException):
+    """an exception raised by user behavior or intention."""
+
+
+class StopImageGenerationException(MFluxUserException):
+    """user has requested to stop a image generation in progress."""

src/mflux/flux/flux.py

@@ -1,10 +1,12 @@
 import mlx.core as mx
+from pathlib import Path
 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.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
@@ -69,7 +71,7 @@ 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()) -> GeneratedImage:
+    def generate_image(self, seed: int, prompt: str, config: Config = Config(), 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))
@@ -86,22 +88,54 @@ class Flux1:
         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:
-            # 3.t Predict the noise
-            noise = self.transformer.predict(
-                t=t,
-                prompt_embeds=prompt_embeds,
-                pooled_prompt_embeds=pooled_prompt_embeds,
-                hidden_states=latents,
-                config=config,
-            )
+            try:
+                # 3.t Predict the noise
+                noise = self.transformer.predict(
+                    t=t,
+                    prompt_embeds=prompt_embeds,
+                    pooled_prompt_embeds=pooled_prompt_embeds,
+                    hidden_states=latents,
+                    config=config,
+                )
 
-            # 4.t Take one denoise step
-            dt = config.sigmas[t + 1] - config.sigmas[t]
-            latents += noise * dt
+                # 4.t Take one denoise step
+                dt = config.sigmas[t + 1] - config.sigmas[t]
+                latents += noise * dt
 
-            # Evaluate to enable progress tracking
-            mx.eval(latents)
+                # Evaluate to enable progress tracking
+                mx.eval(latents)
+
+                if stepwise_output_dir:
+                    stepwise_decoded = self.vae.decode(Flux1._unpack_latents(latents, config.height, config.width))
+                    # 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 = Flux1._unpack_latents(latents, config.height, config.width)

src/mflux/generate.py

@@ -1,7 +1,8 @@
 import argparse
 import time
+from pathlib import Path
 
-from mflux import Flux1, Config, ModelConfig
+from mflux import Flux1, Config, ModelConfig, StopImageGenerationException
 
 
 def main():
@@ -14,6 +15,7 @@ def main():
     parser.add_argument("--height", type=int, default=1024, help="Image height (Default is 1024)")
     parser.add_argument("--width", type=int, default=1024, help="Image width (Default is 1024)")
     parser.add_argument("--steps", type=int, default=None, help="Inference Steps")
+    parser.add_argument('--stepwise-image-output-dir', type=str, default=None, help='Output dir to write step-wise images and their final composite image to.')
     parser.add_argument("--guidance", type=float, default=3.5, help="Guidance Scale (Default is 3.5)")
     parser.add_argument("--quantize",  "-q", type=int, choices=[4, 8], default=None, help="Quantize the model (4 or 8, Default is None)")
     parser.add_argument("--path", type=str, default=None, help="Local path for loading a model from disk")
@@ -39,20 +41,24 @@ def main():
         lora_scales=args.lora_scales,
     )
 
-    # Generate an image
-    image = flux.generate_image(
-        seed=int(time.time()) if args.seed is None else args.seed,
-        prompt=args.prompt,
-        config=Config(
-            num_inference_steps=args.steps,
-            height=args.height,
-            width=args.width,
-            guidance=args.guidance,
-        ),
-    )
+    try:
+        # Generate an image
+        image = flux.generate_image(
+            seed=int(time.time()) if args.seed is None else args.seed,
+            prompt=args.prompt,
+            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
-    image.save(path=args.output, export_json_metadata=args.metadata)
+        # Save the image
+        image.save(path=args.output, export_json_metadata=args.metadata)
+    except StopImageGenerationException as stop_exc:
+        print(stop_exc)
 
 
 if __name__ == "__main__":

src/mflux/generate_controlnet.py

@@ -1,7 +1,8 @@
 import argparse
 import time
+from pathlib import Path
 
-from mflux import Flux1Controlnet, ConfigControlnet, ModelConfig
+from mflux import Flux1Controlnet, ConfigControlnet, ModelConfig, StopImageGenerationException
 
 
 def main():
@@ -17,6 +18,7 @@ def main():
     parser.add_argument("--height", type=int, default=1024, help="Image height (Default is 1024)")
     parser.add_argument("--width", type=int, default=1024, help="Image width (Default is 1024)")
     parser.add_argument("--steps", type=int, default=None, help="Inference Steps")
+    parser.add_argument('--stepwise-image-output-dir', type=str, default=None, help='Output dir to write step-wise images and their final composite image to.')
     parser.add_argument("--guidance", type=float, default=3.5, help="Guidance Scale (Default is 3.5)")
     parser.add_argument("--quantize",  "-q", type=int, choices=[4, 8], default=None, help="Quantize the model (4 or 8, Default is None)")
     parser.add_argument("--path", type=str, default=None, help="Local path for loading a model from disk")
@@ -42,24 +44,28 @@ def main():
         lora_scales=args.lora_scales,
     )
 
-    # Generate an image
-    image = flux.generate_image(
-        seed=int(time.time()) if args.seed is None else args.seed,
-        prompt=args.prompt,
-        output=args.output,
-        controlnet_image_path=args.controlnet_image_path,
-        controlnet_save_canny=args.controlnet_save_canny,
-        config=ConfigControlnet(
-            num_inference_steps=args.steps,
-            height=args.height,
-            width=args.width,
-            guidance=args.guidance,
-            controlnet_strength=args.controlnet_strength,
-        ),
-    )
+    try:
+        # Generate an image
+        image = flux.generate_image(
+            seed=int(time.time()) if args.seed is None else args.seed,
+            prompt=args.prompt,
+            output=args.output,
+            controlnet_image_path=args.controlnet_image_path,
+            controlnet_save_canny=args.controlnet_save_canny,
+            config=ConfigControlnet(
+                num_inference_steps=args.steps,
+                height=args.height,
+                width=args.width,
+                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
-    image.save(path=args.output, export_json_metadata=args.metadata)
+        # Save the image
+        image.save(path=args.output, export_json_metadata=args.metadata)
+    except StopImageGenerationException as stop_exc:
+        print(stop_exc)
 
 
 if __name__ == "__main__":

src/mflux/post_processing/image_util.py

@@ -1,7 +1,7 @@
+import typing as t
 import json
 import logging
 from pathlib import Path
-
 import PIL
 import PIL.Image
 import mlx.core as mx
@@ -47,6 +47,17 @@ class ImageUtil:
             controlnet_strength=config.controlnet_strength if isinstance(config.config, ConfigControlnet) else None,
         )
 
+    def to_composite_image(generated_images: t.List[GeneratedImage]) -> Image:
+        # stitch horizontally
+        total_width = sum(gen_img.image.width for gen_img in generated_images)
+        max_height = max(gen_img.image.height for gen_img in generated_images)
+        composite_img = Image.new("RGB", (total_width, max_height))
+        current_x = 0
+        for index, gen_img in enumerate(generated_images):
+            composite_img.paste(gen_img.image, (current_x, 0))
+            current_x += gen_img.image.width
+        return composite_img
+
     @staticmethod
     def _denormalize(images: mx.array) -> mx.array:
         return mx.clip((images / 2 + 0.5), 0, 1)
@@ -119,7 +130,7 @@ class ImageUtil:
             if metadata is not None:
                 ImageUtil._embed_metadata(metadata, file_path)
                 log.info(f"Metadata embedded successfully at: {file_path}")
-        except Exception as e:
+        except Exception as e:  # noqa: BLE001
             log.error(f"Error saving image: {e}")
 
     @staticmethod
@@ -145,5 +156,5 @@ class ImageUtil:
             # Save the image with metadata
             image.save(path, exif=exif_bytes)
 
-        except Exception as e:
+        except Exception as e:  # noqa: BLE001
             log.error(f"Error embedding metadata: {e}")

tests/test_readme_example.sh

@@ -1,7 +1,39 @@
+#!/bin/zsh -e
+# ^ safe to assume Mac devs have zsh installed
+#   default since Catalina in 2019
+
+mkdir -p /tmp/mflux-test
+
 mflux-generate \
     --prompt "Luxury food photograph" \
     --model schnell \
     --steps 2 \
     --seed 2 \
-    --height 1024 \
-    --width 1024
+    --height 512 \
+    --width 512 \
+    --output /tmp/mflux-test/luxury_food.png
+
+# generate an image of a blue bird, then use it as input for the following test
+mflux-generate \
+    --prompt "blue bird, morning, spring" \
+    --model schnell \
+    --steps 2 \
+    --seed 24 \
+    --height 512 \
+    --width 512 \
+    --stepwise-image-output-dir /tmp/mflux-test \
+    --output /tmp/mflux-test/sf_blue_bird.png
+
+# use the image from the prior test, generate an image with similar visual structure
+mflux-generate-controlnet \
+    --prompt "yellow bird, afternoon, snowy mountain" \
+    --model schnell \
+    --controlnet-image-path /tmp/mflux-test/sf_blue_bird.png \
+    --controlnet-strength 0.7 \
+    --controlnet-save-canny \
+    --steps 2 \
+    --seed 42 \
+    --height 512 \
+    --width 512 \
+    --output /tmp/mflux-test/controlnet_sf_yellow_bird.png \
+    --stepwise-image-output-dir /tmp/mflux-test