90 lines
2.9 KiB
Python
90 lines
2.9 KiB
Python
# ======================================================================
|
|
# Node Pack: artyclaw-comfy
|
|
# Author: artyclaw
|
|
# Version: 1.0.0
|
|
# Description: A power-user toolkit for prompt logic, IO, and RNG.
|
|
# ======================================================================
|
|
|
|
import os
|
|
import torch
|
|
import numpy as np
|
|
from PIL import Image, ImageOps
|
|
|
|
class SimpleDirLoader:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"directory": ("STRING", {
|
|
"default": "",
|
|
"multiline": False,
|
|
"tooltip": "Absolute path to the image folder."
|
|
}),
|
|
},
|
|
"optional": {
|
|
"limit": ("INT", {"default": 0, "min": 0, "max": 9999, "tooltip": "Limit number of images (0 = all)."}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "STRING")
|
|
RETURN_NAMES = ("image_list", "path_list")
|
|
FUNCTION = "load_images"
|
|
CATEGORY = "Custom/Loading"
|
|
|
|
# CRITICAL CHANGE: This tells ComfyUI "The output is a LIST of separate items, not one big item"
|
|
OUTPUT_IS_LIST = (True, True)
|
|
|
|
def load_images(self, directory, limit):
|
|
if not os.path.isdir(directory):
|
|
raise FileNotFoundError(f"Directory not found: {directory}")
|
|
|
|
valid_exts = ('.png', '.jpg', '.jpeg', '.webp', '.bmp', '.tiff')
|
|
files = [f for f in os.listdir(directory) if f.lower().endswith(valid_exts)]
|
|
files.sort()
|
|
|
|
if limit > 0:
|
|
files = files[:limit]
|
|
|
|
if not files:
|
|
raise ValueError("No images found in the specified directory.")
|
|
|
|
images = []
|
|
paths = []
|
|
|
|
print(f"[SimpleDirLoader] Loading {len(files)} images as a LIST...")
|
|
|
|
for file in files:
|
|
full_path = os.path.join(directory, file)
|
|
try:
|
|
img = Image.open(full_path)
|
|
img = ImageOps.exif_transpose(img)
|
|
img = img.convert("RGB")
|
|
|
|
# Normalize to 0-1
|
|
img_np = np.array(img).astype(np.float32) / 255.0
|
|
|
|
# Create a Tensor for THIS SINGLE IMAGE [1, H, W, C]
|
|
# We do NOT resize it. It keeps its native H/W.
|
|
img_tensor = torch.from_numpy(img_np)[None,]
|
|
|
|
images.append(img_tensor)
|
|
paths.append(full_path)
|
|
|
|
except Exception as e:
|
|
print(f" [!] Error loading '{file}': {e}")
|
|
continue
|
|
|
|
# We return the Python LIST of tensors, not a single torch.cat() tensor.
|
|
return (images, paths)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"SimpleDirLoader": SimpleDirLoader
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"LoadImages": "Load Images from Folder (List Mode)"
|
|
}
|