Skip to content

Commit

Permalink
Add a RebatchImages node.
Browse files Browse the repository at this point in the history
  • Loading branch information
comfyanonymous committed Dec 20, 2023
1 parent e82942c commit 5f54614
Showing 1 changed file with 31 additions and 1 deletion.
32 changes: 31 additions & 1 deletion comfy_extras/nodes_rebatch.py
Original file line number Diff line number Diff line change
Expand Up @@ -99,10 +99,40 @@ def rebatch(self, latents, batch_size):

return (output_list,)

class ImageRebatch:
@classmethod
def INPUT_TYPES(s):
return {"required": { "images": ("IMAGE",),
"batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}),
}}
RETURN_TYPES = ("IMAGE",)
INPUT_IS_LIST = True
OUTPUT_IS_LIST = (True, )

FUNCTION = "rebatch"

CATEGORY = "image/batch"

def rebatch(self, images, batch_size):
batch_size = batch_size[0]

output_list = []
all_images = []
for img in images:
for i in range(img.shape[0]):
all_images.append(img[i:i+1])

for i in range(0, len(all_images), batch_size):
output_list.append(torch.cat(all_images[i:i+batch_size], dim=0))

return (output_list,)

NODE_CLASS_MAPPINGS = {
"RebatchLatents": LatentRebatch,
"RebatchImages": ImageRebatch,
}

NODE_DISPLAY_NAME_MAPPINGS = {
"RebatchLatents": "Rebatch Latents",
}
"RebatchImages": "Rebatch Images",
}

0 comments on commit 5f54614

Please sign in to comment.