perf(nodes): compute the mask fade only on the fade band

invokeai/app/invocations/image.py

@@ -1348,64 +1348,68 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
             image_dto = context.images.save(image=pil_mask, image_category=ImageCategory.MASK)
             return ImageOutput.build(image_dto)
 
-        np_mask = numpy.array(pil_mask)
+        np_mask = numpy.asarray(pil_mask)
 
         # Threshold the mask to create a binary mask - 0 for black, 255 for white
         # If we don't threshold we can get some weird artifacts
-        np_mask = numpy.where(np_mask > self.threshold, 255, 0).astype(numpy.uint8)
+        _, np_result = cv2.threshold(np_mask, self.threshold, 255, cv2.THRESH_BINARY)
 
-        # Create a mask for the black region (1 where black, 0 otherwise)
-        black_mask = (np_mask == 0).astype(numpy.uint8)
+        # Create a distance transform of the thresholded mask. Distances are measured to the
+        # nearest black pixel, so the black region is exactly where dist == 0.
+        dist = cv2.distanceTransform(np_result, cv2.DIST_L2, 5)
 
-        # Invert the black region
-        bg_mask = 1 - black_mask
-
-        # Create a distance transform of the inverted mask
-        dist = cv2.distanceTransform(bg_mask, cv2.DIST_L2, 5)
-
-        # Normalize distances so that pixels <fade_size_px become a linear gradient (0 to 1)
-        d_norm = numpy.clip(dist / self.fade_size_px, 0, 1)
-
-        # Control points: x values (normalized distance) and corresponding fade pct y values.
-
-        # There are some magic numbers here that are used to create a smooth transition:
-        # - The first point is at 0% of fade size from edge of mask (meaning the edge of the mask), and is 0% fade (black)
-        # - The second point is 1px from the edge of the mask and also has 0% fade, effectively expanding the mask
-        #   by 1px. This fixes an issue where artifacts can occur at the edge of the mask
-        # - The third point is at 20% of the fade size from the edge of the mask and has 20% fade
-        # - The fourth point is at 80% of the fade size from the edge of the mask and has 90% fade
-        # - The last point is at 100% of the fade size from the edge of the mask and has 100% fade (white)
-
-        # x values: 0 = mask edge, 1 = fade_size_px from edge
-        x_control = numpy.array([0.0, 1.0 / self.fade_size_px, 0.2, 0.8, 1.0])
-        # y values: 0 = black, 1 = white
-        y_control = numpy.array([0.0, 0.0, 0.2, 0.9, 1.0])
-
-        # Fit a cubic polynomial that smoothly passes through the control points
-        coeffs = numpy.polyfit(x_control, y_control, 3)
-        poly = numpy.poly1d(coeffs)
-
-        # Evaluate the polynomial
-        feather = poly(d_norm)
+        # np_result is also the final image for everywhere except the fade band. The black
+        # region is already 0, and everything at or beyond the fade distance is already 255,
+        # which is what the forced 1.0 below works out to. Only the band still needs the
+        # polynomial, and for the default 16px fade on a 1024x1024 mask that is a few
+        # percent of the pixels rather than all of them.
 
         # The polynomial fit isn't perfect. Points beyond the fade distance are likely to be slightly less than 1.0,
         # even though the control points indicate that they should be exactly 1.0. This is due to the nature of the
         # polynomial fit, which is a best approximation of the control points but not an exact match.
 
         # When this occurs, the area outside the mask and fade-out will not be 100% transparent. For example, it may
         # have an alpha value of 1 instead of 0. So we must force pixels at or beyond the fade distance to exactly 1.0.
+        # That forced 1.0 is 255, which those pixels already hold, so they are left untouched.
+
+        # The fade band is the only part that still needs the polynomial. Building the mask
+        # in place avoids a second full-size temporary.
+        band = dist > 0
+        band &= dist < self.fade_size_px
+
+        if band.any():
+            # Control points: x values (normalized distance) and corresponding fade pct y values.
+
+            # There are some magic numbers here that are used to create a smooth transition:
+            # - The first point is at 0% of fade size from edge of mask (meaning the edge of the mask), and is 0% fade (black)
+            # - The second point is 1px from the edge of the mask and also has 0% fade, effectively expanding the mask
+            #   by 1px. This fixes an issue where artifacts can occur at the edge of the mask
+            # - The third point is at 20% of the fade size from the edge of the mask and has 20% fade
+            # - The fourth point is at 80% of the fade size from the edge of the mask and has 90% fade
+            # - The last point is at 100% of the fade size from the edge of the mask and has 100% fade (white)
+
+            # x values: 0 = mask edge, 1 = fade_size_px from edge
+            x_control = numpy.array([0.0, 1.0 / self.fade_size_px, 0.2, 0.8, 1.0])
+            # y values: 0 = black, 1 = white
+            y_control = numpy.array([0.0, 0.0, 0.2, 0.9, 1.0])
+
+            # Fit a cubic polynomial that smoothly passes through the control points
+            coeffs = numpy.polyfit(x_control, y_control, 3)
+            poly = numpy.poly1d(coeffs)
 
-        # Force pixels at or beyond the fade distance to exactly 1.0
-        feather = numpy.where(d_norm >= 1.0, 1.0, feather)
+            # Evaluate the polynomial on the band only. Normalizing just those distances is
+            # the same as normalizing the whole image and clipping, because 0 < dist < fade_size_px
+            # there. Calling the same poly() keeps the arithmetic identical to before rather
+            # than depending on how a hand-written expression promotes dtypes.
+            feather = poly(dist[band] / self.fade_size_px)
 
