59 lines
2.1 KiB
Python
59 lines
2.1 KiB
Python
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,
|
|
height=341,
|
|
width=768,
|
|
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,
|
|
height=341,
|
|
width=768,
|
|
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,
|
|
height=341,
|
|
width=768,
|
|
prompt="mkym this is made of wool, burger",
|
|
lora_paths=["FLUX-dev-lora-MiaoKa-Yarn-World.safetensors"],
|
|
lora_scales=[1.0],
|
|
)
|
|
|
|
def test_image_generation_dev_image_to_image(self):
|
|
ImageGeneratorTestHelper.assert_matches_reference_image(
|
|
reference_image_path="reference_dev_image_to_image_result.png",
|
|
init_image_path="reference_dev_image_to_image_init.png",
|
|
init_image_strength=0.126, # fill 1/8 steps
|
|
output_image_path=TestImageGenerator.OUTPUT_IMAGE_FILENAME,
|
|
model_config=ModelConfig.FLUX1_DEV,
|
|
steps=8,
|
|
seed=42,
|
|
height=768,
|
|
width=768,
|
|
prompt="astronauts in a jungle",
|
|
)
|