""" 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", }