-        # Clip any other values to ensure they're in the valid range [0,1]
-        feather = numpy.clip(feather, 0, 1)
+            # Clip any other values to ensure they're in the valid range [0,1]
+            numpy.clip(feather, 0, 1, out=feather)
 
-        # Build final image.
-        np_result = numpy.where(black_mask == 1, 0, (feather * 255).astype(numpy.uint8))
+            np_result[band] = (feather * 255).astype(numpy.uint8)
 
         # Convert back to PIL, grayscale
-        pil_result = Image.fromarray(np_result.astype(numpy.uint8), mode="L")
+        pil_result = Image.fromarray(np_result, mode="L")
 
         image_dto = context.images.save(image=pil_result, image_category=ImageCategory.MASK)
 

tests/app/invocations/test_expand_mask_with_fade.py

@@ -0,0 +1,109 @@
+"""The faded mask must come out byte-for-byte the same as the original implementation.
+
+`ExpandMaskWithFadeInvocation` feeds canvas compositing, so a one-level change in the
+feather ramp shows up as a seam in a paste-back. This test pins the output against a
+transcription of the previous implementation rather than against stored fixtures, so it
+keeps working if the fixtures are ever regenerated.
+"""
+
+import cv2
+import numpy
+import pytest
+from PIL import Image
+
+from invokeai.app.invocations.fields import ImageField
+from invokeai.app.invocations.image import ExpandMaskWithFadeInvocation
+
+
+def _previous_implementation(pil_mask: Image.Image, threshold: int, fade_size_px: int) -> Image.Image:
+    """The pre-optimisation body of ExpandMaskWithFadeInvocation.invoke, verbatim."""
+    if fade_size_px == 0:
+        return pil_mask
+
+    np_mask = numpy.array(pil_mask)
+    np_mask = numpy.where(np_mask > threshold, 255, 0).astype(numpy.uint8)
+    black_mask = (np_mask == 0).astype(numpy.uint8)
+    bg_mask = 1 - black_mask
+    dist = cv2.distanceTransform(bg_mask, cv2.DIST_L2, 5)
+    d_norm = numpy.clip(dist / fade_size_px, 0, 1)
+    x_control = numpy.array([0.0, 1.0 / fade_size_px, 0.2, 0.8, 1.0])
+    y_control = numpy.array([0.0, 0.0, 0.2, 0.9, 1.0])
+    poly = numpy.poly1d(numpy.polyfit(x_control, y_control, 3))
+    feather = poly(d_norm)
+    feather = numpy.where(d_norm >= 1.0, 1.0, feather)
+    feather = numpy.clip(feather, 0, 1)
+    np_result = numpy.where(black_mask == 1, 0, (feather * 255).astype(numpy.uint8))
+    return Image.fromarray(np_result.astype(numpy.uint8), mode="L")
+
+
+class _Saved:
+    def __init__(self, image: Image.Image) -> None:
+        self.image_name = "mask.png"
+        self.width = image.width
+        self.height = image.height
+
+
+class _Images:
+    def __init__(self, source: Image.Image) -> None:
+        self.source = source
+        self.saved: Image.Image | None = None
+
+    def get_pil(self, image_name: str, mode=None) -> Image.Image:
+        image = self.source
+        if mode and mode != image.mode:
+            image = image.convert(mode)
+        return image
+
+    def save(self, image: Image.Image, image_category=None, **kwargs) -> _Saved:
+        self.saved = image
+        return _Saved(image)
+
+
+class _Context:
+    def __init__(self, images: _Images) -> None:
+        self.images = images
+
+
+def _mask(kind: str, height: int, width: int) -> Image.Image:
+    rng = numpy.random.default_rng(0)
+    array = numpy.zeros((height, width), dtype=numpy.uint8)
+    if kind == "blobs":
+        for _ in range(3):
+            cx = int(rng.integers(width // 4, 3 * width // 4))
+            cy = int(rng.integers(height // 4, 3 * height // 4))
+            cv2.ellipse(array, (cx, cy), (width // 6, height // 6), 0, 0, 360, 255, -1)
+    elif kind == "band":
+        array[:, : width // 3] = 255
+    elif kind == "empty":
+        pass
+    elif kind == "full":
+        array[:] = 255
+    elif kind == "hairline":
+        cv2.line(array, (0, height // 2), (width - 1, height // 2), 255, 1)
+    elif kind == "grey_ramp":
+        array[:] = numpy.linspace(0, 255, width, dtype=numpy.uint8)[None, :]
+    else:
+        raise AssertionError(kind)
+    return Image.fromarray(array, mode="L")
+
+
+@pytest.mark.parametrize("kind", ["blobs", "band", "empty", "full", "hairline", "grey_ramp"])
+@pytest.mark.parametrize("size", [(64, 64), (129, 97)])
+@pytest.mark.parametrize("threshold,fade_size_px", [(0, 32), (0, 1), (127, 8), (254, 64), (0, 0)])
+def test_matches_previous_implementation(kind, size, threshold, fade_size_px):
+    height, width = size
+    source = _mask(kind, height, width)
+
+    images = _Images(source)
+    node = ExpandMaskWithFadeInvocation(
+        mask=ImageField(image_name="mask.png"), threshold=threshold, fade_size_px=fade_size_px
+    )
+    output = node.invoke(_Context(images))
+
+    assert images.saved is not None
+    assert images.saved.mode == "L"
+    assert images.saved.size == source.size
+    assert (output.width, output.height) == source.size
+
+    expected = _previous_implementation(source, threshold, fade_size_px)
+    assert numpy.array_equal(numpy.array(images.saved), numpy.array(expected))