shootthesound
840d1dbfaa
New STRING output (appended last so existing workflows keep their link slots) carrying the absolute path of the file loaded this run, alongside the existing filename output. Bump version to 1.4.0. Closes #1.
297 lines
10 KiB
Python
297 lines
10 KiB
Python
"""Sequential Image Loader — load PNGs from a folder, one per queue.
|
|
|
|
Paste a directory path into the node; it scans for PNG files, sorts them
|
|
in natural order, and on each Queue press hands back the *next* image in
|
|
the sequence. A "Reset to start" button (frontend) rewinds to the first
|
|
file. Index advancing is driven by the JS queue hook; this module just
|
|
resolves index → file → tensors and exposes a listing route the UI uses
|
|
to show progress.
|
|
"""
|
|
|
|
import io
|
|
import os
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image, ImageOps
|
|
|
|
|
|
# --- Shared listing helpers -----------------------------------------------
|
|
|
|
# Selectable filetypes -> the extensions each option matches. "jpg" and
|
|
# "jpeg" both match .jpg/.jpeg (same format) so either choice behaves as
|
|
# users expect. "all" matches every supported image extension.
|
|
FILETYPE_EXTS = {
|
|
"png": (".png",),
|
|
"jpg": (".jpg", ".jpeg"),
|
|
"jpeg": (".jpg", ".jpeg"),
|
|
"webp": (".webp",),
|
|
"bmp": (".bmp",),
|
|
"gif": (".gif",),
|
|
"tiff": (".tif", ".tiff"),
|
|
}
|
|
FILETYPE_CHOICES = ["png", "jpg", "jpeg", "webp", "bmp", "gif", "tiff", "all"]
|
|
_ALL_EXTS = tuple(sorted({e for exts in FILETYPE_EXTS.values() for e in exts}))
|
|
|
|
|
|
def _exts_for(filetype):
|
|
"""Resolve a filetype choice to the tuple of extensions to match."""
|
|
if not filetype or filetype == "all":
|
|
return _ALL_EXTS
|
|
return FILETYPE_EXTS.get(filetype, (".png",))
|
|
|
|
|
|
def _natural_key(name):
|
|
"""Sort key so 'img2.png' precedes 'img10.png' (numeric-aware)."""
|
|
import re
|
|
parts = re.split(r"(\d+)", name.lower())
|
|
return [int(p) if p.isdigit() else p for p in parts]
|
|
|
|
|
|
def list_images(directory, include_subfolders=False, filetype="png",
|
|
reverse=False):
|
|
"""Return image file paths under `directory`, natural-sorted.
|
|
|
|
`filetype` is one of FILETYPE_CHOICES (or "all"). With
|
|
include_subfolders the walk is recursive and entries are sorted by
|
|
their path relative to `directory` so ordering is stable/intuitive.
|
|
`reverse` flips the final order (so the sequence runs last->first).
|
|
Returns [] for a missing/empty/invalid directory (never raises).
|
|
"""
|
|
directory = (directory or "").strip().strip('"').strip("'")
|
|
if not directory or not os.path.isdir(directory):
|
|
return []
|
|
|
|
exts = _exts_for(filetype)
|
|
|
|
def _matches(name):
|
|
return name.lower().endswith(exts)
|
|
|
|
found = []
|
|
if include_subfolders:
|
|
for root, _dirs, files in os.walk(directory):
|
|
for f in files:
|
|
if _matches(f):
|
|
found.append(os.path.join(root, f))
|
|
found.sort(key=lambda p: _natural_key(os.path.relpath(p, directory)))
|
|
else:
|
|
try:
|
|
names = os.listdir(directory)
|
|
except OSError:
|
|
return []
|
|
names = [n for n in names if _matches(n)
|
|
and os.path.isfile(os.path.join(directory, n))]
|
|
names.sort(key=_natural_key)
|
|
found = [os.path.join(directory, n) for n in names]
|
|
|
|
if reverse:
|
|
found.reverse()
|
|
return found
|
|
|
|
|
|
def _load_image_tensors(path):
|
|
"""Mirror ComfyUI's LoadImage: -> (IMAGE [1,H,W,3], MASK [1,H,W])."""
|
|
img = Image.open(path)
|
|
img = ImageOps.exif_transpose(img)
|
|
|
|
rgb = img.convert("RGB")
|
|
arr = np.array(rgb).astype(np.float32) / 255.0
|
|
image = torch.from_numpy(arr)[None,]
|
|
|
|
if "A" in img.getbands():
|
|
alpha = np.array(img.getchannel("A")).astype(np.float32) / 255.0
|
|
mask = 1.0 - torch.from_numpy(alpha)
|
|
else:
|
|
mask = torch.zeros((rgb.height, rgb.width), dtype=torch.float32)
|
|
mask = mask[None,]
|
|
return image, mask
|
|
|
|
|
|
# --- Node ------------------------------------------------------------------
|
|
|
|
class SequentialImageLoader:
|
|
"""Load PNGs from a folder in order, advancing one image per queue."""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"directory": ("STRING", {
|
|
"default": "",
|
|
"multiline": False,
|
|
"placeholder": "C:\\path\\to\\folder",
|
|
"tooltip": "Folder to scan for .png files.",
|
|
}),
|
|
"include_subfolders": ("BOOLEAN", {"default": False}),
|
|
# Mechanical: driven by the JS queue hook / Reset button.
|
|
"index": ("INT", {
|
|
"default": 0, "min": 0, "max": 0x7FFFFFFF, "step": 1,
|
|
}),
|
|
# Declared LAST (newest widgets) so older saved workflows keep
|
|
# their positional widget order; absent -> defaults below.
|
|
"filetype": (FILETYPE_CHOICES, {"default": "png"}),
|
|
"reverse": ("BOOLEAN", {
|
|
"default": False,
|
|
"label_on": "reverse (last→first)",
|
|
"label_off": "forward (first→last)",
|
|
}),
|
|
}
|
|
}
|
|
|
|
# file_path is declared last (newest output) so links in older saved
|
|
# workflows keep their slot positions.
|
|
RETURN_TYPES = ("IMAGE", "MASK", "STRING", "INT", "INT", "STRING")
|
|
RETURN_NAMES = ("image", "mask", "filename", "index", "total", "file_path")
|
|
FUNCTION = "load"
|
|
CATEGORY = "image"
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, directory, include_subfolders, index, filetype="png",
|
|
reverse=False):
|
|
# index changes every queue (JS hook), so this always reruns; keep
|
|
# directory in the key too so editing the path forces a fresh scan.
|
|
return f"{directory}|{include_subfolders}|{filetype}|{reverse}|{index}"
|
|
|
|
def load(self, directory, include_subfolders, index, filetype="png",
|
|
reverse=False):
|
|
files = list_images(directory, include_subfolders, filetype, reverse)
|
|
total = len(files)
|
|
if total == 0:
|
|
label = "image" if filetype == "all" else filetype.upper()
|
|
raise ValueError(
|
|
f"[Sequential Image Loader] No {label} files found in: "
|
|
f"{directory!r}"
|
|
)
|
|
|
|
pos = int(index) % total # wrap around at the end
|
|
path = os.path.abspath(files[pos])
|
|
filename = os.path.basename(path)
|
|
|
|
image, mask = _load_image_tensors(path)
|
|
|
|
return {
|
|
"ui": {"seqloader": [{
|
|
"filename": filename,
|
|
"path": path,
|
|
"index": pos,
|
|
"total": total,
|
|
}]},
|
|
"result": (image, mask, filename, pos, total, path),
|
|
}
|
|
|
|
|
|
# --- Directory browsing (server-side folder picker) -----------------------
|
|
|
|
def _list_roots():
|
|
"""Top-level entries: drive letters on Windows, '/' + home elsewhere."""
|
|
roots = []
|
|
if os.name == "nt":
|
|
import string
|
|
for c in string.ascii_uppercase:
|
|
d = f"{c}:\\"
|
|
if os.path.exists(d):
|
|
roots.append({"name": d, "path": d})
|
|
else:
|
|
roots.append({"name": "/", "path": "/"})
|
|
home = os.path.expanduser("~")
|
|
if os.path.isdir(home):
|
|
roots.append({"name": f"~ ({home})", "path": home})
|
|
return roots
|
|
|
|
|
|
def browse_dir(path, filetype="png"):
|
|
"""Resolve a folder for the picker UI.
|
|
|
|
Returns the current path, its parent ("" => the roots list), the
|
|
immediate subdirectories, and how many matching images live here.
|
|
An empty/invalid path returns the roots listing.
|
|
"""
|
|
path = (path or "").strip().strip('"').strip("'")
|
|
if not path or not os.path.isdir(path):
|
|
return {"cwd": "", "parent": None, "is_roots": True,
|
|
"dirs": _list_roots(), "image_count": 0}
|
|
|
|
norm = os.path.normpath(path)
|
|
parent = os.path.dirname(norm)
|
|
if parent == norm: # at a drive root -> "up" goes to the roots list
|
|
parent = ""
|
|
|
|
dirs = []
|
|
try:
|
|
for name in os.listdir(norm):
|
|
full = os.path.join(norm, name)
|
|
try:
|
|
if os.path.isdir(full):
|
|
dirs.append({"name": name, "path": full})
|
|
except OSError:
|
|
continue
|
|
except OSError:
|
|
pass
|
|
dirs.sort(key=lambda d: _natural_key(d["name"]))
|
|
|
|
image_count = len(list_images(norm, False, filetype, False))
|
|
return {"cwd": norm, "parent": parent, "is_roots": False,
|
|
"dirs": dirs, "image_count": image_count}
|
|
|
|
|
|
# --- Server routes (UI progress display + folder picker) ------------------
|
|
|
|
try:
|
|
from server import PromptServer
|
|
from aiohttp import web
|
|
|
|
@PromptServer.instance.routes.post("/seqloader/list")
|
|
async def _seqloader_list(request):
|
|
try:
|
|
data = await request.json()
|
|
except Exception:
|
|
return web.json_response({"error": "invalid JSON"}, status=400)
|
|
directory = (data or {}).get("directory", "")
|
|
include_subfolders = bool((data or {}).get("include_subfolders", False))
|
|
filetype = (data or {}).get("filetype", "png")
|
|
reverse = bool((data or {}).get("reverse", False))
|
|
files = list_images(directory, include_subfolders, filetype, reverse)
|
|
return web.json_response({
|
|
"count": len(files),
|
|
"files": [os.path.basename(p) for p in files],
|
|
"paths": files,
|
|
})
|
|
|
|
@PromptServer.instance.routes.get("/seqloader/thumb")
|
|
async def _seqloader_thumb(request):
|
|
path = (request.query.get("path", "") or "").strip()
|
|
if (not path or not os.path.isfile(path)
|
|
or not path.lower().endswith(_ALL_EXTS)):
|
|
return web.Response(status=404)
|
|
try:
|
|
img = Image.open(path)
|
|
img = ImageOps.exif_transpose(img)
|
|
img = img.convert("RGB")
|
|
img.thumbnail((256, 256))
|
|
buf = io.BytesIO()
|
|
img.save(buf, format="JPEG", quality=85)
|
|
return web.Response(body=buf.getvalue(), content_type="image/jpeg")
|
|
except Exception as e:
|
|
return web.Response(status=500, text=str(e))
|
|
|
|
@PromptServer.instance.routes.post("/seqloader/browse")
|
|
async def _seqloader_browse(request):
|
|
try:
|
|
data = await request.json()
|
|
except Exception:
|
|
return web.json_response({"error": "invalid JSON"}, status=400)
|
|
path = (data or {}).get("path", "")
|
|
filetype = (data or {}).get("filetype", "png")
|
|
try:
|
|
return web.json_response(browse_dir(path, filetype))
|
|
except Exception as e:
|
|
return web.json_response({"error": str(e)}, status=500)
|
|
except Exception as _e: # pragma: no cover - server optional
|
|
print(f"[Sequential Image Loader] routes not registered: {_e}")
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {"SequentialImageLoader": SequentialImageLoader}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"SequentialImageLoader": "Sequential Image Loader (folder)",
|
|
}
|