Merge pull request #70 from filipstrand/add-stepwise-handler

Add stepwise handler and tests

Makefile

@@ -2,6 +2,7 @@
 
 PYTHON_VERSION = 3.11
 VENV_DIR = .venv
+PYTHON = $(VENV_DIR)/bin/python
 
 # Default target
 .PHONY: all
@@ -39,6 +40,14 @@ ensure-ruff:
 		uv tool install ruff; \
 	fi
 
+# ensure pytest is available
+.PHONY: ensure-pytest
+ensure-pytest:
+	@if ! $(PYTHON) -c "import pytest" 2>/dev/null; then \
+		echo "pytest required for testing. Installing pytest..."; \
+		uv pip install pytest; \
+	fi
+
 # Create virtual environment with uv
 .PHONY: venv-init
 venv-init: expect-arm64 expect-uv
@@ -81,10 +90,9 @@ check: ensure-ruff
 
 # Run tests
 .PHONY: test
-test:
+test: ensure-pytest
 	# 🏗️ Running tests...
-	# mock success stub for future test suite 😜
-	@true
+	$(PYTHON) -m pytest
 	# ✅ Tests completed
 
 # Clean up

README.md

@@ -496,7 +496,6 @@ with different prompts and LoRA adapters active.
 
 ### ✅ TODO
 
-- [ ] Establish unit test suite
 - [ ] LoRA fine-tuning
 - [ ] Frontend support (Gradio/Streamlit/Other?)
 

pyproject.toml

@@ -25,6 +25,12 @@ dependencies = [
     "tqdm>=4.66.5,<5.0",
     "transformers>=4.44.0,<5.0",
 ]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=8.0.0,<9.0"
