From e4833300eb970cfce468bef0716558a28f447191 Mon Sep 17 00:00:00 2001 From: Balaji-ADOLF Date: Thu, 2 Apr 2026 14:32:46 +0530 Subject: [PATCH] added the main function --- __init__.py | 3 + multi_image_loader.py | 245 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 248 insertions(+) create mode 100644 __init__.py create mode 100644 multi_image_loader.py diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..c01f17f --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ +from .multi_image_loader import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/multi_image_loader.py b/multi_image_loader.py new file mode 100644 index 0000000..7f7fd95 --- /dev/null +++ b/multi_image_loader.py @@ -0,0 +1,245 @@ +""" +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", +}