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
This commit is contained in:
parent
66756de100
commit
1c8ae9dbb2
|
@ -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),)
|
||||
|
||||
|
||||
|
||||
|
|
Loading…
Reference in New Issue