added the main function
This commit is contained in:
@@ -0,0 +1,3 @@
|
||||
from .multi_image_loader import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
||||
@@ -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",
|
||||
}
|
||||
Reference in New Issue
Block a user