diff --git a/pyproject.toml b/pyproject.toml index 7b07041..519c963 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -137,7 +137,7 @@ docstring-code-line-length = "dynamic" [tool.pytest.ini_options] testpaths = ["tests"] python_files = "test_*.py" -addopts = "-v --exitfirst --failed-first --showlocals --tb=long --full-trace" +addopts = "-v --failed-first --showlocals --tb=long --full-trace" # https://docs.astral.sh/ruff/settings/#lintisort [tool.ruff.lint.isort] diff --git a/tests/image_generation/helpers/image_compare.py b/tests/image_generation/helpers/image_compare.py new file mode 100644 index 0000000..eb71676 --- /dev/null +++ b/tests/image_generation/helpers/image_compare.py @@ -0,0 +1,46 @@ +import os +from pathlib import Path + +import cv2 +import numpy as np +from PIL import Image + + +class ReferenceVsOutputImageError(AssertionError): ... + + +# How we determined DEFAULT_MISMATCH_THRESHOLD value: by eye test +# successive mlx/mflux versions have generated images that are visually close enough +# to consider as a valid upgrade. Minor visual differences are attributable to mlx updates +DEFAULT_MISMATCH_THRESHOLD = 0.15 +ENV_MISMATCH_THRESHOLD = float(os.environ.get("MFLUX_IMAGE_MISMATCH_THRESHOLD", DEFAULT_MISMATCH_THRESHOLD)) + + +def check_images_close_enough( + image1_path: str | Path, + image2_path: str | Path, + error_message_prefix: str, + mismatch_threshold: float = ENV_MISMATCH_THRESHOLD, +): + image1_path = Path(image1_path) + image2_path = Path(image2_path) + image1_data = np.array(Image.open(image1_path)) + image2_data = np.array(Image.open(image2_path)) + closeness_array = np.isclose( + image1_data, + image2_data, + rtol=float(os.environ.get("MFLUX_IMAGE_ALLCLOSE_RTOL", 0.1)), + atol=0, # ignore absolute tolerance + ) + num_mismatched = np.count_nonzero(~closeness_array) + total_elements = image1_data.size + mismatch_ratio = num_mismatched / total_elements + if mismatch_ratio > mismatch_threshold: + diff = image1_data - image2_data + # Scale the difference so it's visible. A difference of 50 will become bright. + # We'll scale it so that a difference of 50 or more becomes pure white (255). + diff_visual = (diff / 50 * 255).clip(0, 255).astype(np.uint8) + cv2.imwrite(image2_path.with_stem(image1_path.stem + "_diff"), diff_visual) + raise AssertionError( + f"{error_message_prefix} Check {image1_path} vs {image2_path} :: their elements are {mismatch_ratio:.1%} different. Fails assertion for {mismatch_threshold=}" + ) diff --git a/tests/image_generation/helpers/image_generation_concept_test_helper.py b/tests/image_generation/helpers/image_generation_concept_test_helper.py index 48d0102..c01185c 100644 --- a/tests/image_generation/helpers/image_generation_concept_test_helper.py +++ b/tests/image_generation/helpers/image_generation_concept_test_helper.py @@ -1,14 +1,13 @@ import os from pathlib import Path -import numpy as np -from PIL import Image - from mflux.community.concept_attention.flux_concept import Flux1Concept from mflux.community.concept_attention.flux_concept_from_image import Flux1ConceptFromImage from mflux.config.config import Config from mflux.config.model_config import ModelConfig +from .image_compare import check_images_close_enough + class ImageGenerationConceptTestHelper: @staticmethod @@ -57,10 +56,11 @@ class ImageGenerationConceptTestHelper: image.save_concept_heatmap(path=output_heatmap_path, overwrite=True) # then - verify the heatmap matches reference - np.testing.assert_array_equal( - np.array(Image.open(output_heatmap_path)), - np.array(Image.open(reference_heatmap_path)), - err_msg=f"Generated concept heatmap doesn't match reference heatmap. Check {output_heatmap_path} vs {reference_heatmap_path}", + check_images_close_enough( + output_heatmap_path, + reference_heatmap_path, + "Generated concept heatmap doesn't match reference heatmap.", + mismatch_threshold=0.25, # special case for heatmap, allow higher threshold ) finally: @@ -117,15 +117,15 @@ class ImageGenerationConceptTestHelper: image.save_concept_heatmap(path=output_heatmap_path, overwrite=True) # then - verify the heatmap matches reference - np.testing.assert_array_equal( - np.array(Image.open(output_heatmap_path)), - np.array(Image.open(reference_heatmap_path)), - err_msg=f"Generated concept from image heatmap doesn't match reference heatmap. Check {output_heatmap_path} vs {reference_heatmap_path}", + check_images_close_enough( + output_heatmap_path, + reference_heatmap_path, + "Generated concept from image heatmap doesn't match reference heatmap.", ) finally: # cleanup - if os.path.exists(output_heatmap_path): + if os.path.exists(output_heatmap_path) and "MFLUX_PRESERVE_TEST_OUTPUT" not in os.environ: os.remove(output_heatmap_path) @staticmethod diff --git a/tests/image_generation/helpers/image_generation_controlnet_test_helper.py b/tests/image_generation/helpers/image_generation_controlnet_test_helper.py index fbf6009..998c24f 100644 --- a/tests/image_generation/helpers/image_generation_controlnet_test_helper.py +++ b/tests/image_generation/helpers/image_generation_controlnet_test_helper.py @@ -1,13 +1,12 @@ import os -import numpy as np -from PIL import Image - from mflux.config.config import Config from mflux.config.model_config import ModelConfig from mflux.controlnet.flux_controlnet import Flux1Controlnet from tests.image_generation.helpers.image_generation_test_helper import ImageGeneratorTestHelper +from .image_compare import check_images_close_enough + class ImageGeneratorControlnetTestHelper: @staticmethod @@ -55,13 +54,13 @@ class ImageGeneratorControlnetTestHelper: image.save(path=output_image_path, overwrite=True) # 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", + check_images_close_enough( + output_image_path, + reference_image_path, + "Generated controlnet image doesn't match reference image", ) finally: # cleanup - if os.path.exists(output_image_path): + if os.path.exists(output_image_path) and "MFLUX_PRESERVE_TEST_OUTPUT" not in os.environ: os.remove(output_image_path) diff --git a/tests/image_generation/helpers/image_generation_depth_test_helper.py b/tests/image_generation/helpers/image_generation_depth_test_helper.py index 9487e1e..5d646d2 100644 --- a/tests/image_generation/helpers/image_generation_depth_test_helper.py +++ b/tests/image_generation/helpers/image_generation_depth_test_helper.py @@ -1,13 +1,12 @@ import os from pathlib import Path -import numpy as np -from PIL import Image - from mflux.config.config import Config from mflux.config.model_config import ModelConfig from mflux.flux_tools.depth.flux_depth import Flux1Depth +from .image_compare import check_images_close_enough + class ImageGeneratorDepthTestHelper: @staticmethod @@ -47,15 +46,14 @@ class ImageGeneratorDepthTestHelper: image.save(path=output_image_path, overwrite=True) # then - np.testing.assert_array_equal( - np.array(Image.open(output_image_path)), - np.array(Image.open(reference_image_path)), - err_msg=f"Generated image doesn't match reference image. Check {output_image_path} vs {reference_image_path}", + check_images_close_enough( + output_image_path, + reference_image_path, + "Generated depth image doesn't match reference image.", ) - finally: # cleanup - if os.path.exists(output_image_path): + if os.path.exists(output_image_path) and "MFLUX_PRESERVE_TEST_OUTPUT" not in os.environ: os.remove(output_image_path) @staticmethod diff --git a/tests/image_generation/helpers/image_generation_fill_test_helper.py b/tests/image_generation/helpers/image_generation_fill_test_helper.py index 5a6a914..7b61b9e 100644 --- a/tests/image_generation/helpers/image_generation_fill_test_helper.py +++ b/tests/image_generation/helpers/image_generation_fill_test_helper.py @@ -1,14 +1,13 @@ import os -import numpy as np -from PIL import Image - from mflux.config.config import Config from mflux.config.model_config import ModelConfig from mflux.flux_tools.fill.flux_fill import Flux1Fill from mflux.ui import defaults as ui_defaults from tests.image_generation.helpers.image_generation_test_helper import ImageGeneratorTestHelper +from .image_compare import check_images_close_enough + class ImageGeneratorFillTestHelper: @staticmethod @@ -57,13 +56,13 @@ class ImageGeneratorFillTestHelper: image.save(path=output_image_path, overwrite=True) # then - np.testing.assert_array_equal( - np.array(Image.open(output_image_path)), - np.array(Image.open(reference_image_path)), - err_msg=f"Generated image doesn't match reference image. Check {output_image_path} vs {reference_image_path}", + check_images_close_enough( + output_image_path, + reference_image_path, + "Generated fill image doesn't match reference image.", ) finally: # cleanup - if os.path.exists(output_image_path): + if os.path.exists(output_image_path) and "MFLUX_PRESERVE_TEST_OUTPUT" not in os.environ: os.remove(output_image_path) diff --git a/tests/image_generation/helpers/image_generation_ic_edit_test_helper.py b/tests/image_generation/helpers/image_generation_ic_edit_test_helper.py index 62a4afe..6b389ed 100644 --- a/tests/image_generation/helpers/image_generation_ic_edit_test_helper.py +++ b/tests/image_generation/helpers/image_generation_ic_edit_test_helper.py @@ -1,14 +1,13 @@ import os from pathlib import Path -import numpy as np -from PIL import Image - from mflux.community.in_context.flux_in_context_fill import Flux1InContextFill from mflux.community.in_context.utils.in_context_loras import prepare_ic_edit_loras from mflux.config.config import Config from mflux.config.model_config import ModelConfig +from .image_compare import check_images_close_enough + class ImageGeneratorICEditTestHelper: @staticmethod @@ -56,15 +55,15 @@ class ImageGeneratorICEditTestHelper: image.get_right_half().save(path=output_image_path, overwrite=True) # then - np.testing.assert_array_equal( - np.array(Image.open(output_image_path)), - np.array(Image.open(reference_image_path)), - err_msg=f"Generated image doesn't match reference image. Check {output_image_path} vs {reference_image_path}", + check_images_close_enough( + output_image_path, + reference_image_path, + "Generated ic-edit image doesn't match reference image.", ) finally: # cleanup - if os.path.exists(output_image_path): + if os.path.exists(output_image_path) and "MFLUX_PRESERVE_TEST_OUTPUT" not in os.environ: os.remove(output_image_path) @staticmethod diff --git a/tests/image_generation/helpers/image_generation_in_context_test_helper.py b/tests/image_generation/helpers/image_generation_in_context_test_helper.py index f912564..df20d0f 100644 --- a/tests/image_generation/helpers/image_generation_in_context_test_helper.py +++ b/tests/image_generation/helpers/image_generation_in_context_test_helper.py @@ -1,14 +1,13 @@ import os from pathlib import Path -import numpy as np -from PIL import Image - from mflux.community.in_context.flux_in_context_dev import Flux1InContextDev from mflux.community.in_context.utils.in_context_loras import LORA_REPO_ID, get_lora_filename from mflux.config.config import Config from mflux.config.model_config import ModelConfig +from .image_compare import check_images_close_enough + class ImageGeneratorInContextTestHelper: @staticmethod @@ -59,15 +58,15 @@ class ImageGeneratorInContextTestHelper: image.get_right_half().save(path=output_image_path, overwrite=True) # then - np.testing.assert_array_equal( - np.array(Image.open(output_image_path)), - np.array(Image.open(reference_image_path)), - err_msg=f"Generated image doesn't match reference image. Check {output_image_path} vs {reference_image_path}", + check_images_close_enough( + output_image_path, + reference_image_path, + "Generated in-context image doesn't match reference image.", ) finally: # cleanup - if os.path.exists(output_image_path): + if os.path.exists(output_image_path) and "MFLUX_PRESERVE_TEST_OUTPUT" not in os.environ: os.remove(output_image_path) @staticmethod diff --git a/tests/image_generation/helpers/image_generation_kontext_test_helper.py b/tests/image_generation/helpers/image_generation_kontext_test_helper.py index acd4e60..c3b6cdc 100644 --- a/tests/image_generation/helpers/image_generation_kontext_test_helper.py +++ b/tests/image_generation/helpers/image_generation_kontext_test_helper.py @@ -1,13 +1,12 @@ import os from pathlib import Path -import numpy as np -from PIL import Image - from mflux.config.config import Config from mflux.config.model_config import ModelConfig from mflux.kontext.flux_kontext import Flux1Kontext +from .image_compare import check_images_close_enough + class ImageGeneratorKontextTestHelper: @staticmethod @@ -46,18 +45,17 @@ class ImageGeneratorKontextTestHelper: image_path=kontext_image_path, ), ) - image.save(path=output_image_path) + image.save(path=output_image_path, overwrite=True) # then - np.testing.assert_array_equal( - np.array(Image.open(output_image_path)), - np.array(Image.open(reference_image_path)), - err_msg=f"Generated image doesn't match reference image. Check {output_image_path} vs {reference_image_path}", + check_images_close_enough( + output_image_path, + reference_image_path, + "Generated kontext image doesn't match reference image.", ) - finally: # cleanup - if os.path.exists(output_image_path): + if os.path.exists(output_image_path) and "MFLUX_PRESERVE_TEST_OUTPUT" not in os.environ: os.remove(output_image_path) @staticmethod diff --git a/tests/image_generation/helpers/image_generation_redux_test_helper.py b/tests/image_generation/helpers/image_generation_redux_test_helper.py index 5dc68b4..16f2dab 100644 --- a/tests/image_generation/helpers/image_generation_redux_test_helper.py +++ b/tests/image_generation/helpers/image_generation_redux_test_helper.py @@ -1,13 +1,12 @@ import os from pathlib import Path -import numpy as np -from PIL import Image - from mflux.config.config import Config from mflux.config.model_config import ModelConfig from mflux.flux_tools.redux.flux_redux import Flux1Redux +from .image_compare import check_images_close_enough + class ImageGeneratorReduxTestHelper: @staticmethod @@ -45,18 +44,18 @@ class ImageGeneratorReduxTestHelper: redux_image_paths=[redux_image_path], ), ) - image.save(path=output_image_path) + image.save(path=output_image_path, overwrite=True) # then - np.testing.assert_array_equal( - np.array(Image.open(output_image_path)), - np.array(Image.open(reference_image_path)), - err_msg=f"Generated image doesn't match reference image. Check {output_image_path} vs {reference_image_path}", + check_images_close_enough( + output_image_path, + reference_image_path, + "Generated redux image doesn't match reference image.", ) finally: # cleanup - if os.path.exists(output_image_path): + if os.path.exists(output_image_path) and "MFLUX_PRESERVE_TEST_OUTPUT" not in os.environ: os.remove(output_image_path) @staticmethod diff --git a/tests/image_generation/helpers/image_generation_test_helper.py b/tests/image_generation/helpers/image_generation_test_helper.py index ac63ba6..a1731df 100644 --- a/tests/image_generation/helpers/image_generation_test_helper.py +++ b/tests/image_generation/helpers/image_generation_test_helper.py @@ -1,13 +1,12 @@ import os from pathlib import Path -import numpy as np -from PIL import Image - from mflux.config.config import Config from mflux.config.model_config import ModelConfig from mflux.flux.flux import Flux1 +from .image_compare import check_images_close_enough + class ImageGeneratorTestHelper: @staticmethod @@ -53,15 +52,14 @@ class ImageGeneratorTestHelper: image.save(path=output_image_path, overwrite=True) # then - np.testing.assert_array_equal( - np.array(Image.open(output_image_path)), - np.array(Image.open(reference_image_path)), - err_msg=f"Generated image doesn't match reference image. Check {output_image_path} vs {reference_image_path}", + 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): + if os.path.exists(output_image_path) and "MFLUX_PRESERVE_TEST_OUTPUT" not in os.environ: os.remove(output_image_path) @staticmethod diff --git a/tests/resources/.gitignore b/tests/resources/.gitignore new file mode 100644 index 0000000..9e9c01f --- /dev/null +++ b/tests/resources/.gitignore @@ -0,0 +1 @@ +output_*.png diff --git a/tests/resources/reference_controlnet_dev_lora_depth.png b/tests/resources/reference_controlnet_dev_lora_depth.png index 0cd2da4..a3fe91c 100644 Binary files a/tests/resources/reference_controlnet_dev_lora_depth.png and b/tests/resources/reference_controlnet_dev_lora_depth.png differ