From 1d36dfb9fe025b716bc66d920b996381f457393d Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Tue, 26 Sep 2023 02:53:57 -0400 Subject: [PATCH] GrowMask now works with mask batches. --- comfy_extras/nodes_mask.py | 22 +++++++++++++--------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/comfy_extras/nodes_mask.py b/comfy_extras/nodes_mask.py index aa13cac0..af7cb07b 100644 --- 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() - while expand < 0: - output = scipy.ndimage.grey_erosion(output, footprint=kernel) - expand += 1 - while expand > 0: - output = scipy.ndimage.grey_dilation(output, footprint=kernel) - expand -= 1 - output = torch.from_numpy(output) - return (output,) + mask = mask.reshape((-1, mask.shape[-2], mask.shape[-1])) + out = [] + for m in mask: + output = m.numpy() + while expand < 0: + output = scipy.ndimage.grey_erosion(output, footprint=kernel) + expand += 1 + while expand > 0: + output = scipy.ndimage.grey_dilation(output, footprint=kernel) + expand -= 1 + output = torch.from_numpy(output) + out.append(output) + return (torch.cat(out, dim=0),)