diff --git a/Makefile b/Makefile index b200161..9d58705 100644 --- a/Makefile +++ b/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 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/src/mflux/__init__.py b/src/mflux/__init__.py index e7c48dc..e53c33f 100644 --- a/src/mflux/__init__.py +++ b/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__ = [ diff --git a/src/mflux/controlnet/flux_controlnet.py b/src/mflux/controlnet/flux_controlnet.py index 3241c30..3cca900 100644 --- a/src/mflux/controlnet/flux_controlnet.py +++ b/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) diff --git a/src/mflux/error/__init__.py b/src/mflux/error/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/mflux/exceptions.py b/src/mflux/error/exceptions.py similarity index 100% rename from src/mflux/exceptions.py rename to src/mflux/error/exceptions.py diff --git a/src/mflux/flux/flux.py b/src/mflux/flux/flux.py index 41de838..634bf32 100644 --- a/src/mflux/flux/flux.py +++ b/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( diff --git a/src/mflux/generate.py b/src/mflux/generate.py index 36d1cf5..d6726f0 100644 --- a/src/mflux/generate.py +++ b/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 diff --git a/src/mflux/generate_controlnet.py b/src/mflux/generate_controlnet.py index c8813d8..470236d 100644 --- a/src/mflux/generate_controlnet.py +++ b/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 diff --git a/src/mflux/post_processing/array_util.py b/src/mflux/post_processing/array_util.py new file mode 100644 index 0000000..9fe82b8 --- /dev/null +++ b/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 diff --git a/src/mflux/post_processing/stepwise_handler.py b/src/mflux/post_processing/stepwise_handler.py new file mode 100644 index 0000000..d3e9cd6 --- /dev/null +++ b/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") 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..5ed4f6a --- /dev/null +++ b/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) diff --git a/tests/helpers/image_generation_test_helper.py b/tests/helpers/image_generation_test_helper.py new file mode 100644 index 0000000..72eb3f7 --- /dev/null +++ b/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 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