Add image generation tests

rename

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"

tests/helpers/image_generation_controlnet_test_helper.py

@@ -0,0 +1,63 @@
+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
+
+        # 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",
+        )
+
+        # cleanup
+        if os.path.exists(output_image_path):
+            os.remove(output_image_path)

tests/helpers/image_generation_test_helper.py

@@ -0,0 +1,60 @@
+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
+
+        # 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",
+        )
+
+        # 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