246 lines
11 KiB
Python
246 lines
11 KiB
Python
"""
|
||
Multi Image Loader Node for ComfyUI
|
||
====================================
|
||
Supports loading 2–8 images dynamically with optional inputs.
|
||
Each loaded image is:
|
||
- Individually accessible as image1..imageN outputs
|
||
- Also available as a batched tensor (all images stacked)
|
||
|
||
Designed to work with TextEncodeQwenImageEditPlus which accepts image1 and image2.
|
||
You can route any individual output to image1 / image2 slots.
|
||
"""
|
||
|
||
import os
|
||
import torch
|
||
import numpy as np
|
||
from PIL import Image, ImageOps, ImageSequence
|
||
import folder_paths
|
||
|
||
|
||
# ──────────────────────────────────────────────────────────────────────────────
|
||
# Helper: load a single image file → (H, W, C) float32 tensor [0,1]
|
||
# ──────────────────────────────────────────────────────────────────────────────
|
||
def load_image_tensor(image_path: str) -> torch.Tensor:
|
||
img = Image.open(image_path)
|
||
img = ImageOps.exif_transpose(img)
|
||
if img.mode == "I":
|
||
img = img.point(lambda i: i * (1 / 255))
|
||
img = img.convert("RGB")
|
||
arr = np.array(img).astype(np.float32) / 255.0
|
||
return torch.from_numpy(arr).unsqueeze(0) # (1, H, W, C)
|
||
|
||
|
||
# ──────────────────────────────────────────────────────────────────────────────
|
||
# Node: MultiImageLoader
|
||
# ──────────────────────────────────────────────────────────────────────────────
|
||
class MultiImageLoader:
|
||
"""
|
||
Load 2–8 images. Each slot is optional beyond the first two.
|
||
Outputs:
|
||
• image_1 … image_N — individual tensors (pass to image1 / image2 slots)
|
||
• image_batch — all loaded images concatenated along batch dim
|
||
• loaded_count — INT: how many images were actually loaded
|
||
"""
|
||
|
||
MAX_IMAGES = 8
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
image_files = folder_paths.get_filename_list("input")
|
||
|
||
inputs = {
|
||
"required": {
|
||
"image_1": (sorted(image_files),),
|
||
"image_2": (sorted(image_files),),
|
||
},
|
||
"optional": {}
|
||
}
|
||
|
||
# image_3 … image_8 are optional
|
||
for i in range(3, cls.MAX_IMAGES + 1):
|
||
inputs["optional"][f"image_{i}"] = (
|
||
["__none__"] + sorted(image_files),
|
||
)
|
||
|
||
return inputs
|
||
|
||
RETURN_TYPES = ("IMAGE",) * MAX_IMAGES + ("IMAGE", "INT")
|
||
RETURN_NAMES = tuple(f"image_{i}" for i in range(1, MAX_IMAGES + 1)) + \
|
||
("image_batch", "loaded_count")
|
||
FUNCTION = "load_images"
|
||
CATEGORY = "image/multi"
|
||
OUTPUT_NODE = False
|
||
|
||
def load_images(self, image_1, image_2, **kwargs):
|
||
tensors = []
|
||
|
||
# Always load required slots
|
||
for name in ("image_1", "image_2"):
|
||
val = image_1 if name == "image_1" else image_2
|
||
path = folder_paths.get_annotated_filepath(val)
|
||
tensors.append(load_image_tensor(path))
|
||
|
||
# Optional slots
|
||
for i in range(3, self.MAX_IMAGES + 1):
|
||
val = kwargs.get(f"image_{i}", "__none__")
|
||
if val and val != "__none__":
|
||
path = folder_paths.get_annotated_filepath(val)
|
||
tensors.append(load_image_tensor(path))
|
||
else:
|
||
tensors.append(None)
|
||
|
||
loaded_count = sum(1 for t in tensors if t is not None)
|
||
|
||
# For individual outputs: if slot is empty, fall back to image_1
|
||
individual = []
|
||
for t in tensors:
|
||
individual.append(t if t is not None else tensors[0])
|
||
|
||
# Batch: only loaded images
|
||
loaded_tensors = [t for t in tensors if t is not None]
|
||
if len(loaded_tensors) == 1:
|
||
batch = loaded_tensors[0]
|
||
else:
|
||
# Resize all to match image_1 dimensions before stacking
|
||
h, w = loaded_tensors[0].shape[1], loaded_tensors[0].shape[2]
|
||
resized = []
|
||
for t in loaded_tensors:
|
||
if t.shape[1] != h or t.shape[2] != w:
|
||
# Use PIL to resize
|
||
arr = (t[0].numpy() * 255).astype(np.uint8)
|
||
img = Image.fromarray(arr).resize((w, h), Image.LANCZOS)
|
||
arr_r = np.array(img).astype(np.float32) / 255.0
|
||
t = torch.from_numpy(arr_r).unsqueeze(0)
|
||
resized.append(t)
|
||
batch = torch.cat(resized, dim=0)
|
||
|
||
return tuple(individual) + (batch, loaded_count)
|
||
|
||
|
||
# ──────────────────────────────────────────────────────────────────────────────
|
||
# Node: ImageSelectFromBatch
|
||
# ──────────────────────────────────────────────────────────────────────────────
|
||
class ImageSelectFromBatch:
|
||
"""
|
||
Pick a single image from a batch by index.
|
||
Useful for routing a specific image from image_batch into image1/image2 slots.
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"image_batch": ("IMAGE",),
|
||
"index": ("INT", {"default": 0, "min": 0, "max": 63, "step": 1}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE",)
|
||
RETURN_NAMES = ("image",)
|
||
FUNCTION = "select"
|
||
CATEGORY = "image/multi"
|
||
|
||
def select(self, image_batch, index):
|
||
count = image_batch.shape[0]
|
||
idx = min(index, count - 1)
|
||
return (image_batch[idx:idx+1],)
|
||
|
||
|
||
# ──────────────────────────────────────────────────────────────────────────────
|
||
# Node: ImageBatchBuilder (progressive: add one image at a time to a batch)
|
||
# ──────────────────────────────────────────────────────────────────────────────
|
||
class ImageBatchBuilder:
|
||
"""
|
||
Append a new image onto an existing batch.
|
||
Chain multiple of these to build a batch of any size.
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"image_new": ("IMAGE",),
|
||
},
|
||
"optional": {
|
||
"image_batch": ("IMAGE",),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE", "INT")
|
||
RETURN_NAMES = ("image_batch", "batch_size")
|
||
FUNCTION = "append"
|
||
CATEGORY = "image/multi"
|
||
|
||
def append(self, image_new, image_batch=None):
|
||
if image_batch is None:
|
||
result = image_new
|
||
else:
|
||
# Resize image_new to match batch spatial dims if needed
|
||
h, w = image_batch.shape[1], image_batch.shape[2]
|
||
if image_new.shape[1] != h or image_new.shape[2] != w:
|
||
arr = (image_new[0].numpy() * 255).astype(np.uint8)
|
||
img = Image.fromarray(arr).resize((w, h), Image.LANCZOS)
|
||
arr_r = np.array(img).astype(np.float32) / 255.0
|
||
image_new = torch.from_numpy(arr_r).unsqueeze(0)
|
||
result = torch.cat([image_batch, image_new], dim=0)
|
||
|
||
return (result, result.shape[0])
|
||
|
||
|
||
# ──────────────────────────────────────────────────────────────────────────────
|
||
# Node: QwenMultiRefEncoder
|
||
# ──────────────────────────────────────────────────────────────────────────────
|
||
class QwenMultiRefEncoder:
|
||
"""
|
||
Convenience wrapper: takes a batch of reference images and a prompt,
|
||
and calls TextEncodeQwenImageEditPlus using image_1 and image_2 from the batch.
|
||
For >2 references you would chain multiple encoders or concatenate prompts.
|
||
|
||
This node auto-selects which two images from your batch to use as
|
||
image1 and image2 for conditioning.
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"prompt": ("STRING", {"multiline": True, "default": ""}),
|
||
"clip": ("CLIP",),
|
||
"vae": ("VAE",),
|
||
"image_batch": ("IMAGE",),
|
||
"ref_index_1": ("INT", {"default": 0, "min": 0, "max": 7, "step": 1,
|
||
"tooltip": "Batch index to use as image1"}),
|
||
"ref_index_2": ("INT", {"default": 1, "min": 0, "max": 7, "step": 1,
|
||
"tooltip": "Batch index to use as image2"}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE", "IMAGE")
|
||
RETURN_NAMES = ("selected_image1", "selected_image2")
|
||
FUNCTION = "pick_refs"
|
||
CATEGORY = "image/multi"
|
||
|
||
def pick_refs(self, prompt, clip, vae, image_batch, ref_index_1, ref_index_2):
|
||
count = image_batch.shape[0]
|
||
i1 = min(ref_index_1, count - 1)
|
||
i2 = min(ref_index_2, count - 1)
|
||
return (image_batch[i1:i1+1], image_batch[i2:i2+1])
|
||
|
||
|
||
# ──────────────────────────────────────────────────────────────────────────────
|
||
# Registration
|
||
# ──────────────────────────────────────────────────────────────────────────────
|
||
NODE_CLASS_MAPPINGS = {
|
||
"MultiImageLoader": MultiImageLoader,
|
||
"ImageSelectFromBatch": ImageSelectFromBatch,
|
||
"ImageBatchBuilder": ImageBatchBuilder,
|
||
"QwenMultiRefEncoder": QwenMultiRefEncoder,
|
||
}
|
||
|
||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||
"MultiImageLoader": "Multi Image Loader (2–8)",
|
||
"ImageSelectFromBatch": "Image Select From Batch",
|
||
"ImageBatchBuilder": "Image Batch Builder",
|
||
"QwenMultiRefEncoder": "Qwen Multi-Ref Picker",
|
||
}
|