Update tests/image_generation image comparison logic to allow similar images that are _close enough_ (#263)

Co-authored-by: Anthony Wu <pls-file-gh-issue@users.noreply.github.com>
Co-authored-by: filipstrand <strand.filip@gmail.com>

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]

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=}"
+        )

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

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)

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

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)

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

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

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

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

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

tests/resources/.gitignore

@@ -0,0 +1 @@
+output_*.png