+]
+
 classifiers = [
     "Intended Audience :: Developers",
     "Operating System :: MacOS",
@@ -93,3 +99,8 @@ docstring-code-format = false
 # This only has an effect when the `docstring-code-format` setting is
 # enabled.
 docstring-code-line-length = "dynamic"
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+python_files = "test_*.py"
+addopts = "-v"

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:
-                    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")
+            except KeyboardInterrupt:
+                stepwise_handler.handle_interruption()
+                raise StopImageGenerationException(f"Stopping image generation at step {t + 1}/{len(time_steps)}")
 
         # 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:
-                    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")
+            except KeyboardInterrupt:
+                stepwise_handler.handle_interruption()
+                raise StopImageGenerationException(f"Stopping image generation at step {t + 1}/{len(time_steps)}")
 
         # 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

@@ -15,7 +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('--stepwise-image-output-dir', type=str, default=None, help='[EXPERIMENTAL] Output dir to write step-wise images and their final composite image to. This feature may change in future versions.')
     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")
@@ -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")

tests/helpers/image_generation_controlnet_test_helper.py

@@ -0,0 +1,65 @@
+import os
+
+import numpy as np
+from PIL import Image
+
+from mflux import ModelConfig, Flux1Controlnet, ConfigControlnet
+from tests.helpers.image_generation_test_helper import ImageGeneratorTestHelper
+
+
+class ImageGeneratorControlnetTestHelper:
+    @staticmethod
+    def assert_matches_reference_image(
+        reference_image_path: str,
+        output_image_path: str,
+        controlnet_image_path: str,
+        model_config: ModelConfig,
+        prompt: str,
+        steps: int,
+        seed: int,
+        controlnet_strength: float,
+        lora_paths: list[str] | None = None,
+        lora_scales: list[float] | None = None,
+    ):
+        # resolve paths
+        reference_image_path = ImageGeneratorTestHelper.resolve_path(reference_image_path)
+        output_image_path = ImageGeneratorTestHelper.resolve_path(output_image_path)
+        controlnet_image_path = str(ImageGeneratorTestHelper.resolve_path(controlnet_image_path))
+        lora_paths = [str(ImageGeneratorTestHelper.resolve_path(p)) for p in lora_paths] if lora_paths else None
+
+        try:
+            # given
+            flux = Flux1Controlnet(
+                model_config=model_config,
+                quantize=8,
+                lora_paths=lora_paths,
+                lora_scales=lora_scales,
+            )
+
+            # when
+            image = flux.generate_image(
+                seed=seed,
+                prompt=prompt,
+                output=str(output_image_path),
+                controlnet_image_path=controlnet_image_path,
+                controlnet_save_canny=False,
+                config=ConfigControlnet(
+                    num_inference_steps=steps,
+                    height=768,
+                    width=493,
+                    controlnet_strength=controlnet_strength,
+                ),
+            )
+            image.save(path=output_image_path)
+
+            # then
+            np.testing.assert_array_equal(
+                np.array(Image.open(output_image_path)),
+                np.array(Image.open(reference_image_path)),
+                err_msg="Generated image doesn't match reference image",
+            )
+
+        finally:
+            # cleanup
+            if os.path.exists(output_image_path):
+                os.remove(output_image_path)

tests/helpers/image_generation_test_helper.py

@@ -0,0 +1,62 @@
+import os
+from pathlib import Path
+
+import numpy as np
+from PIL import Image
+
+from mflux import Flux1, Config, ModelConfig
+
+
+class ImageGeneratorTestHelper:
+    @staticmethod
+    def assert_matches_reference_image(
+        reference_image_path: str,
+        output_image_path: str,
+        model_config: ModelConfig,
+        prompt: str,
+        steps: int,
+        seed: int,
+        lora_paths: list[str] | None = None,
+        lora_scales: list[float] | None = None,
+    ):
+        # resolve paths
+        reference_image_path = ImageGeneratorTestHelper.resolve_path(reference_image_path)
+        output_image_path = ImageGeneratorTestHelper.resolve_path(output_image_path)
+        lora_paths = [str(ImageGeneratorTestHelper.resolve_path(p)) for p in lora_paths] if lora_paths else None
+
+        try:
+            # given
+            flux = Flux1(
+                model_config=model_config,
+                quantize=8,
+                lora_paths=lora_paths,
+                lora_scales=lora_scales
+            )  # fmt: off
+
+            # when
+            image = flux.generate_image(
+                seed=seed,
+                prompt=prompt,
+                config=Config(
+                    num_inference_steps=steps,
+                    height=341,
+                    width=768,
+                ),
+            )
+            image.save(path=output_image_path)
+
+            # then
+            np.testing.assert_array_equal(
+                np.array(Image.open(output_image_path)),
+                np.array(Image.open(reference_image_path)),
+                err_msg="Generated image doesn't match reference image",
+            )
+
+        finally:
+            # cleanup
+            if os.path.exists(output_image_path):
+                os.remove(output_image_path)
+
+    @staticmethod
+    def resolve_path(path) -> Path:
+        return Path(__file__).parent.parent / "resources" / path

tests/test_generate_image.py

@@ -0,0 +1,38 @@
+from mflux import ModelConfig
+from tests.helpers.image_generation_test_helper import ImageGeneratorTestHelper
+
+
+class TestImageGenerator:
+    OUTPUT_IMAGE_FILENAME = "output.png"
+
+    def test_image_generation_schnell(self):
+        ImageGeneratorTestHelper.assert_matches_reference_image(
+            reference_image_path="reference_schnell.png",
+            output_image_path=TestImageGenerator.OUTPUT_IMAGE_FILENAME,
+            model_config=ModelConfig.FLUX1_SCHNELL,
+            steps=2,
+            seed=42,
+            prompt="Luxury food photograph",
+        )
+
+    def test_image_generation_dev(self):
+        ImageGeneratorTestHelper.assert_matches_reference_image(
+            reference_image_path="reference_dev.png",
+            output_image_path=TestImageGenerator.OUTPUT_IMAGE_FILENAME,
+            model_config=ModelConfig.FLUX1_DEV,
+            steps=15,
+            seed=42,
+            prompt="Luxury food photograph",
+        )
+
+    def test_image_generation_dev_lora(self):
+        ImageGeneratorTestHelper.assert_matches_reference_image(
+            reference_image_path="reference_dev_lora.png",
+            output_image_path=TestImageGenerator.OUTPUT_IMAGE_FILENAME,
+            model_config=ModelConfig.FLUX1_DEV,
+            steps=15,
+            seed=42,
+            prompt="mkym this is made of wool, burger",
+            lora_paths=["FLUX-dev-lora-MiaoKa-Yarn-World.safetensors"],
+            lora_scales=[1.0],
+        )

tests/test_generate_image_controlnet.py

@@ -0,0 +1,45 @@
+from mflux import ModelConfig
+from tests.helpers.image_generation_controlnet_test_helper import ImageGeneratorControlnetTestHelper
+
+
+class TestImageGeneratorControlnet:
+    OUTPUT_IMAGE_FILENAME = "output.png"
+    CONTROLNET_REFERENCE_FILENAME = "controlnet_reference.png"
+
+    def test_image_generation_schnell_controlnet(self):
+        ImageGeneratorControlnetTestHelper.assert_matches_reference_image(
+            reference_image_path="reference_controlnet_schnell.png",
+            output_image_path=TestImageGeneratorControlnet.OUTPUT_IMAGE_FILENAME,
+            controlnet_image_path=TestImageGeneratorControlnet.CONTROLNET_REFERENCE_FILENAME,
+            model_config=ModelConfig.FLUX1_SCHNELL,
+            steps=2,
+            seed=43,
+            prompt="The joker with a hat and a cane",
+            controlnet_strength=0.4,
+        )
+
+    def test_image_generation_dev_controlnet(self):
+        ImageGeneratorControlnetTestHelper.assert_matches_reference_image(
+            reference_image_path="reference_controlnet_dev.png",
+            output_image_path=TestImageGeneratorControlnet.OUTPUT_IMAGE_FILENAME,
+            controlnet_image_path=TestImageGeneratorControlnet.CONTROLNET_REFERENCE_FILENAME,
+            model_config=ModelConfig.FLUX1_DEV,
+            steps=15,
+            seed=42,
+            prompt="The joker with a hat and a cane",
+            controlnet_strength=0.4,
+        )
+
+    def test_image_generation_dev_lora_controlnet(self):
+        ImageGeneratorControlnetTestHelper.assert_matches_reference_image(
+            reference_image_path="reference_controlnet_dev_lora.png",
+            output_image_path=TestImageGeneratorControlnet.OUTPUT_IMAGE_FILENAME,
+            controlnet_image_path=TestImageGeneratorControlnet.CONTROLNET_REFERENCE_FILENAME,
+            model_config=ModelConfig.FLUX1_DEV,
+            steps=15,
+            seed=43,
+            prompt="mkym this is made of wool, The joker with a hat and a cane",
+            lora_paths=["FLUX-dev-lora-MiaoKa-Yarn-World.safetensors"],
+            lora_scales=[1.0],
+            controlnet_strength=0.4,
+        )

tests/test_readme_example.sh

@@ -1,39 +0,0 @@
-#!/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 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