528 lines
17 KiB
Python
528 lines
17 KiB
Python
import os
|
|
import time
|
|
import glob
|
|
import random
|
|
import uuid
|
|
from typing import List, Optional, Tuple
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image, ImageOps
|
|
|
|
import folder_paths
|
|
|
|
|
|
IMAGE_EXTS = (".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tiff", ".tif")
|
|
|
|
|
|
def _resolve_folder(path: str) -> str:
|
|
path = path.strip()
|
|
if not path:
|
|
path = ".uploading"
|
|
if not os.path.isabs(path):
|
|
path = os.path.abspath(os.path.join(os.getcwd(), path))
|
|
return path
|
|
|
|
|
|
def _scan_images(folder: str, recursive: bool) -> List[str]:
|
|
if not os.path.isdir(folder):
|
|
return []
|
|
pattern = "**/*" if recursive else "*"
|
|
files = glob.glob(os.path.join(folder, pattern), recursive=recursive)
|
|
out = []
|
|
for f in files:
|
|
if os.path.isfile(f) and os.path.splitext(f)[1].lower() in IMAGE_EXTS:
|
|
out.append(f)
|
|
out.sort(key=lambda x: x.lower())
|
|
return out
|
|
|
|
|
|
def _pil_to_comfy_tensor(img: Image.Image) -> torch.Tensor:
|
|
img = ImageOps.exif_transpose(img)
|
|
if img.mode not in ("RGB", "RGBA"):
|
|
img = img.convert("RGB")
|
|
if img.mode == "RGBA":
|
|
img = img.convert("RGB")
|
|
|
|
arr = np.asarray(img).astype(np.float32) / 255.0
|
|
return torch.from_numpy(arr).unsqueeze(0) # [1,H,W,3]
|
|
|
|
|
|
def _title_from_path(path: str) -> str:
|
|
# Base filename without extension, e.g. "cat.png" -> "cat"
|
|
return os.path.splitext(os.path.basename(path))[0]
|
|
|
|
|
|
class BatchLoaderFromFolder:
|
|
|
|
def __init__(self):
|
|
self._folder: Optional[str] = None
|
|
self._recursive: Optional[bool] = None
|
|
self._files: List[str] = []
|
|
self._seq_index = 0
|
|
self._freeze_left = 0
|
|
self._frozen_selection: Optional[List[str]] = None
|
|
self._random_step = 0
|
|
self._pool: List[str] = []
|
|
self._pool_cycle = 0
|
|
self._pool_seed_last: Optional[int] = None
|
|
self._pool_signature: Optional[Tuple[str, ...]] = None
|
|
self._size_cache: dict[str, Tuple[int, int]] = {}
|
|
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"folder": ("STRING", {"default": ".uploading"}),
|
|
"mode": (["random", "sequential", "batch_freeze"], {"default": "random"}),
|
|
"batch_size": ("INT", {"default": 1, "min": 1, "max": 64, "step": 1}),
|
|
"freeze_count": ("INT", {"default": 1, "min": 1, "max": 999, "step": 1}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0x7FFFFFFF}),
|
|
"no_repeat_random": ("BOOLEAN", {"default": True}),
|
|
"recursive": ("BOOLEAN", {"default": False}),
|
|
"reload": ("BOOLEAN", {"default": False}),
|
|
"reset": ("BOOLEAN", {"default": False}),
|
|
"skip_exact_mode": (
|
|
["disabled", "width or height", "width and height", "width only", "height only"],
|
|
{"default": "disabled"},
|
|
),
|
|
"skip_exact_width": ("INT", {"default": 0, "min": 0, "max": 65535, "step": 1}),
|
|
"skip_exact_height":("INT", {"default": 0, "min": 0, "max": 65535, "step": 1}),
|
|
"move_processed_to_subfolder": (
|
|
"BOOLEAN",
|
|
{"default": False, "label_on": "move on", "label_off": "move off"},
|
|
),
|
|
"processed_subfolder_name": ("STRING", {"default": "_processed"}),
|
|
"resize_to_first": ("BOOLEAN", {"default": True}),
|
|
"show_preview": ("BOOLEAN", {"default": True}),
|
|
"preview_max": ("INT", {"default": 4, "min": 1, "max": 12, "step": 1}),
|
|
}
|
|
}
|
|
|
|
|
|
RETURN_TYPES = ("IMAGE", "STRING", "STRING", "INT", "INT")
|
|
RETURN_NAMES = ("images", "filename", "format", "width", "height")
|
|
FUNCTION = "load"
|
|
CATEGORY = "AntiMatter/Image"
|
|
OUTPUT_NODE = False
|
|
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs):
|
|
return float("inf") # force re-execution on every queue
|
|
|
|
|
|
def _reset_runtime_state(self):
|
|
self._seq_index = 0
|
|
self._freeze_left = 0
|
|
self._frozen_selection = None
|
|
self._random_step = 0
|
|
self._pool = []
|
|
self._pool_cycle = 0
|
|
self._pool_seed_last = None
|
|
self._pool_signature = None
|
|
|
|
|
|
def _ensure_files(self, folder: str, recursive: bool, reload: bool, reset: bool):
|
|
folder_abs = _resolve_folder(folder)
|
|
folder_changed = (self._folder != folder_abs) or (self._recursive != recursive)
|
|
|
|
if reset or folder_changed:
|
|
self._reset_runtime_state()
|
|
|
|
if reload or folder_changed or not self._files:
|
|
self._files = _scan_images(folder_abs, recursive)
|
|
self._pool = []
|
|
self._pool_cycle = 0
|
|
self._pool_signature = None
|
|
|
|
self._folder = folder_abs
|
|
self._recursive = recursive
|
|
|
|
|
|
def _pool_refill_if_needed(self, seed: int, files: List[str]):
|
|
if not self._pool:
|
|
self._pool = list(files)
|
|
rng = random.Random(seed + self._pool_cycle)
|
|
rng.shuffle(self._pool)
|
|
self._pool_cycle += 1
|
|
|
|
|
|
def _choose_random_no_repeat(self, files: List[str], seed: int, n: int) -> List[str]:
|
|
if not files:
|
|
return []
|
|
out = []
|
|
while len(out) < n:
|
|
self._pool_refill_if_needed(seed, files)
|
|
take = min(n - len(out), len(self._pool))
|
|
out.extend(self._pool[:take])
|
|
self._pool = self._pool[take:]
|
|
return out
|
|
|
|
|
|
def _choose_random_with_repeats(self, files: List[str], seed: int, n: int) -> List[str]:
|
|
if not files:
|
|
return []
|
|
rng = random.Random(seed + self._random_step)
|
|
self._random_step += 1
|
|
if n == 1:
|
|
return [rng.choice(files)]
|
|
if len(files) >= n:
|
|
return rng.sample(files, n)
|
|
return [rng.choice(files) for _ in range(n)]
|
|
|
|
|
|
def _choose_random(self, files: List[str], seed: int, n: int, no_repeat: bool) -> List[str]:
|
|
signature = tuple(files)
|
|
if no_repeat and (
|
|
self._pool_seed_last is None
|
|
or self._pool_seed_last != seed
|
|
or self._pool_signature != signature
|
|
):
|
|
self._pool = []
|
|
self._pool_cycle = 0
|
|
self._pool_seed_last = seed
|
|
self._pool_signature = signature
|
|
|
|
if no_repeat:
|
|
return self._choose_random_no_repeat(files, seed, n)
|
|
return self._choose_random_with_repeats(files, seed, n)
|
|
|
|
|
|
def _choose_sequential(self, files: List[str], n: int) -> List[str]:
|
|
if not files:
|
|
return []
|
|
available = set(files)
|
|
out = []
|
|
total = len(self._files) if self._files else len(files)
|
|
checked = 0
|
|
while len(out) < n and checked < total:
|
|
idx = self._seq_index % total
|
|
candidate = self._files[idx] if self._files else files[idx % len(files)]
|
|
self._seq_index += 1
|
|
checked += 1
|
|
if candidate in available:
|
|
out.append(candidate)
|
|
return out
|
|
|
|
|
|
def _processed_subfolder_path(self, folder_abs: str, processed_subfolder_name: str) -> Optional[str]:
|
|
subfolder = processed_subfolder_name.strip()
|
|
if not subfolder:
|
|
return None
|
|
|
|
if os.path.isabs(subfolder):
|
|
return subfolder
|
|
|
|
return os.path.abspath(os.path.join(folder_abs, subfolder))
|
|
|
|
|
|
def _is_in_processed_subfolder(
|
|
self,
|
|
path: str,
|
|
folder_abs: str,
|
|
processed_subfolder_name: str,
|
|
) -> bool:
|
|
processed_dir = self._processed_subfolder_path(folder_abs, processed_subfolder_name)
|
|
if not processed_dir:
|
|
return False
|
|
|
|
try:
|
|
return os.path.commonpath((processed_dir, os.path.abspath(path))) == processed_dir
|
|
except ValueError:
|
|
return False
|
|
|
|
|
|
def _move_to_processed_subfolder(
|
|
self,
|
|
selected_paths: List[str],
|
|
folder_abs: str,
|
|
processed_subfolder_name: str,
|
|
):
|
|
processed_dir = self._processed_subfolder_path(folder_abs, processed_subfolder_name)
|
|
if not processed_dir:
|
|
return
|
|
|
|
os.makedirs(processed_dir, exist_ok=True)
|
|
|
|
updated_files: List[str] = []
|
|
selected_set = set(selected_paths)
|
|
for path in self._files:
|
|
if path not in selected_set:
|
|
updated_files.append(path)
|
|
|
|
for src in selected_paths:
|
|
if not os.path.isfile(src):
|
|
continue
|
|
|
|
base_name = os.path.basename(src)
|
|
dst = os.path.join(processed_dir, base_name)
|
|
|
|
if os.path.abspath(src) == os.path.abspath(dst):
|
|
continue
|
|
|
|
if os.path.exists(dst):
|
|
name, ext = os.path.splitext(base_name)
|
|
suffix = 1
|
|
while True:
|
|
candidate = os.path.join(processed_dir, f"{name}_{suffix}{ext}")
|
|
if not os.path.exists(candidate):
|
|
dst = candidate
|
|
break
|
|
suffix += 1
|
|
|
|
os.replace(src, dst)
|
|
|
|
self._files = updated_files
|
|
self._pool = []
|
|
self._pool_cycle = 0
|
|
self._pool_signature = None
|
|
|
|
|
|
def _get_image_size(self, path: str) -> Optional[Tuple[int, int]]:
|
|
cached = self._size_cache.get(path)
|
|
if cached is not None:
|
|
return cached
|
|
|
|
try:
|
|
with Image.open(path) as img:
|
|
img = ImageOps.exif_transpose(img)
|
|
size = img.size
|
|
except Exception:
|
|
return None
|
|
|
|
self._size_cache[path] = size
|
|
return size
|
|
|
|
|
|
def _should_skip_exact(
|
|
self,
|
|
path: str,
|
|
skip_exact_mode: str,
|
|
skip_exact_width: int,
|
|
skip_exact_height: int,
|
|
) -> bool:
|
|
if skip_exact_mode == "disabled":
|
|
return False
|
|
if skip_exact_width <= 0 and skip_exact_height <= 0:
|
|
return False
|
|
|
|
size = self._get_image_size(path)
|
|
if size is None:
|
|
return False
|
|
|
|
width, height = size
|
|
width_match = skip_exact_width > 0 and width == skip_exact_width
|
|
height_match = skip_exact_height > 0 and height == skip_exact_height
|
|
|
|
if skip_exact_mode == "width and height":
|
|
width_ok = width_match if skip_exact_width > 0 else True
|
|
height_ok = height_match if skip_exact_height > 0 else True
|
|
return width_ok and height_ok
|
|
if skip_exact_mode == "width only":
|
|
return width_match
|
|
if skip_exact_mode == "height only":
|
|
return height_match
|
|
return width_match or height_match
|
|
|
|
|
|
def _get_candidate_files(
|
|
self,
|
|
folder_abs: str,
|
|
skip_exact_mode: str,
|
|
skip_exact_width: int,
|
|
skip_exact_height: int,
|
|
processed_subfolder_name: str,
|
|
) -> List[str]:
|
|
return [
|
|
path
|
|
for path in self._files
|
|
if not self._is_in_processed_subfolder(path, folder_abs, processed_subfolder_name)
|
|
if not self._should_skip_exact(path, skip_exact_mode, skip_exact_width, skip_exact_height)
|
|
]
|
|
|
|
|
|
def _select_files(
|
|
self,
|
|
files: List[str],
|
|
mode: str,
|
|
batch_size: int,
|
|
freeze_count: int,
|
|
seed: int,
|
|
no_repeat_random: bool,
|
|
) -> List[str]:
|
|
if self._freeze_left > 0 and self._frozen_selection is not None:
|
|
available = set(files)
|
|
if all(path in available for path in self._frozen_selection):
|
|
self._freeze_left -= 1
|
|
return self._frozen_selection
|
|
self._freeze_left = 0
|
|
self._frozen_selection = None
|
|
|
|
if mode == "random":
|
|
selection = self._choose_random(files, seed, batch_size, no_repeat_random)
|
|
elif mode == "sequential":
|
|
selection = self._choose_sequential(files, batch_size)
|
|
elif mode == "batch_freeze":
|
|
one = self._choose_random(files, seed, 1, no_repeat_random) if files else []
|
|
selection = one * batch_size if one else []
|
|
else:
|
|
selection = self._choose_random(files, seed, batch_size, no_repeat_random)
|
|
|
|
self._frozen_selection = selection
|
|
self._freeze_left = max(0, freeze_count - 1)
|
|
return selection
|
|
|
|
|
|
def _save_preview(self, images: List[Image.Image], max_count: int) -> dict:
|
|
if not images:
|
|
return {}
|
|
|
|
temp_dir = folder_paths.get_temp_directory()
|
|
os.makedirs(temp_dir, exist_ok=True)
|
|
|
|
ui_images = []
|
|
for img in images[:max(1, max_count)]:
|
|
name = f"preview_{uuid.uuid4().hex[:12]}.png"
|
|
path = os.path.join(temp_dir, name)
|
|
img.save(path, "PNG")
|
|
ui_images.append({
|
|
"filename": name,
|
|
"subfolder": "",
|
|
"type": "temp"
|
|
})
|
|
|
|
return {"images": ui_images}
|
|
|
|
|
|
def load(
|
|
self,
|
|
folder: str,
|
|
mode: str,
|
|
batch_size: int,
|
|
freeze_count: int,
|
|
seed: int,
|
|
no_repeat_random: bool,
|
|
recursive: bool,
|
|
reload: bool,
|
|
reset: bool,
|
|
skip_exact_mode: str,
|
|
skip_exact_width: int,
|
|
skip_exact_height: int,
|
|
move_processed_to_subfolder: bool,
|
|
processed_subfolder_name: str,
|
|
resize_to_first: bool,
|
|
show_preview: bool,
|
|
preview_max: int,
|
|
) -> Tuple[torch.Tensor, str, str, int, int] | dict:
|
|
|
|
self._ensure_files(folder, recursive, reload, reset)
|
|
|
|
if not self._files:
|
|
raise ValueError(f"No images found in folder: {self._folder}")
|
|
|
|
candidate_files = self._get_candidate_files(
|
|
self._folder,
|
|
skip_exact_mode,
|
|
skip_exact_width,
|
|
skip_exact_height,
|
|
processed_subfolder_name,
|
|
)
|
|
|
|
if not candidate_files:
|
|
raise ValueError(
|
|
"No images left after skip filters. "
|
|
"Change filters or add more images."
|
|
)
|
|
|
|
if move_processed_to_subfolder and mode != "batch_freeze" and len(candidate_files) < batch_size:
|
|
raise ValueError(
|
|
f"Only {len(candidate_files)} eligible images left, but batch_size={batch_size}. "
|
|
"Lower batch_size or add more images."
|
|
)
|
|
|
|
selected_paths = self._select_files(
|
|
candidate_files,
|
|
mode,
|
|
batch_size,
|
|
freeze_count,
|
|
seed,
|
|
no_repeat_random,
|
|
)
|
|
|
|
if not selected_paths:
|
|
raise ValueError("Failed to select any images")
|
|
|
|
titles = [_title_from_path(p) for p in selected_paths]
|
|
if not titles:
|
|
title_out = ""
|
|
elif len(set(titles)) == 1:
|
|
title_out = titles[0]
|
|
else:
|
|
title_out = ", ".join(titles)
|
|
|
|
exts = [os.path.splitext(p)[1].lower().lstrip(".") for p in selected_paths]
|
|
exts = [e for e in exts if e]
|
|
if not exts:
|
|
format_out = ""
|
|
elif len(set(exts)) == 1:
|
|
format_out = exts[0]
|
|
else:
|
|
format_out = ", ".join(exts)
|
|
|
|
tensors = []
|
|
preview_images = []
|
|
base_size = None
|
|
|
|
for i, path in enumerate(selected_paths):
|
|
img = Image.open(path)
|
|
img = ImageOps.exif_transpose(img)
|
|
|
|
if img.mode == "RGBA":
|
|
img = img.convert("RGB")
|
|
elif img.mode != "RGB":
|
|
img = img.convert("RGB")
|
|
|
|
if i == 0:
|
|
base_size = img.size
|
|
elif img.size != base_size:
|
|
if resize_to_first:
|
|
img = img.resize(base_size, Image.LANCZOS)
|
|
else:
|
|
raise ValueError(
|
|
f"Image size mismatch in batch: first={base_size}, current={img.size}. "
|
|
"Enable resize_to_first or use batch_size=1."
|
|
)
|
|
|
|
if show_preview:
|
|
preview_images.append(img.copy())
|
|
|
|
tensors.append(_pil_to_comfy_tensor(img))
|
|
|
|
batch = torch.cat(tensors, dim=0)
|
|
width_out = int(batch.shape[2])
|
|
height_out = int(batch.shape[1])
|
|
|
|
if move_processed_to_subfolder:
|
|
self._move_to_processed_subfolder(
|
|
selected_paths,
|
|
self._folder,
|
|
processed_subfolder_name,
|
|
)
|
|
|
|
if show_preview and preview_images:
|
|
ui = self._save_preview(preview_images, preview_max)
|
|
return {"ui": ui, "result": (batch, title_out, format_out, width_out, height_out)}
|
|
|
|
return (batch, title_out, format_out, width_out, height_out)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"Batch_Loader_From_Folder": BatchLoaderFromFolder,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"Batch_Loader_From_Folder": "Batch Loader from folder",
|
|
}
|