diff --git a/README.md b/README.md index f5236b4..03aa39c 100644 --- a/README.md +++ b/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?) diff --git a/pyproject.toml b/pyproject.toml index 3bd2267..edfba11 100644 --- a/pyproject.toml +++ b/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" \ No newline at end of file diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/helpers/__init__.py b/tests/helpers/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/helpers/image_generation_controlnet_test_helper.py b/tests/helpers/image_generation_controlnet_test_helper.py new file mode 100644 index 0000000..c629424 --- /dev/null +++ b/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) diff --git a/tests/helpers/image_generation_test_helper.py b/tests/helpers/image_generation_test_helper.py new file mode 100644 index 0000000..2c5666b --- /dev/null +++ b/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 diff --git a/tests/resources/FLUX-dev-lora-MiaoKa-Yarn-World.safetensors b/tests/resources/FLUX-dev-lora-MiaoKa-Yarn-World.safetensors new file mode 100644 index 0000000..36998aa Binary files /dev/null and b/tests/resources/FLUX-dev-lora-MiaoKa-Yarn-World.safetensors differ diff --git a/tests/resources/controlnet_reference.png b/tests/resources/controlnet_reference.png new file mode 100644 index 0000000..be6792b Binary files /dev/null and b/tests/resources/controlnet_reference.png differ diff --git a/tests/resources/reference_controlnet_dev.png b/tests/resources/reference_controlnet_dev.png new file mode 100644 index 0000000..e62f8b4 Binary files /dev/null and b/tests/resources/reference_controlnet_dev.png differ diff --git a/tests/resources/reference_controlnet_dev_lora.png b/tests/resources/reference_controlnet_dev_lora.png new file mode 100644 index 0000000..086c4dc Binary files /dev/null and b/tests/resources/reference_controlnet_dev_lora.png differ diff --git a/tests/resources/reference_controlnet_schnell.png b/tests/resources/reference_controlnet_schnell.png new file mode 100644 index 0000000..0ba9dd3 Binary files /dev/null and b/tests/resources/reference_controlnet_schnell.png differ diff --git a/tests/resources/reference_dev.png b/tests/resources/reference_dev.png new file mode 100644 index 0000000..eceaf81 Binary files /dev/null and b/tests/resources/reference_dev.png differ diff --git a/tests/resources/reference_dev_lora.png b/tests/resources/reference_dev_lora.png new file mode 100644 index 0000000..bd1844a Binary files /dev/null and b/tests/resources/reference_dev_lora.png differ diff --git a/tests/resources/reference_schnell.png b/tests/resources/reference_schnell.png new file mode 100644 index 0000000..ea5d799 Binary files /dev/null and b/tests/resources/reference_schnell.png differ diff --git a/tests/test_generate_image.py b/tests/test_generate_image.py new file mode 100644 index 0000000..c52d5a8 --- /dev/null +++ b/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], + ) diff --git a/tests/test_generate_image_controlnet.py b/tests/test_generate_image_controlnet.py new file mode 100644 index 0000000..35b1f84 --- /dev/null +++ b/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, + ) diff --git a/tests/test_readme_example.sh b/tests/test_readme_example.sh deleted file mode 100755 index 6346e3b..0000000 --- a/tests/test_readme_example.sh +++ /dev/null @@ -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