From 1c8ae9dbb249ed5326d61d16b4e6b5807c09c0e1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jukka=20Sepp=C3=A4nen?= <40791699+kijai@users.noreply.github.com> Date: Fri, 29 Sep 2023 05:01:19 +0300 Subject: [PATCH] Allow GrowMask node to work with batches (for AnimateDiff) (#1623) * Allow mask batches This allows LatentCompositeMasked -node to work with AnimateDiff. I tried to keep old functionality too, unsure if it's correct, but both single mask and batch of masks seems to work with this change. * Update nodes_mask.py --- comfy_extras/nodes_mask.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/comfy_extras/nodes_mask.py b/comfy_extras/nodes_mask.py index af7cb07bf..cdf762ffd 100644 --- a/comfy_extras/nodes_mask.py +++ b/comfy_extras/nodes_mask.py @@ -331,15 +331,14 @@ class GrowMask: 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 + for _ in range(abs(expand)): + if expand < 0: + output = scipy.ndimage.grey_erosion(output, footprint=kernel) + else: + output = scipy.ndimage.grey_dilation(output, footprint=kernel) output = torch.from_numpy(output) out.append(output) - return (torch.cat(out, dim=0),) + return (torch.stack(out, dim=0),)