diff --git a/nodes/batch_image_loader.py b/nodes/batch_image_loader.py new file mode 100644 index 0000000..f70c2b4 --- /dev/null +++ b/nodes/batch_image_loader.py @@ -0,0 +1,167 @@ +# Batch Image Loader nodes — part of ComfyUI-AI2Go-Utils. GPL-3.0. +# +# Two loaders that emit one output socket group per dropped file (design: +# docs/superpowers/specs/2026-08-12-batch-loaders-design.md). Simple = image_N + count; +# Advanced adds mask_N + filename_N. Outputs are declared at the MAX_FILES ceiling and the +# front-end (web/js/batch_loader.js) trims the unused tail — slot order is count first, then +# grouped per file, so used slots stay contiguous (ComfyUI validates output types by position). +# Mask quirks copied from stock LoadImage: mask = 1.0 - alpha; no alpha -> 64x64 zeros. +import os + +import numpy as np +import torch +from PIL import Image, ImageOps + +import folder_paths +from comfy_api.latest import io + +from .batch_loader_core import DOWNSCALE_MODES, MAX_FILES, downscale_size, parse_files + + +def _input_path(f): + """Resolve a files_json entry to a path inside the input directory (traversal-safe).""" + base = os.path.abspath(folder_paths.get_input_directory()) + path = os.path.abspath(os.path.join(base, f["subfolder"], f["name"])) + if os.path.commonpath([base, path]) != base: + raise ValueError(f"File path escapes the input folder: {f['name']!r}") + return path + + +def _load_image(path, mode, max_size): + """Load one image -> (IMAGE [1,H,W,3], MASK). Applies downscale geometry to both.""" + try: + img = Image.open(path) + img = ImageOps.exif_transpose(img) + except Exception as e: + raise ValueError(f"Could not read image {os.path.basename(path)!r}: {e}") from e + + alpha = img.getchannel("A") if "A" in img.getbands() else None + rgb = img.convert("RGB") + + (tw, th), box = downscale_size(rgb.width, rgb.height, mode, max_size) + if box is not None: + rgb = rgb.crop(box) + alpha = alpha.crop(box) if alpha is not None else None + if (tw, th) != rgb.size: + rgb = rgb.resize((tw, th), Image.LANCZOS) + alpha = alpha.resize((tw, th), Image.LANCZOS) if alpha is not None else None + + image = torch.from_numpy(np.array(rgb).astype(np.float32) / 255.0)[None,] + if alpha is not None: + mask = 1.0 - torch.from_numpy(np.array(alpha).astype(np.float32) / 255.0) + else: + mask = torch.zeros((64, 64), dtype=torch.float32) # stock LoadImage placeholder + return image, mask.unsqueeze(0) + + +def _run_image(files_json, downscale_mode, max_size, advanced): + """Shared engine for both nodes. Returns the padded output list (count first).""" + files = parse_files(files_json) + group = 3 if advanced else 1 + outputs = [None] * (1 + MAX_FILES * group) + outputs[0] = len(files) + for i, f in enumerate(files): + path = _input_path(f) + if not os.path.isfile(path): + raise ValueError(f"File not found in the input folder: {f['name']!r} — re-add it to the node.") + image, mask = _load_image(path, downscale_mode, max_size) + base = 1 + i * group + outputs[base] = image + if advanced: + outputs[base + 1] = mask + outputs[base + 2] = f["name"] + return outputs + + +def _fingerprint(files_json, downscale_mode, max_size): + """Cache key: file list + size/mtime per file, so an on-disk change re-runs the node.""" + try: + files = parse_files(files_json) + except ValueError: + return files_json + sig = [str(downscale_mode), str(max_size)] + for f in files: + try: + st = os.stat(_input_path(f)) + sig.append(f"{f['subfolder']}/{f['name']}:{st.st_size}:{st.st_mtime_ns}") + except (OSError, ValueError): + sig.append(f"{f['subfolder']}/{f['name']}:missing") + return "|".join(sig) + + +def _image_inputs(): + return [ + io.String.Input( + "files_json", default="[]", + tooltip="Authoritative file list as JSON. Hidden in the UI and kept in sync by the " + "front-end — drop files onto the node instead of editing this.", + ), + io.Combo.Input( + "downscale_mode", options=list(DOWNSCALE_MODES), default="off", + tooltip="off = images pass through untouched. fit = shrink keeping shape so the longest " + "edge is max_size. crop = centered square, at most max_size. stretch = force a " + "max_size square (ignores shape). Never upscales (fit/crop); masks follow their image.", + ), + io.Int.Input( + "max_size", default=1200, min=8, max=8192, step=8, + tooltip="Size threshold/target in pixels for downscaling. Ignored while downscale_mode is off.", + ), + ] + + +def _image_outputs(advanced): + outs = [io.Int.Output(display_name="count")] + for i in range(1, MAX_FILES + 1): + outs.append(io.Image.Output(display_name=f"image_{i}")) + if advanced: + outs.append(io.Mask.Output(display_name=f"mask_{i}")) + outs.append(io.String.Output(display_name=f"filename_{i}")) + return outs + + +class AI2GoBatchImageLoader(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="AI2GoBatchImageLoader", + display_name="AI2Go Batch Image Loader", + category="AI2Go/image", + search_aliases=["batch", "load", "images", "multi", "drop", "upload"], + description="Drop up to 8 images onto the node; each gets its own image_N output " + "socket (sockets appear/disappear with the list). Optional downscaling: " + "set downscale_mode to fit/crop/stretch and images larger than max_size " + "shrink on load. count = number of files loaded.", + inputs=_image_inputs(), + outputs=_image_outputs(advanced=False), + ) + + @classmethod + def execute(cls, files_json="[]", downscale_mode="off", max_size=1200) -> io.NodeOutput: + return io.NodeOutput(*_run_image(files_json, downscale_mode, max_size, advanced=False)) + + @classmethod + def fingerprint_inputs(cls, files_json="[]", downscale_mode="off", max_size=1200): + return _fingerprint(files_json, downscale_mode, max_size) + + +class AI2GoBatchImageLoaderAdvanced(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="AI2GoBatchImageLoaderAdvanced", + display_name="AI2Go Batch Image Loader Advanced", + category="AI2Go/image", + search_aliases=["batch", "load", "images", "multi", "drop", "upload", "mask", "filename"], + description="Batch Image Loader plus a mask_N (inverted alpha, stock LoadImage " + "behavior) and filename_N output per file.", + inputs=_image_inputs(), + outputs=_image_outputs(advanced=True), + ) + + @classmethod + def execute(cls, files_json="[]", downscale_mode="off", max_size=1200) -> io.NodeOutput: + return io.NodeOutput(*_run_image(files_json, downscale_mode, max_size, advanced=True)) + + @classmethod + def fingerprint_inputs(cls, files_json="[]", downscale_mode="off", max_size=1200): + return _fingerprint(files_json, downscale_mode, max_size) diff --git a/tests/test_batch_image_loader.py b/tests/test_batch_image_loader.py new file mode 100644 index 0000000..ae88929 --- /dev/null +++ b/tests/test_batch_image_loader.py @@ -0,0 +1,54 @@ +# Integration test for the Batch Image Loader nodes — part of ComfyUI-AI2Go-Utils. GPL-3.0. +import json + +import pytest + +pytest.importorskip("comfy_api") # only runs inside a ComfyUI environment +import torch # noqa: E402 +from PIL import Image # noqa: E402 + + +@pytest.fixture +def input_dir(tmp_path, monkeypatch): + import folder_paths + monkeypatch.setattr(folder_paths, "get_input_directory", lambda: str(tmp_path)) + return tmp_path + + +def test_advanced_images_masks_filenames(input_dir): + from nodes.batch_image_loader import _run_image + Image.new("RGBA", (32, 16), (255, 0, 0, 128)).save(input_dir / "a.png") + Image.new("RGB", (8, 8), (0, 255, 0)).save(input_dir / "b.png") + out = _run_image(json.dumps([{"name": "a.png"}, {"name": "b.png"}]), "off", 1200, advanced=True) + + assert len(out) == 25 and out[0] == 2 + img1, mask1, name1 = out[1], out[2], out[3] + assert img1.shape == (1, 16, 32, 3) # [1,H,W,C] + assert mask1.shape == (1, 16, 32) # inverted alpha, image-sized + assert torch.allclose(mask1, torch.full_like(mask1, 1.0 - 128 / 255), atol=1e-3) + assert name1 == "a.png" + img2, mask2 = out[4], out[5] + assert img2.shape == (1, 8, 8, 3) + assert mask2.shape == (1, 64, 64) and mask2.max() == 0 # no alpha -> stock 64x64 zeros + assert out[7] is None # padding starts after file 2's group + + +def test_simple_layout(input_dir): + from nodes.batch_image_loader import _run_image + Image.new("RGB", (8, 8), (0, 0, 255)).save(input_dir / "c.png") + out = _run_image(json.dumps([{"name": "c.png"}]), "off", 1200, advanced=False) + assert len(out) == 9 and out[0] == 1 + assert out[1].shape == (1, 8, 8, 3) and out[2] is None + + +def test_downscale_fit_applies(input_dir): + from nodes.batch_image_loader import _run_image + Image.new("RGB", (2400, 1200), (9, 9, 9)).save(input_dir / "big.png") + out = _run_image(json.dumps([{"name": "big.png"}]), "fit", 1200, advanced=False) + assert out[1].shape == (1, 600, 1200, 3) + + +def test_missing_file_raises(input_dir): + from nodes.batch_image_loader import _run_image + with pytest.raises(ValueError, match="gone.png"): + _run_image(json.dumps([{"name": "gone.png"}]), "off", 1200, advanced=False)