added the main function

This commit is contained in:
Balaji-ADOLF
2026-04-02 14:32:46 +05:30
parent 9a6220d991
commit e4833300eb
2 changed files with 248 additions and 0 deletions
+3
View File
@@ -0,0 +1,3 @@
from .multi_image_loader import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
+245
View File
@@ -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",
}