97 lines
3.4 KiB
Python
97 lines
3.4 KiB
Python
import os
|
|
from pathlib import Path
|
|
from typing import Type, Union
|
|
|
|
from mflux.config.config import Config
|
|
from mflux.config.model_config import ModelConfig
|
|
from mflux.models.flux.variants.txt2img.flux import Flux1
|
|
from mflux.models.qwen.variants.txt2img.qwen_image import QwenImage
|
|
from mflux.utils.image_compare import ImageCompare
|
|
|
|
|
|
class ImageGeneratorTestHelper:
|
|
@staticmethod
|
|
def assert_matches_reference_image(
|
|
reference_image_path: str,
|
|
output_image_path: str,
|
|
model_class: Type[Union[Flux1, QwenImage]],
|
|
model_config: ModelConfig,
|
|
prompt: str,
|
|
steps: int,
|
|
seed: int,
|
|
quantize=8,
|
|
height: int | None = None,
|
|
width: int | None = None,
|
|
image_path: str | None = None,
|
|
image_strength: float | None = None,
|
|
lora_paths: list[str] | None = None,
|
|
lora_scales: list[float] | None = None,
|
|
lora_names: list[str] | None = None,
|
|
lora_repo_id: str | None = None,
|
|
negative_prompt: str | None = None,
|
|
guidance: 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
|
|
model_kwargs = {
|
|
"model_config": model_config,
|
|
"quantize": quantize,
|
|
"lora_paths": lora_paths,
|
|
"lora_scales": lora_scales,
|
|
}
|
|
|
|
# Add HuggingFace LoRA parameters if provided
|
|
if lora_names is not None:
|
|
model_kwargs["lora_names"] = lora_names
|
|
if lora_repo_id is not None:
|
|
model_kwargs["lora_repo_id"] = lora_repo_id
|
|
|
|
model = model_class(**model_kwargs)
|
|
config_kwargs = {
|
|
"num_inference_steps": steps,
|
|
"image_path": ImageGeneratorTestHelper.resolve_path(image_path),
|
|
"image_strength": image_strength,
|
|
"height": height,
|
|
"width": width,
|
|
}
|
|
|
|
# Add guidance if provided
|
|
if guidance is not None:
|
|
config_kwargs["guidance"] = guidance
|
|
|
|
generate_kwargs = {
|
|
"seed": seed,
|
|
"prompt": prompt,
|
|
"config": Config(**config_kwargs),
|
|
}
|
|
|
|
# Add negative_prompt for Qwen models
|
|
if model_class == QwenImage and negative_prompt is not None:
|
|
generate_kwargs["negative_prompt"] = negative_prompt
|
|
|
|
# when
|
|
image = model.generate_image(**generate_kwargs)
|
|
image.save(path=output_image_path, overwrite=True)
|
|
|
|
# then
|
|
ImageCompare.check_images_close_enough(
|
|
output_image_path,
|
|
reference_image_path,
|
|
"Generated image doesn't match reference image.",
|
|
)
|
|
finally:
|
|
# cleanup
|
|
if os.path.exists(output_image_path) and "MFLUX_PRESERVE_TEST_OUTPUT" not in os.environ:
|
|
os.remove(output_image_path)
|
|
|
|
@staticmethod
|
|
def resolve_path(path) -> Path | None:
|
|
if path is None:
|
|
return None
|
|
return Path(__file__).parent.parent.parent / "resources" / path
|