40 lines
1.5 KiB
Python
40 lines
1.5 KiB
Python
import pytest
|
|
|
|
from mflux.models.common.config import ModelConfig
|
|
from tests.image_generation.helpers.image_generation_concept_test_helper import ImageGenerationConceptTestHelper
|
|
|
|
|
|
class TestImageGeneratorConcept:
|
|
@pytest.mark.slow
|
|
def test_concept_attention_generation(self):
|
|
ImageGenerationConceptTestHelper.assert_matches_reference_image_concept(
|
|
reference_heatmap_path="reference_concept_schnell_heatmap.png",
|
|
output_heatmap_path="output_concept_schnell_heatmap.png",
|
|
model_config=ModelConfig.schnell(),
|
|
prompt="A dragon on a hill",
|
|
concept="dragon",
|
|
steps=4,
|
|
seed=44,
|
|
height=512,
|
|
width=512,
|
|
heatmap_layer_indices=[15, 16, 17, 18],
|
|
heatmap_timesteps=[0, 1, 2, 3],
|
|
)
|
|
|
|
@pytest.mark.slow
|
|
def test_concept_attention_from_image(self):
|
|
ImageGenerationConceptTestHelper.assert_matches_reference_image_concept_from_image(
|
|
reference_heatmap_path="reference_concept_from_image_schnell_heatmap.png",
|
|
output_heatmap_path="output_concept_from_image_schnell_heatmap.png",
|
|
input_image_path="reference_depth_dev_from_image.png",
|
|
model_config=ModelConfig.schnell(),
|
|
prompt="A photo of cartoon of Albert Einstein",
|
|
concept="man",
|
|
steps=4,
|
|
seed=42,
|
|
height=512,
|
|
width=320,
|
|
heatmap_layer_indices=[15, 16, 17, 18],
|
|
heatmap_timesteps=[0, 1, 2, 3],
|
|
)
|