Skip to content

Commit 025e679

Browse files
authored
Batch broadcasting in JoinImageWithAlpha node (Comfy-Org#13686)
* Batch broadcasting in JoinImageWithAlpha node
1 parent 867b8d2 commit 025e679

1 file changed

Lines changed: 4 additions & 7 deletions

File tree

comfy_extras/nodes_compositing.py

Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -202,14 +202,11 @@ def define_schema(cls):
202202

203203
@classmethod
204204
def execute(cls, image: torch.Tensor, alpha: torch.Tensor) -> io.NodeOutput:
205-
batch_size = min(len(image), len(alpha))
206-
out_images = []
207-
205+
batch_size = max(len(image), len(alpha))
208206
alpha = 1.0 - resize_mask(alpha, image.shape[1:])
209-
for i in range(batch_size):
210-
out_images.append(torch.cat((image[i][:,:,:3], alpha[i].unsqueeze(2)), dim=2))
211-
212-
return io.NodeOutput(torch.stack(out_images))
207+
alpha = comfy.utils.repeat_to_batch_size(alpha, batch_size)
208+
image = comfy.utils.repeat_to_batch_size(image, batch_size)
209+
return io.NodeOutput(torch.cat((image[..., :3], alpha.unsqueeze(-1)), dim=-1))
213210

214211

215212
class CompositingExtension(ComfyExtension):

0 commit comments

Comments
 (0)