GrowMask now works with mask batches.

1d36dfb9 · comfyanonymous · d76d71de · 1d36dfb9
Commit 1d36dfb9 authored Sep 26, 2023 by comfyanonymous
Hide whitespace changes
Inline Side-by-side

Showing with 13 additions and 9 deletions

comfy_extras/nodes_mask.py comfy_extras/nodes_mask.py +13 -9

No files found.
--- a/comfy_extras/nodes_mask.py
+++ b/comfy_extras/nodes_mask.py
@@ -327,15 +327,19 @@ class GrowMask:
        kernel = np.array([[c, 1, c],
                           [1, 1, 1],
                           [c, 1, c]])
-        output = mask.numpy().copy()
+        mask = mask.reshape((-1, mask.shape[-2], mask.shape[-1]))
-        while expand < 0:
+        out = []
-            output = scipy.ndimage.grey_erosion(output, footprint=kernel)
+        for m in mask:
-            expand += 1
+            output = m.numpy()
-        while expand > 0:
+            while expand < 0:
-            output = scipy.ndimage.grey_dilation(output, footprint=kernel)
+                output = scipy.ndimage.grey_erosion(output, footprint=kernel)
-            expand -= 1
+                expand += 1
-        output = torch.from_numpy(output)
+            while expand > 0:
-        return (output,)
+                output = scipy.ndimage.grey_dilation(output, footprint=kernel)
+                expand -= 1
+            output = torch.from_numpy(output)
+            out.append(output)
+        return (torch.cat(out, dim=0),)