Compare commits
6
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
359a3b56dd | ||
|
|
d23c54eda3 | ||
|
|
e8a0f8c005 | ||
|
|
d128024dff | ||
|
|
e21e2ee7ef | ||
|
|
71e42b303b |
@@ -44,14 +44,16 @@ example, ONNX Runtime sessions in the vendored code are created through
|
||||
## Models downloaded at runtime
|
||||
|
||||
These files are **not** stored in this repository. The nodes download them from
|
||||
Hugging Face on first use.
|
||||
Hugging Face on first use, pinned to the revisions below (only the files the nodes need).
|
||||
They are looked up first in the folders registered as `vton_mask` and `nh_nodes`
|
||||
(configurable in `extra_model_paths.yaml`).
|
||||
|
||||
| Node | Hugging Face repository | Contents | Terms |
|
||||
|------|-------------------------|----------|-------|
|
||||
| VTON Ultimate Processor | [`kg-09/kg-vton-mask`](https://huggingface.co/kg-09/kg-vton-mask) | `dwpose/yolox_l.onnx`, `dwpose/dw-ll_ucoco_384.onnx` | DWPose / YOLOX weights, Apache-2.0 upstream |
|
||||
| VTON Ultimate Processor | [`kg-09/kg-vton-mask`](https://huggingface.co/kg-09/kg-vton-mask) @ `4648778` | `dwpose/yolox_l.onnx`, `dwpose/dw-ll_ucoco_384.onnx` | DWPose / YOLOX weights, Apache-2.0 upstream |
|
||||
| VTON Ultimate Processor | same | `humanparsing/parsing_atr.onnx`, `humanparsing/parsing_lip.onnx` | From the OOTDiffusion/FitDiT checkpoints (CC BY-NC-SA 4.0); trained on the ATR and LIP datasets, which are distributed for research use |
|
||||
| Portrait Segment | [`jonathandinu/face-parsing`](https://huggingface.co/jonathandinu/face-parsing) | SegFormer fine-tuned on CelebAMask-HQ | Model card: "non-commercial research and educational purposes" |
|
||||
| Garment Segment | [`1038lab/segformer_fashion`](https://huggingface.co/1038lab/segformer_fashion) | SegFormer trained for Fashionpedia classes | No license declared by the uploader (Fashionpedia annotations are CC BY 4.0) |
|
||||
| Portrait Segment | [`jonathandinu/face-parsing`](https://huggingface.co/jonathandinu/face-parsing) @ `758b82e` | SegFormer fine-tuned on CelebAMask-HQ | Model card: "non-commercial research and educational purposes" |
|
||||
| Garment Segment | [`1038lab/segformer_fashion`](https://huggingface.co/1038lab/segformer_fashion) @ `e5bae2d` | SegFormer trained for Fashionpedia classes | No license declared by the uploader (Fashionpedia annotations are CC BY 4.0) |
|
||||
|
||||
## Python dependencies
|
||||
|
||||
|
||||
+3
-45
@@ -45,58 +45,16 @@ for _category in pkgutil.iter_modules(_nodes_pkg.__path__):
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(getattr(_module, "NODE_DISPLAY_NAME_MAPPINGS", {}))
|
||||
_module_count += 1
|
||||
|
||||
# Route-only module: registers /nh-nodes/large-preview/* on import.
|
||||
# Route-only modules: they register the /nh-nodes/* HTTP routes on import.
|
||||
try:
|
||||
from .core import large_preview_routes # noqa: F401
|
||||
except Exception as e:
|
||||
print(f"[NH-Nodes] Failed to import core.large_preview_routes: {e}")
|
||||
|
||||
try:
|
||||
from aiohttp import web
|
||||
from server import PromptServer
|
||||
|
||||
@PromptServer.instance.routes.get("/nh-nodes/load-images-folder/count")
|
||||
async def nh_load_images_folder_count(request):
|
||||
folder_path = request.query.get("folder_path", "")
|
||||
recursive = request.query.get("recursive", "").lower() in {"1", "true", "yes", "on"}
|
||||
try:
|
||||
from .nodes.image.image_tools_nodes import _collect_image_files, _resolve_load_dir
|
||||
|
||||
resolved_folder = _resolve_load_dir(folder_path)
|
||||
image_count = len(_collect_image_files(resolved_folder, recursive=recursive))
|
||||
return web.json_response({"count": image_count, "path": resolved_folder})
|
||||
except Exception as exc:
|
||||
return web.json_response({"count": 0, "error": str(exc)})
|
||||
|
||||
@PromptServer.instance.routes.get("/nh-nodes/smart-resolution-picker/presets")
|
||||
async def nh_smart_resolution_picker_presets(request):
|
||||
try:
|
||||
from .core.resolution_data import DEFAULT_MODEL_LABEL, DEFAULT_PRESET, MODEL_LABELS, PRESET_LABELS_BY_MODEL
|
||||
from .nodes.resolution.smart_resolution_picker import TARGET_RESOLUTION_LEVELS, _resolution_candidates
|
||||
|
||||
ratios_by_model_level = {}
|
||||
for model_label in MODEL_LABELS:
|
||||
ratios_by_model_level[model_label] = {}
|
||||
for resolution_level in TARGET_RESOLUTION_LEVELS:
|
||||
candidates = _resolution_candidates(model_label, resolution_level)
|
||||
ratios = []
|
||||
for entry in candidates:
|
||||
if entry["aspect"] not in ratios:
|
||||
ratios.append(entry["aspect"])
|
||||
ratios_by_model_level[model_label][resolution_level] = ratios
|
||||
|
||||
return web.json_response({
|
||||
"models": MODEL_LABELS,
|
||||
"default_model": DEFAULT_MODEL_LABEL,
|
||||
"default_preset": DEFAULT_PRESET,
|
||||
"presets_by_model": PRESET_LABELS_BY_MODEL,
|
||||
"resolution_levels": TARGET_RESOLUTION_LEVELS,
|
||||
"ratios_by_model_level": ratios_by_model_level,
|
||||
})
|
||||
except Exception as exc:
|
||||
return web.json_response({"error": str(exc)}, status=500)
|
||||
from .core import folder_routes # noqa: F401
|
||||
except Exception as e:
|
||||
print(f"[NH-Nodes] Failed to register web routes: {e}")
|
||||
print(f"[NH-Nodes] Failed to import core.folder_routes: {e}")
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]
|
||||
|
||||
|
||||
@@ -0,0 +1,103 @@
|
||||
"""HTTP routes behind the NH-Nodes frontend extensions.
|
||||
|
||||
GET /nh-nodes/load-images-folder/count how many images a folder loader sees (auto increment / wrap)
|
||||
GET /nh-nodes/smart-resolution-picker/presets the preset lists for the resolution pickers' dropdowns
|
||||
|
||||
Importing this module registers both on ComfyUI's PromptServer, like core.large_preview_routes. The handlers are plain
|
||||
module-level functions, so a test can serve them from any aiohttp app through ``register_routes``.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
|
||||
from aiohttp import web
|
||||
|
||||
from .image_folders import collect_image_files, resolve_load_dir
|
||||
|
||||
COUNT_PATH = "/nh-nodes/load-images-folder/count"
|
||||
PRESETS_PATH = "/nh-nodes/smart-resolution-picker/presets"
|
||||
|
||||
_MAX_FOLDER_PATH_CHARS = 4096
|
||||
# Spellings of the ``recursive`` query parameter. Anything else is a bad request, not a silent "false".
|
||||
_FLAGS = {"1": True, "true": True, "yes": True, "on": True, "": False, "0": False, "false": False, "no": False, "off": False}
|
||||
|
||||
|
||||
def _error(status, message):
|
||||
# Fixed short texts only: ComfyUI can be reachable from the network (--listen), so an answer never carries a
|
||||
# resolved path or the text of an exception.
|
||||
return web.json_response({"error": message}, status=status)
|
||||
|
||||
|
||||
def _count_images(folder_path, recursive):
|
||||
"""How many images a loader with these settings sees. Blocking (a recursive walk): run it in a worker thread."""
|
||||
folder = resolve_load_dir(folder_path)
|
||||
if not os.path.isdir(folder):
|
||||
raise FileNotFoundError("folder not found")
|
||||
try:
|
||||
return len(collect_image_files(folder, recursive=recursive))
|
||||
except ValueError:
|
||||
return 0 # the folder exists but holds no image
|
||||
|
||||
|
||||
async def nh_load_images_folder_count(request):
|
||||
folder_path = request.query.get("folder_path", "")
|
||||
recursive = _FLAGS.get(request.query.get("recursive", "").strip().lower())
|
||||
if recursive is None:
|
||||
return _error(400, "recursive must be true or false")
|
||||
if len(folder_path) > _MAX_FOLDER_PATH_CHARS or "\0" in folder_path:
|
||||
return _error(400, "invalid folder_path")
|
||||
|
||||
try:
|
||||
count = await asyncio.to_thread(_count_images, folder_path, recursive)
|
||||
except FileNotFoundError:
|
||||
return _error(404, "folder not found")
|
||||
except PermissionError:
|
||||
return _error(403, "folder is not readable")
|
||||
except Exception as exc:
|
||||
print(f"[NH-Nodes] load-images-folder/count failed: {type(exc).__name__}: {exc}")
|
||||
return _error(500, "could not read the folder")
|
||||
return web.json_response({"count": count})
|
||||
|
||||
|
||||
async def nh_smart_resolution_picker_presets(request):
|
||||
try:
|
||||
# Imported here: core must not depend on the node modules at import time.
|
||||
from ..nodes.resolution.smart_resolution_picker import TARGET_RESOLUTION_LEVELS, _resolution_candidates
|
||||
from .resolution_data import DEFAULT_MODEL_LABEL, DEFAULT_PRESET, MODEL_LABELS, PRESET_LABELS_BY_MODEL
|
||||
|
||||
ratios_by_model_level = {}
|
||||
for model_label in MODEL_LABELS:
|
||||
ratios_by_model_level[model_label] = {}
|
||||
for resolution_level in TARGET_RESOLUTION_LEVELS:
|
||||
candidates = _resolution_candidates(model_label, resolution_level)
|
||||
ratios = []
|
||||
for entry in candidates:
|
||||
if entry["aspect"] not in ratios:
|
||||
ratios.append(entry["aspect"])
|
||||
ratios_by_model_level[model_label][resolution_level] = ratios
|
||||
|
||||
return web.json_response({
|
||||
"models": MODEL_LABELS,
|
||||
"default_model": DEFAULT_MODEL_LABEL,
|
||||
"default_preset": DEFAULT_PRESET,
|
||||
"presets_by_model": PRESET_LABELS_BY_MODEL,
|
||||
"resolution_levels": TARGET_RESOLUTION_LEVELS,
|
||||
"ratios_by_model_level": ratios_by_model_level,
|
||||
})
|
||||
except Exception as exc:
|
||||
print(f"[NH-Nodes] smart-resolution-picker/presets failed: {type(exc).__name__}: {exc}")
|
||||
return _error(500, "could not build the preset lists")
|
||||
|
||||
|
||||
def register_routes(routes):
|
||||
"""Attach the handlers to an aiohttp RouteTableDef (ComfyUI's ``PromptServer.instance.routes``)."""
|
||||
routes.get(COUNT_PATH)(nh_load_images_folder_count)
|
||||
routes.get(PRESETS_PATH)(nh_smart_resolution_picker_presets)
|
||||
|
||||
|
||||
try:
|
||||
from server import PromptServer
|
||||
|
||||
register_routes(PromptServer.instance.routes)
|
||||
except Exception as exc:
|
||||
print(f"[NH-Nodes] Failed to register folder routes: {exc}")
|
||||
@@ -0,0 +1,78 @@
|
||||
"""Image folder helpers shared by the folder loader nodes and the folder-count HTTP route.
|
||||
|
||||
``resolve_load_dir`` turns the folder text of a loader into a directory, ``collect_image_files`` lists the images in
|
||||
it in natural order. Nothing here needs torch or Pillow.
|
||||
"""
|
||||
|
||||
import os
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
import folder_paths
|
||||
|
||||
IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp", ".avif", ".bmp", ".tif", ".tiff"}
|
||||
|
||||
|
||||
def natural_key(value):
|
||||
"""Sort key that puts 'img2' before 'img10' and ignores case."""
|
||||
return [int(part) if part.isdigit() else part.lower() for part in re.split(r"(\d+)", value)]
|
||||
|
||||
|
||||
def resolve_load_dir(folder_path, image="", upload_probe=""):
|
||||
"""The folder a loader reads.
|
||||
|
||||
``folder_path`` when given (relative to ComfyUI's input folder unless absolute; a file stands for its folder),
|
||||
else the folder of the uploaded ``image`` / ``upload_probe``, else the input folder.
|
||||
"""
|
||||
folder_path = (folder_path or "").strip()
|
||||
if folder_path:
|
||||
resolved = folder_path
|
||||
if not os.path.isabs(resolved):
|
||||
resolved = os.path.join(folder_paths.get_input_directory(), resolved)
|
||||
if os.path.isfile(resolved):
|
||||
resolved = os.path.dirname(resolved)
|
||||
return resolved
|
||||
|
||||
upload_name = image or upload_probe
|
||||
if upload_name:
|
||||
return os.path.dirname(folder_paths.get_annotated_filepath(upload_name))
|
||||
|
||||
return folder_paths.get_input_directory()
|
||||
|
||||
|
||||
def collect_image_files(folder_path, recursive=False):
|
||||
"""The images of ``folder_path`` in natural order, as names (paths relative to the folder below it).
|
||||
|
||||
Sub-folders are read when ``recursive`` is set, and also when the folder holds no image itself. ValueError for a
|
||||
folder that does not exist or has no images.
|
||||
"""
|
||||
if not os.path.isdir(folder_path):
|
||||
raise ValueError(f"Folder does not exist: {folder_path}")
|
||||
|
||||
files = []
|
||||
for entry in os.listdir(folder_path):
|
||||
full_path = os.path.join(folder_path, entry)
|
||||
if os.path.isfile(full_path) and Path(entry).suffix.lower() in IMAGE_EXTENSIONS:
|
||||
files.append(entry)
|
||||
|
||||
if recursive or not files:
|
||||
seen = set(files)
|
||||
models_dir = os.path.normcase(os.path.abspath(folder_paths.models_dir))
|
||||
for current_root, dir_names, file_names in os.walk(folder_path):
|
||||
# Natural order, and never into the ComfyUI models folder (a walk from one of its parents would crawl it).
|
||||
dir_names[:] = sorted(
|
||||
(name for name in dir_names if os.path.normcase(os.path.abspath(os.path.join(current_root, name))) != models_dir),
|
||||
key=natural_key,
|
||||
)
|
||||
for file_name in sorted(file_names, key=natural_key):
|
||||
full_path = os.path.join(current_root, file_name)
|
||||
if os.path.isfile(full_path) and Path(file_name).suffix.lower() in IMAGE_EXTENSIONS:
|
||||
rel_path = os.path.relpath(full_path, folder_path)
|
||||
if rel_path not in seen:
|
||||
seen.add(rel_path)
|
||||
files.append(rel_path)
|
||||
|
||||
files.sort(key=natural_key)
|
||||
if not files:
|
||||
raise ValueError(f"No images found in folder: {folder_path}")
|
||||
return files
|
||||
@@ -0,0 +1,47 @@
|
||||
"""Pillow / tensor helpers shared by the image nodes (save, load, compare, label, resize, layer composite).
|
||||
|
||||
IMAGE tensors become PIL images with core.utils.tensor_to_pil (rgb=True, plus keep_alpha=True for layers);
|
||||
``pil_to_image_tensor`` converts back in ComfyUI's intermediate dtype.
|
||||
"""
|
||||
|
||||
from PIL import Image, ImageColor
|
||||
|
||||
import comfy.model_management
|
||||
|
||||
from .utils import parse_hex_color, pil_to_tensor
|
||||
|
||||
RESAMPLING = getattr(Image, "Resampling", Image)
|
||||
|
||||
|
||||
def parse_color(value, default):
|
||||
"""Parse a colour into an (r, g, b, a) tuple of 0-255 ints; ``default`` when the text is not a colour.
|
||||
|
||||
'#RGB', '#RRGGBB' and '#RRGGBBAA' go through core.utils.parse_hex_color (opaque unless AA is given). Other CSS
|
||||
colours Pillow understands ('red', 'rgb(1, 2, 3)') keep working, as they did before.
|
||||
"""
|
||||
text = str(value or "").strip()
|
||||
color = parse_hex_color(text, ())
|
||||
if color:
|
||||
return color
|
||||
try:
|
||||
return ImageColor.getcolor(text, "RGBA")
|
||||
except ValueError:
|
||||
return tuple(default) + (255,) * (4 - len(default))
|
||||
|
||||
|
||||
|
||||
|
||||
def pil_to_image_tensor(image):
|
||||
"""A PIL image as a [1, H, W, C] float IMAGE in ComfyUI's intermediate dtype."""
|
||||
return pil_to_tensor(image).to(dtype=comfy.model_management.intermediate_dtype())
|
||||
|
||||
|
||||
def resize_pil(image, size, method):
|
||||
"""``image`` resized to ``size`` (width, height) with the 'lanczos', 'bicubic', 'bilinear' or 'nearest' filter."""
|
||||
resample = {
|
||||
"lanczos": RESAMPLING.LANCZOS,
|
||||
"bicubic": RESAMPLING.BICUBIC,
|
||||
"bilinear": RESAMPLING.BILINEAR,
|
||||
"nearest": RESAMPLING.NEAREST,
|
||||
}[method]
|
||||
return image.resize(size, resample)
|
||||
@@ -10,10 +10,11 @@ import uuid
|
||||
from pathlib import Path
|
||||
|
||||
import folder_paths
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image, ImageChops, ImageDraw
|
||||
|
||||
from .utils import tensor_to_pil
|
||||
|
||||
|
||||
CACHE_SUBDIR = "nh_large_preview"
|
||||
CACHE_ID_RE = re.compile(r"^[a-f0-9]{32}$")
|
||||
@@ -139,21 +140,6 @@ def delete_cache(cache_id):
|
||||
}
|
||||
|
||||
|
||||
def _tensor_to_pil(image_tensor):
|
||||
if not isinstance(image_tensor, torch.Tensor):
|
||||
raise TypeError(f"Expected IMAGE tensor, got {type(image_tensor).__name__}")
|
||||
|
||||
array = image_tensor.detach().cpu().numpy()
|
||||
array = np.clip(array * 255.0, 0, 255).astype(np.uint8)
|
||||
if array.ndim == 2:
|
||||
return Image.fromarray(array).convert("RGB")
|
||||
if array.ndim == 3 and array.shape[-1] == 1:
|
||||
return Image.fromarray(array[:, :, 0]).convert("RGB")
|
||||
if array.ndim == 3 and array.shape[-1] == 4:
|
||||
return Image.fromarray(array).convert("RGBA")
|
||||
return Image.fromarray(array).convert("RGB")
|
||||
|
||||
|
||||
def _preview_dimensions(width, height, preview_megapixels):
|
||||
width = max(1, int(width))
|
||||
height = max(1, int(height))
|
||||
@@ -292,7 +278,9 @@ def create_preview_cache(
|
||||
batch_index=0,
|
||||
batch_count=1,
|
||||
):
|
||||
source = _tensor_to_pil(image_tensor)
|
||||
if not isinstance(image_tensor, torch.Tensor):
|
||||
raise TypeError(f"Expected IMAGE tensor, got {type(image_tensor).__name__}")
|
||||
source = tensor_to_pil(image_tensor, rgb=True, keep_alpha=True)
|
||||
source_width, source_height = source.size
|
||||
preview_width, preview_height = _preview_dimensions(source_width, source_height, preview_megapixels)
|
||||
preview = source.copy()
|
||||
|
||||
@@ -1,67 +1,106 @@
|
||||
"""HTTP routes for NH large image preview cache."""
|
||||
"""HTTP routes for NH large image preview cache.
|
||||
|
||||
Importing this module registers the routes on ComfyUI's PromptServer. The handlers are plain module-level functions,
|
||||
so a test can serve them from any aiohttp app through ``register_routes``.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
|
||||
from aiohttp import web
|
||||
|
||||
from .large_preview_cache import cache_info, cleanup_cache, delete_cache, get_thumb_path, get_tile_path, read_manifest
|
||||
|
||||
|
||||
def _without_paths(value):
|
||||
"""Answers go to the browser: drop absolute filesystem paths (the Cache Manager node still shows them)."""
|
||||
if isinstance(value, dict):
|
||||
return {key: _without_paths(item) for key, item in value.items() if key not in ("path", "root")}
|
||||
if isinstance(value, list):
|
||||
return [_without_paths(item) for item in value]
|
||||
return value
|
||||
|
||||
|
||||
async def _read_json_object(request):
|
||||
"""The JSON object in the request body ({} when the body is empty). ValueError for anything else."""
|
||||
try:
|
||||
text = await request.text()
|
||||
payload = json.loads(text) if text.strip() else {}
|
||||
except ValueError: # not JSON, or not UTF-8 text: json.JSONDecodeError and UnicodeDecodeError both are ValueErrors
|
||||
raise ValueError("request body is not valid JSON") from None
|
||||
if not isinstance(payload, dict):
|
||||
raise ValueError("request body must be a JSON object")
|
||||
return payload
|
||||
|
||||
|
||||
async def nh_large_preview_manifest(request):
|
||||
try:
|
||||
manifest = read_manifest(request.match_info["cache_id"])
|
||||
return web.json_response(manifest)
|
||||
except FileNotFoundError:
|
||||
return web.json_response({"error": "manifest not found"}, status=404)
|
||||
except ValueError as exc:
|
||||
return web.json_response({"error": str(exc)}, status=400)
|
||||
|
||||
|
||||
async def nh_large_preview_thumb(request):
|
||||
try:
|
||||
return web.FileResponse(get_thumb_path(request.match_info["cache_id"]))
|
||||
except FileNotFoundError:
|
||||
return web.json_response({"error": "thumbnail not found"}, status=404)
|
||||
except ValueError as exc:
|
||||
return web.json_response({"error": str(exc)}, status=400)
|
||||
|
||||
|
||||
async def nh_large_preview_tile(request):
|
||||
try:
|
||||
return web.FileResponse(get_tile_path(
|
||||
request.match_info["cache_id"],
|
||||
request.match_info["level"],
|
||||
request.match_info["tile"],
|
||||
))
|
||||
except FileNotFoundError:
|
||||
return web.json_response({"error": "tile not available in thumbnail-only preview"}, status=404)
|
||||
except ValueError as exc:
|
||||
return web.json_response({"error": str(exc)}, status=400)
|
||||
|
||||
|
||||
async def nh_large_preview_cache_info(request):
|
||||
try:
|
||||
# Directory walks run in a worker thread so a large cache never stalls the server's event loop.
|
||||
return web.json_response(_without_paths(await asyncio.to_thread(cache_info)))
|
||||
except ValueError as exc:
|
||||
return web.json_response({"error": str(exc)}, status=400)
|
||||
|
||||
|
||||
async def nh_large_preview_cleanup(request):
|
||||
try:
|
||||
payload = await _read_json_object(request)
|
||||
except ValueError as exc:
|
||||
return web.json_response({"error": str(exc)}, status=400)
|
||||
try:
|
||||
cache_id = payload.get("cache_id", "")
|
||||
if cache_id:
|
||||
result = await asyncio.to_thread(delete_cache, cache_id)
|
||||
else:
|
||||
result = await asyncio.to_thread(cleanup_cache, payload.get("max_age_hours", 24), payload.get("max_total_mb", 0))
|
||||
return web.json_response(_without_paths(result))
|
||||
except (ValueError, TypeError, OverflowError) as exc: # a bad id, or a limit that is not a number
|
||||
return web.json_response({"error": str(exc)}, status=400)
|
||||
|
||||
|
||||
def register_routes(routes):
|
||||
"""Attach the handlers to an aiohttp RouteTableDef (ComfyUI's ``PromptServer.instance.routes``)."""
|
||||
routes.get("/nh-nodes/large-preview/manifest/{cache_id}")(nh_large_preview_manifest)
|
||||
routes.get("/nh-nodes/large-preview/thumb/{cache_id}")(nh_large_preview_thumb)
|
||||
routes.get("/nh-nodes/large-preview/tile/{cache_id}/{level}/{tile}")(nh_large_preview_tile)
|
||||
routes.get("/nh-nodes/large-preview/cache-info")(nh_large_preview_cache_info)
|
||||
routes.post("/nh-nodes/large-preview/cleanup")(nh_large_preview_cleanup)
|
||||
|
||||
|
||||
try:
|
||||
from server import PromptServer
|
||||
|
||||
@PromptServer.instance.routes.get("/nh-nodes/large-preview/manifest/{cache_id}")
|
||||
async def nh_large_preview_manifest(request):
|
||||
try:
|
||||
manifest = read_manifest(request.match_info["cache_id"])
|
||||
return web.json_response(manifest)
|
||||
except FileNotFoundError:
|
||||
return web.json_response({"error": "manifest not found"}, status=404)
|
||||
except ValueError as exc:
|
||||
return web.json_response({"error": str(exc)}, status=400)
|
||||
|
||||
@PromptServer.instance.routes.get("/nh-nodes/large-preview/thumb/{cache_id}")
|
||||
async def nh_large_preview_thumb(request):
|
||||
try:
|
||||
return web.FileResponse(get_thumb_path(request.match_info["cache_id"]))
|
||||
except FileNotFoundError:
|
||||
return web.json_response({"error": "thumbnail not found"}, status=404)
|
||||
except ValueError as exc:
|
||||
return web.json_response({"error": str(exc)}, status=400)
|
||||
|
||||
@PromptServer.instance.routes.get("/nh-nodes/large-preview/tile/{cache_id}/{level}/{tile}")
|
||||
async def nh_large_preview_tile(request):
|
||||
try:
|
||||
return web.FileResponse(get_tile_path(
|
||||
request.match_info["cache_id"],
|
||||
request.match_info["level"],
|
||||
request.match_info["tile"],
|
||||
))
|
||||
except FileNotFoundError:
|
||||
return web.json_response({"error": "tile not available in thumbnail-only preview"}, status=404)
|
||||
except ValueError as exc:
|
||||
return web.json_response({"error": str(exc)}, status=400)
|
||||
|
||||
@PromptServer.instance.routes.get("/nh-nodes/large-preview/cache-info")
|
||||
async def nh_large_preview_cache_info(request):
|
||||
try:
|
||||
return web.json_response(cache_info())
|
||||
except ValueError as exc:
|
||||
return web.json_response({"error": str(exc)}, status=400)
|
||||
|
||||
@PromptServer.instance.routes.post("/nh-nodes/large-preview/cleanup")
|
||||
async def nh_large_preview_cleanup(request):
|
||||
try:
|
||||
payload = await request.json()
|
||||
except Exception:
|
||||
payload = {}
|
||||
try:
|
||||
cache_id = payload.get("cache_id", "")
|
||||
if cache_id:
|
||||
result = delete_cache(cache_id)
|
||||
else:
|
||||
result = cleanup_cache(payload.get("max_age_hours", 24), payload.get("max_total_mb", 0))
|
||||
return web.json_response(result)
|
||||
except ValueError as exc:
|
||||
return web.json_response({"error": str(exc)}, status=400)
|
||||
register_routes(PromptServer.instance.routes)
|
||||
except Exception as exc:
|
||||
print(f"[NH-Nodes] Failed to register large preview routes: {exc}")
|
||||
|
||||
|
||||
@@ -0,0 +1,25 @@
|
||||
"""Mask helpers shared by the NH mask and image nodes.
|
||||
|
||||
NH-Nodes' own code (MIT). It is deliberately not part of core/utils_mask.py: that file is third-party
|
||||
FitDiT code under CC BY-NC-SA 4.0 (see THIRD_PARTY_NOTICES.md), and plain MIT nodes must not depend on it.
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def mask_bbox(mask, threshold=0.5):
|
||||
"""Inclusive bounding box ``(x1, y1, x2, y2)`` of the pixels of a 2-D mask that are above ``threshold``.
|
||||
|
||||
``mask`` is an [H, W] array (a torch mask goes through ``.cpu().numpy()`` first). A pixel counts only when it is
|
||||
strictly greater than ``threshold``, so a value of exactly 0.5 is outside, as in every NH mask node. The result is
|
||||
four Python ints; (x2, y2) is the last pixel inside, so width = x2 - x1 + 1 and height = y2 - y1 + 1.
|
||||
Returns None when no pixel is above the threshold, i.e. for an empty mask. Anything that is not 2-D raises ValueError.
|
||||
"""
|
||||
inside = np.asarray(mask) > threshold
|
||||
if inside.ndim != 2:
|
||||
raise ValueError(f"mask_bbox expects a 2-D [H, W] mask, got shape {inside.shape}")
|
||||
rows = np.flatnonzero(inside.any(axis=1))
|
||||
if rows.size == 0:
|
||||
return None
|
||||
cols = np.flatnonzero(inside.any(axis=0))
|
||||
return int(cols[0]), int(rows[0]), int(cols[-1]), int(rows[-1])
|
||||
@@ -0,0 +1,131 @@
|
||||
"""Model-folder, download and VRAM helpers shared by the SegFormer and VTON model loaders.
|
||||
|
||||
Both model families follow the same rules:
|
||||
|
||||
* Model folders are registered with ComfyUI's ``folder_paths`` under the names
|
||||
``nh_nodes`` (SegFormer) and ``vton_mask`` (VTON). Users can add more
|
||||
locations with those names in ``extra_model_paths.yaml``; every registered
|
||||
location is searched before anything is downloaded. Each name is registered
|
||||
with the extension of its model files (``.safetensors``, ``.onnx``), so
|
||||
ComfyUI's model lists show those files and not the Hugging Face cache
|
||||
metadata that downloads leave in the same folder.
|
||||
* Downloads always go to the default folder under the models dir, pinned to a
|
||||
fixed Hugging Face revision.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
import torch
|
||||
|
||||
try:
|
||||
import folder_paths
|
||||
except Exception: # running outside ComfyUI (tests)
|
||||
folder_paths = None
|
||||
|
||||
try:
|
||||
import comfy.model_management as _mm
|
||||
except Exception:
|
||||
_mm = None
|
||||
|
||||
|
||||
def models_dir():
|
||||
"""ComfyUI's models directory, read on every call (--models-directory, tests)."""
|
||||
if folder_paths is not None:
|
||||
return folder_paths.models_dir
|
||||
return os.path.join(os.path.dirname(__file__), "models")
|
||||
|
||||
|
||||
def unique_dirs(paths):
|
||||
"""Drop repeats (same folder spelled differently) but keep the first spelling and the order."""
|
||||
seen, unique = set(), []
|
||||
for path in paths:
|
||||
key = os.path.normcase(os.path.abspath(path))
|
||||
if key not in seen:
|
||||
seen.add(key)
|
||||
unique.append(path)
|
||||
return unique
|
||||
|
||||
|
||||
def _list_only(folder_name, extensions):
|
||||
"""Add `extensions` to the file types ComfyUI lists for the registered `folder_name`.
|
||||
|
||||
add_model_folder_path registers a name with an empty extension set, which ComfyUI reads as "list
|
||||
every file": /models and the Model Library then show the .cache/huggingface metadata of our
|
||||
downloads and a repo's .gitattributes. The set is replaced, never changed in place (ComfyUI shares
|
||||
one set between many folder names), and the locations list stays the very same object, so every
|
||||
extra_model_paths.yaml entry survives.
|
||||
"""
|
||||
if not extensions:
|
||||
return
|
||||
registry = folder_paths.folder_names_and_paths
|
||||
paths, current = registry[folder_name]
|
||||
merged = {*current, *extensions}
|
||||
if merged != set(current):
|
||||
registry[folder_name] = (paths, merged)
|
||||
|
||||
|
||||
def register_model_folder(folder_name, default_dir, extensions=()):
|
||||
"""Register `folder_name` with ComfyUI and return every folder searched for it, best first.
|
||||
|
||||
ComfyUI applies extra_model_paths.yaml before custom nodes load, so locations the user
|
||||
configured under `folder_name` are already registered; they come first. `default_dir`
|
||||
is added too (a no-op when already present) and is always part of the result, so models
|
||||
that live there are found whether or not the user configured anything. Safe to call repeatedly.
|
||||
|
||||
`extensions` (".onnx") are the file types of the models in these folders: ComfyUI's lists for
|
||||
`folder_name` then show only those, also for the locations the user added.
|
||||
"""
|
||||
registered = []
|
||||
if folder_paths is not None:
|
||||
try:
|
||||
folder_paths.add_model_folder_path(folder_name, default_dir)
|
||||
registered = folder_paths.get_folder_paths(folder_name)
|
||||
_list_only(folder_name, extensions) # last: if it fails, the locations above are still returned
|
||||
except Exception as exc: # an unusual ComfyUI build: fall back to the default location alone
|
||||
print(f"[NH-Nodes] Could not register model folder '{folder_name}' with ComfyUI: {exc}")
|
||||
return unique_dirs([*registered, default_dir])
|
||||
|
||||
|
||||
def download_pinned(repo_id, revision, target, allow_patterns):
|
||||
"""snapshot_download of `allow_patterns` from `repo_id` at `revision` into `target`.
|
||||
|
||||
Finished files are kept and partial ones resume, so calling it again after an interrupted
|
||||
download completes the folder. Failures are re-raised with where to get the files by hand.
|
||||
"""
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
os.makedirs(target, exist_ok=True)
|
||||
print(f"[NH-Nodes] Downloading {repo_id}@{revision[:8]} -> {target}")
|
||||
try:
|
||||
snapshot_download(repo_id=repo_id, revision=revision, local_dir=target, allow_patterns=list(allow_patterns))
|
||||
except Exception as exc:
|
||||
raise RuntimeError(
|
||||
f"[NH-Nodes] Downloading {repo_id} (revision {revision[:8]}) into {target} failed: {exc}\n"
|
||||
f"Run again once the connection works (finished files are kept), or download "
|
||||
f"{', '.join(allow_patterns)} from https://huggingface.co/{repo_id}/tree/{revision} into that folder."
|
||||
) from exc
|
||||
|
||||
|
||||
def make_room_on_gpu(nbytes):
|
||||
"""Ask ComfyUI to unload its own models until about `nbytes` of VRAM are free.
|
||||
|
||||
ComfyUI's own loaders ask for their size plus the VRAM it keeps in reserve (extra_reserved_memory(),
|
||||
set by --reserve-vram or its built-in default), so that reserve is added here as well: asking for
|
||||
`nbytes` alone would leave us that much short. Also hands PyTorch's cached blocks back to the
|
||||
driver: libraries that allocate outside PyTorch (ONNX Runtime) cannot use them. Nothing happens
|
||||
without a CUDA device, when ComfyUI itself runs on the CPU, or outside ComfyUI. A failure is
|
||||
reported, not raised: the caller still tries its allocation.
|
||||
"""
|
||||
if _mm is None or not torch.cuda.is_available():
|
||||
return
|
||||
try:
|
||||
# ComfyUI's loaded models carry an indexed device (cuda:0); an index-less
|
||||
# torch.device("cuda") compares unequal to it and would match none of them.
|
||||
device = _mm.get_torch_device()
|
||||
if device.type != "cuda":
|
||||
return
|
||||
reserved = getattr(_mm, "extra_reserved_memory", None) # absent in older ComfyUI builds
|
||||
_mm.free_memory(nbytes + (reserved() if reserved else 0), device)
|
||||
_mm.soft_empty_cache()
|
||||
except Exception as exc:
|
||||
print(f"[NH-Nodes] Could not ask ComfyUI to free VRAM ({exc}); continuing")
|
||||
@@ -2,6 +2,10 @@
|
||||
|
||||
Data is kept separate from node logic so new model families can be added
|
||||
without changing execution code.
|
||||
|
||||
min_res / max_res are the clamp range for *custom* sizes. Curated presets may go
|
||||
beyond max_res (e.g. 4 MP tiers); their own reliability/warning fields say how
|
||||
safe they are, and the node reports that warning in its info output.
|
||||
"""
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,126 @@
|
||||
"""Arithmetic expression evaluator shared by Math Eval (NH) and layer position expressions.
|
||||
|
||||
Accepted: numbers, named variables, + - * / // % **, unary + and -, and an optional
|
||||
whitelist of functions. Everything is computed in floating point, so no expression can
|
||||
build a huge integer and freeze the worker (``9**9**9`` simply overflows and is reported).
|
||||
"""
|
||||
|
||||
import ast
|
||||
import math
|
||||
import numbers
|
||||
import operator
|
||||
|
||||
# Python's ** returns a complex number for a negative base with a fractional exponent;
|
||||
# math.pow raises instead, which is what a node input needs.
|
||||
_BINARY = {
|
||||
ast.Add: operator.add,
|
||||
ast.Sub: operator.sub,
|
||||
ast.Mult: operator.mul,
|
||||
ast.Div: operator.truediv,
|
||||
ast.FloorDiv: operator.floordiv,
|
||||
ast.Mod: operator.mod,
|
||||
ast.Pow: math.pow,
|
||||
}
|
||||
_UNARY = {
|
||||
ast.UAdd: operator.pos,
|
||||
ast.USub: operator.neg,
|
||||
}
|
||||
_DIVISIONS = (ast.Div, ast.FloorDiv, ast.Mod)
|
||||
|
||||
MATH_FUNCTIONS = {
|
||||
"min": min,
|
||||
"max": max,
|
||||
"abs": abs,
|
||||
"round": lambda value, digits=0: round(value, int(digits)),
|
||||
"sqrt": math.sqrt,
|
||||
"floor": math.floor,
|
||||
"ceil": math.ceil,
|
||||
"clamp": lambda value, low, high: max(low, min(high, value)),
|
||||
}
|
||||
|
||||
|
||||
class ExpressionError(ValueError):
|
||||
"""Raised for any expression that cannot be evaluated to a finite number."""
|
||||
|
||||
|
||||
def evaluate(expression, variables=None, functions=None, *, zero_division="zero",
|
||||
case_insensitive=False, on_warning=None):
|
||||
"""Evaluate ``expression`` and return a finite float.
|
||||
|
||||
zero_division: "zero" makes x/0, x//0 and x%0 evaluate to 0 (reported through
|
||||
``on_warning``); "error" raises ExpressionError.
|
||||
case_insensitive: match variable names with str.casefold (keys must already be folded).
|
||||
"""
|
||||
variables = variables or {}
|
||||
functions = functions or {}
|
||||
text = str(expression).strip()
|
||||
if not text:
|
||||
raise ExpressionError("Expression is empty")
|
||||
|
||||
def warn(message):
|
||||
if on_warning is not None:
|
||||
on_warning(message)
|
||||
|
||||
def number(value, what):
|
||||
# numpy scalars and 0-d torch tensors arrive from upstream nodes on FLOAT/INT sockets.
|
||||
if getattr(value, "ndim", None) == 0 and callable(getattr(value, "item", None)):
|
||||
value = value.item()
|
||||
if not isinstance(value, numbers.Real): # bool is a Real (True + 1 == 2, as before)
|
||||
raise ExpressionError(f"{what} must be a number, got {type(value).__name__}")
|
||||
try:
|
||||
return float(value)
|
||||
except OverflowError as exc:
|
||||
raise ExpressionError(f"{what} is too large") from exc
|
||||
|
||||
def lookup(name):
|
||||
key = name.casefold() if case_insensitive else name
|
||||
if key in variables:
|
||||
return number(variables[key], f"Variable '{name}'")
|
||||
raise ExpressionError(f"Unknown variable '{name}'")
|
||||
|
||||
def visit(node):
|
||||
if isinstance(node, ast.Constant):
|
||||
return number(node.value, f"Constant {node.value!r}")
|
||||
if isinstance(node, ast.Name):
|
||||
return lookup(node.id)
|
||||
if isinstance(node, ast.BinOp):
|
||||
op = _BINARY.get(type(node.op))
|
||||
if op is None:
|
||||
raise ExpressionError(f"Operator '{type(node.op).__name__}' is not allowed")
|
||||
left, right = visit(node.left), visit(node.right)
|
||||
if isinstance(node.op, _DIVISIONS) and right == 0:
|
||||
if zero_division == "zero":
|
||||
warn("division by zero, using 0")
|
||||
return 0.0
|
||||
raise ExpressionError("Division by zero")
|
||||
return float(op(left, right))
|
||||
if isinstance(node, ast.UnaryOp):
|
||||
op = _UNARY.get(type(node.op))
|
||||
if op is None:
|
||||
raise ExpressionError(f"Operator '{type(node.op).__name__}' is not allowed")
|
||||
return op(visit(node.operand))
|
||||
if isinstance(node, ast.Call):
|
||||
name = node.func.id if isinstance(node.func, ast.Name) else None
|
||||
if name not in functions or node.keywords:
|
||||
raise ExpressionError(f"Function '{name}' is not allowed")
|
||||
return number(functions[name](*[visit(arg) for arg in node.args]), f"{name}()")
|
||||
raise ExpressionError(f"'{type(node).__name__}' is not allowed in an expression")
|
||||
|
||||
try:
|
||||
result = visit(ast.parse(text, mode="eval").body)
|
||||
except ExpressionError:
|
||||
raise
|
||||
except SyntaxError as exc:
|
||||
raise ExpressionError(f"Invalid expression: {text}") from exc
|
||||
except OverflowError as exc:
|
||||
raise ExpressionError(f"Result is too large: {text}") from exc
|
||||
except (ValueError, TypeError, ZeroDivisionError) as exc:
|
||||
# math domain errors, wrong argument counts, 0 ** -1 ...
|
||||
raise ExpressionError(f"Cannot evaluate {text}: {exc}") from exc
|
||||
except (RecursionError, MemoryError) as exc:
|
||||
# The parser raises MemoryError ("parser stack overflowed") for absurdly nested input.
|
||||
raise ExpressionError("Expression is nested too deeply") from exc
|
||||
|
||||
if not math.isfinite(result):
|
||||
raise ExpressionError(f"Result is not a finite number: {text}")
|
||||
return result
|
||||
+62
-31
@@ -3,6 +3,14 @@
|
||||
Handles model discovery (reuse dirs already on disk before downloading),
|
||||
per-model caching with CPU offload between runs, and full-resolution label
|
||||
map prediction.
|
||||
|
||||
The model-folder, download and VRAM helpers it shares with the VTON node live
|
||||
in model_files.py, so both model families follow the same rules. SegFormer
|
||||
models are registered with ComfyUI's ``folder_paths`` under the name
|
||||
``nh_nodes``; users can add more locations with that name in
|
||||
``extra_model_paths.yaml`` and every registered location is searched before
|
||||
anything is downloaded. Downloads always go to the default folder under the
|
||||
models dir, pinned to a fixed Hugging Face revision.
|
||||
"""
|
||||
|
||||
import os
|
||||
@@ -12,35 +20,34 @@ import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
try:
|
||||
import folder_paths
|
||||
|
||||
_MODELS_DIR = folder_paths.models_dir
|
||||
except Exception: # running outside ComfyUI (tests)
|
||||
_MODELS_DIR = os.path.join(os.path.dirname(__file__), "models")
|
||||
|
||||
try:
|
||||
import comfy.model_management as _mm
|
||||
except Exception:
|
||||
_mm = None
|
||||
from .model_files import download_pinned, make_room_on_gpu, models_dir, register_model_folder, unique_dirs
|
||||
|
||||
_REQUIRED_FILES = ("config.json", "model.safetensors", "preprocessor_config.json")
|
||||
# Logits are interpolated at most to this size before argmax; the label map is
|
||||
# then nearest-resized to the input. Keeps VRAM bounded on very large inputs.
|
||||
_MAX_LOGIT_SIDE = 2048
|
||||
|
||||
# name -> HF repo, local dirs to probe before downloading (relative to models
|
||||
# dir; other packs may already hold the same weights) and a label signature
|
||||
# that guards against a different checkpoint sitting in a probed dir.
|
||||
# extra_model_paths.yaml key for the SegFormer models. Each registered location holds
|
||||
# one sub-folder per model (see MODEL_SPECS), e.g. <location>/segformer_fashion/config.json.
|
||||
FOLDER_NAME = "nh_nodes"
|
||||
FOLDER_EXTENSIONS = (".safetensors",) # what ComfyUI lists for FOLDER_NAME: the weights, not config.json or cache metadata
|
||||
_DEFAULT_FOLDER = "NH-Nodes" # under the models dir: the download target, searched first by default
|
||||
|
||||
# name -> HF repo + pinned revision, legacy local dirs to probe after the registered
|
||||
# folders (relative to the models dir; older versions of this pack and other packs may
|
||||
# already hold the same weights) and a label signature that guards against a different
|
||||
# checkpoint sitting in a probed dir.
|
||||
MODEL_SPECS = {
|
||||
"segformer_fashion": {
|
||||
"repo_id": "1038lab/segformer_fashion",
|
||||
"probe_dirs": ("NH-Nodes/segformer_fashion", "RMBG/segformer_fashion"),
|
||||
"revision": "e5bae2d8e3333be9f12eb6aae454c0542c1371d3",
|
||||
"legacy_dirs": ("RMBG/segformer_fashion",),
|
||||
"signature": {1: "shirt, blouse", 24: "shoe", 46: "tassel"},
|
||||
},
|
||||
"face_parsing": {
|
||||
"repo_id": "jonathandinu/face-parsing",
|
||||
"probe_dirs": ("NH-Nodes/face_parsing", "face_parsing", "RMBG/segformer_face"),
|
||||
"revision": "758b82e15a0178c9db39c1ff666a8b56e3a550c8",
|
||||
"legacy_dirs": ("face_parsing", "RMBG/segformer_face"),
|
||||
"signature": {1: "skin", 13: "hair", 17: "neck"},
|
||||
},
|
||||
}
|
||||
@@ -49,35 +56,54 @@ _LOADED = {}
|
||||
_warned_no_cuda = False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Model folders and downloads (the shared helpers are in model_files.py)
|
||||
# ---------------------------------------------------------------------------
|
||||
def default_folder():
|
||||
"""Where SegFormer models are downloaded and searched first: <models dir>/NH-Nodes."""
|
||||
return os.path.join(models_dir(), _DEFAULT_FOLDER)
|
||||
|
||||
|
||||
def _has_required_files(path):
|
||||
return all(os.path.isfile(os.path.join(path, name)) for name in _REQUIRED_FILES)
|
||||
|
||||
|
||||
def model_search_dirs(model_name):
|
||||
"""Folders that may hold `model_name`, best match first.
|
||||
|
||||
The registered ``nh_nodes`` locations come first (models/NH-Nodes by default, plus whatever
|
||||
extra_model_paths.yaml adds), then the legacy dirs under the models dir.
|
||||
"""
|
||||
spec = MODEL_SPECS[model_name]
|
||||
dirs = [os.path.join(base, model_name)
|
||||
for base in register_model_folder(FOLDER_NAME, default_folder(), FOLDER_EXTENSIONS)]
|
||||
dirs += [os.path.normpath(os.path.join(models_dir(), *rel.split("/"))) for rel in spec["legacy_dirs"]]
|
||||
return unique_dirs(dirs)
|
||||
|
||||
|
||||
def resolve_model_dir(model_name, download=True):
|
||||
"""Return a local dir holding the model, downloading into models/NH-Nodes if absent."""
|
||||
spec = MODEL_SPECS[model_name]
|
||||
for rel in spec["probe_dirs"]:
|
||||
candidate = os.path.join(_MODELS_DIR, rel)
|
||||
for candidate in model_search_dirs(model_name):
|
||||
if _has_required_files(candidate):
|
||||
return candidate
|
||||
|
||||
target = os.path.join(_MODELS_DIR, spec["probe_dirs"][0])
|
||||
if not download:
|
||||
return None
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
os.makedirs(target, exist_ok=True)
|
||||
print(f"[NH-Nodes] Downloading {spec['repo_id']} -> {target}")
|
||||
snapshot_download(
|
||||
repo_id=spec["repo_id"],
|
||||
local_dir=target,
|
||||
allow_patterns=list(_REQUIRED_FILES),
|
||||
)
|
||||
target = os.path.join(default_folder(), model_name)
|
||||
download_pinned(spec["repo_id"], spec["revision"], target, _REQUIRED_FILES)
|
||||
if not _has_required_files(target):
|
||||
raise RuntimeError(f"[NH-Nodes] Model files missing after download: {target}")
|
||||
return target
|
||||
|
||||
|
||||
# Make the folder name known to ComfyUI as soon as the pack loads.
|
||||
register_model_folder(FOLDER_NAME, default_folder(), FOLDER_EXTENSIONS)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# SegFormer inference
|
||||
# ---------------------------------------------------------------------------
|
||||
def _torch_device(device):
|
||||
global _warned_no_cuda
|
||||
if device == "cuda":
|
||||
@@ -121,14 +147,19 @@ def load_segformer(model_name, device="cuda"):
|
||||
return bundle
|
||||
|
||||
|
||||
def _is_on(model_device, wanted):
|
||||
"""True when `model_device` satisfies `wanted`; an index-less "cuda" accepts any cuda:N."""
|
||||
return model_device.type == wanted.type and (wanted.index is None or model_device.index == wanted.index)
|
||||
|
||||
|
||||
def _ensure_on_device(bundle):
|
||||
model, device = bundle["model"], bundle["device"]
|
||||
if model.device == device:
|
||||
if _is_on(model.device, device):
|
||||
return
|
||||
if device.type == "cuda" and _mm is not None:
|
||||
if device.type == "cuda":
|
||||
# Let ComfyUI evict its own models first instead of OOM-ing on ours.
|
||||
size = sum(p.numel() * p.element_size() for p in model.parameters())
|
||||
_mm.free_memory(size * 2, device)
|
||||
make_room_on_gpu(size * 2)
|
||||
model.to(device)
|
||||
|
||||
|
||||
|
||||
+45
-6
@@ -11,12 +11,34 @@ def resize_image(image, target_size=768):
|
||||
new_h, new_w = target_size, int(w * target_size / h)
|
||||
return image.resize((new_w, new_h), Image.LANCZOS)
|
||||
|
||||
def tensor_to_pil(tensor_image):
|
||||
"""Convert tensor to PIL image"""
|
||||
def tensor_to_pil(tensor_image, *, rgb=False, keep_alpha=False):
|
||||
"""Convert ONE ComfyUI image ([H,W,C] or [1,H,W,C]) to PIL.
|
||||
|
||||
Values are clipped to [0, 1] first, so overshoot from resampling cannot wrap around
|
||||
in uint8. 1/3/4 channels become L/RGB/RGBA. Batches must be converted per image.
|
||||
|
||||
``rgb=True`` always gives an RGB picture: grey is expanded and alpha is dropped. With ``keep_alpha=True`` as
|
||||
well, a 4-channel image stays RGBA instead (every other image is still RGB).
|
||||
"""
|
||||
if tensor_image.dim() == 4:
|
||||
tensor_image = tensor_image.squeeze(0)
|
||||
np_image = (tensor_image.cpu().numpy() * 255).astype(np.uint8)
|
||||
return Image.fromarray(np_image)
|
||||
if tensor_image.shape[0] != 1:
|
||||
raise ValueError(
|
||||
f"tensor_to_pil expects a single image but got a batch of {tensor_image.shape[0]}; "
|
||||
"convert each image of the batch separately"
|
||||
)
|
||||
tensor_image = tensor_image[0]
|
||||
pixels = tensor_image.detach().cpu()
|
||||
if pixels.dtype == torch.bfloat16: # numpy has no bfloat16; other dtypes keep their own rounding
|
||||
pixels = pixels.float()
|
||||
np_image = np.clip(pixels.numpy() * 255.0, 0, 255).astype(np.uint8)
|
||||
if np_image.ndim == 3 and np_image.shape[2] == 1:
|
||||
np_image = np_image[:, :, 0]
|
||||
pil_image = Image.fromarray(np_image)
|
||||
if rgb:
|
||||
kept_modes = ("RGB", "RGBA") if keep_alpha else ("RGB",)
|
||||
if pil_image.mode not in kept_modes:
|
||||
pil_image = pil_image.convert("RGB")
|
||||
return pil_image
|
||||
|
||||
def pil_to_tensor(pil_image):
|
||||
"""Convert PIL image to tensor in ComfyUI format for IMAGE"""
|
||||
@@ -32,4 +54,21 @@ def pil_to_mask(pil_image):
|
||||
if len(np_image.shape) == 3: # Ensure single channel
|
||||
np_image = np_image[:,:,0]
|
||||
tensor = torch.from_numpy(np_image).unsqueeze(0)
|
||||
return tensor
|
||||
return tensor
|
||||
|
||||
def parse_hex_color(value, default=(0, 0, 0, 255)):
|
||||
"""Parse '#RGB', '#RRGGBB' or '#RRGGBBAA' (the '#' is optional) into an (r, g, b, a) tuple of 0-255 ints.
|
||||
|
||||
Without an alpha part the colour is opaque. Text that is not a valid colour returns ``default``.
|
||||
"""
|
||||
text = str(value or "").strip().lstrip("#")
|
||||
if len(text) == 3:
|
||||
text = "".join(ch * 2 for ch in text)
|
||||
if len(text) == 6:
|
||||
text += "ff"
|
||||
if len(text) != 8:
|
||||
return tuple(default)
|
||||
try:
|
||||
return tuple(int(text[i:i + 2], 16) for i in range(0, 8, 2))
|
||||
except ValueError:
|
||||
return tuple(default)
|
||||
|
||||
+92
-103
@@ -1,7 +1,8 @@
|
||||
# Garment mask construction for VTON Ultimate Processor (NH).
|
||||
# Third-party code: FitDiT src/utils_mask.py (https://github.com/BoyuanJiang/FitDiT),
|
||||
# received via ComfyUI-Vton-Mask. License: CC BY-NC-SA 4.0 (non-commercial), not the
|
||||
# pack's MIT license. NH-Nodes changes: lint cleanups and bug fixes. See THIRD_PARTY_NOTICES.md.
|
||||
# pack's MIT license. NH-Nodes changes: lint cleanups, bug fixes and Image.fromarray calls without the mode
|
||||
# argument that Pillow 11.3 deprecated. See THIRD_PARTY_NOTICES.md.
|
||||
import numpy as np
|
||||
import cv2
|
||||
from PIL import Image, ImageDraw
|
||||
@@ -15,11 +16,27 @@ def remove_small(binary_mask, min_area=10000):
|
||||
small_white_regions = [region for region in regions if region.area < min_area]
|
||||
|
||||
for region in small_white_regions:
|
||||
min_row, min_col, max_row, max_col = region.bbox
|
||||
binary_mask[min_row:max_row, min_col:max_col] = 0
|
||||
# Only this region's pixels: its bounding box can overlap a large region.
|
||||
binary_mask[tuple(region.coords.T)] = 0
|
||||
return binary_mask
|
||||
|
||||
|
||||
def _merge_bound(keypoint_bound, cloth_bound, pick):
|
||||
"""Combine a keypoint bound (None when the whole keypoint group is missing) with the garment contour bound."""
|
||||
if keypoint_bound is None:
|
||||
return cloth_bound
|
||||
return pick(keypoint_bound, cloth_bound)
|
||||
|
||||
|
||||
def _fallback_box(x1, y1, x2, y2, shape):
|
||||
"""Whole-image edges only for bounds that neither keypoints nor the garment contour provided."""
|
||||
height, width = shape[:2]
|
||||
return (0 if x1 is None else x1, 0 if y1 is None else y1,
|
||||
width if x2 is None else x2, height if y2 is None else y2)
|
||||
|
||||
|
||||
# pose_data rows: body 0-17, foot 18-23, face 24-91, hands 92-133. A keypoint that was not detected is
|
||||
# negative (DWPose marks it -1 and the node then scales it by the width/height), so only values > 0 count.
|
||||
def get_img_agnostic_upper_rectangle(im_parse, pose_data, offset_top, offset_bottom, offset_left, offset_right):
|
||||
|
||||
faces = pose_data[24:92]
|
||||
@@ -69,48 +86,37 @@ def get_img_agnostic_upper_rectangle(im_parse, pose_data, offset_top, offset_bot
|
||||
except Exception:
|
||||
cloth_exist = False
|
||||
|
||||
x1 = [body[i, 0] for i in [2, 3, 4] if body[i, 0] != 0]+\
|
||||
[hands2[i, 0] for i in [5,9,13] if hands2[i, 0] != 0]+\
|
||||
[hands1[i, 0] for i in [5,9,13] if hands1[i, 0] != 0]
|
||||
if x1:
|
||||
x1 = np.min(x1)
|
||||
else:
|
||||
x1 = 0
|
||||
y1 = [body[i, 1] for i in [2, 5] if body[i, 1] != 0]+\
|
||||
[hands2[i, 1] for i in [5,9,13] if hands2[i, 1] != 0]+\
|
||||
[hands1[i, 1] for i in [5,9,13] if hands1[i, 1] != 0]
|
||||
if y1:
|
||||
y1 = np.min(y1)
|
||||
else:
|
||||
y1 = 0
|
||||
x2 = [body[i, 0] for i in [5, 6, 7] if body[i, 0] != 0]+\
|
||||
[hands1[i, 0] for i in [5,9,13] if hands1[i, 0] != 0]+\
|
||||
[hands2[i, 0] for i in [5,9,13] if hands2[i, 0] != 0]
|
||||
if x2:
|
||||
x2 = np.max(x2)
|
||||
else:
|
||||
x2 = parse_array.shape[1]
|
||||
y2 = [body[i, 1] for i in [8, 11] if body[i, 1] != 0]+\
|
||||
[hands2[i, 1] for i in [5,9,13] if hands2[i, 1] != 0]+\
|
||||
[hands1[i, 1] for i in [5,9,13] if hands1[i, 1] != 0]
|
||||
if y2:
|
||||
y2 = np.max(y2)
|
||||
else:
|
||||
y2 = parse_array.shape[0]
|
||||
x1 = [body[i, 0] for i in [2, 3, 4] if body[i, 0] > 0]+\
|
||||
[hands2[i, 0] for i in [5,9,13] if hands2[i, 0] > 0]+\
|
||||
[hands1[i, 0] for i in [5,9,13] if hands1[i, 0] > 0]
|
||||
x1 = np.min(x1) if x1 else None
|
||||
y1 = [body[i, 1] for i in [2, 5] if body[i, 1] > 0]+\
|
||||
[hands2[i, 1] for i in [5,9,13] if hands2[i, 1] > 0]+\
|
||||
[hands1[i, 1] for i in [5,9,13] if hands1[i, 1] > 0]
|
||||
y1 = np.min(y1) if y1 else None
|
||||
x2 = [body[i, 0] for i in [5, 6, 7] if body[i, 0] > 0]+\
|
||||
[hands1[i, 0] for i in [5,9,13] if hands1[i, 0] > 0]+\
|
||||
[hands2[i, 0] for i in [5,9,13] if hands2[i, 0] > 0]
|
||||
x2 = np.max(x2) if x2 else None
|
||||
y2 = [body[i, 1] for i in [8, 11] if body[i, 1] > 0]+\
|
||||
[hands2[i, 1] for i in [5,9,13] if hands2[i, 1] > 0]+\
|
||||
[hands1[i, 1] for i in [5,9,13] if hands1[i, 1] > 0]
|
||||
y2 = np.max(y2) if y2 else None
|
||||
|
||||
pad_y1 = 20
|
||||
pad_y2 = 10
|
||||
pad_x1 = pad_x2 = 25
|
||||
|
||||
if cloth_exist:
|
||||
x1 = min(x1, _x1)
|
||||
x2 = max(x2, _x2)
|
||||
y1 = min(y1, _y1)
|
||||
if y2>_y2:
|
||||
x1 = _merge_bound(x1, _x1, min)
|
||||
x2 = _merge_bound(x2, _x2, max)
|
||||
y1 = _merge_bound(y1, _y1, min)
|
||||
if y2 is not None and y2>_y2:
|
||||
pad_y2 = 0
|
||||
y2 = max(y2, _y2)
|
||||
y2 = _merge_bound(y2, _y2, max)
|
||||
x1, y1, x2, y2 = _fallback_box(x1, y1, x2, y2, parse_array.shape)
|
||||
|
||||
y_face = [faces[i, 1] for i in [5, 11] if faces[i, 1] != 0]
|
||||
y_face = [faces[i, 1] for i in [5, 11] if faces[i, 1] > 0]
|
||||
|
||||
if y_face:
|
||||
y_face = np.mean(y_face)
|
||||
@@ -120,8 +126,8 @@ def get_img_agnostic_upper_rectangle(im_parse, pose_data, offset_top, offset_bot
|
||||
|
||||
|
||||
agnostic_draw.rectangle((max(0, x1-pad_x1+offset_left), max(0, y1-pad_y1+offset_top), min(x2+pad_x2+offset_right, parse_array.shape[1]), min(y2+pad_y2+offset_bottom, parse_array.shape[0])), 'gray', 'gray')
|
||||
agnostic.paste(img_black, None, Image.fromarray(np.uint8(parse_head * 255), 'L'))
|
||||
# agnostic.paste(img_black, None, Image.fromarray(np.uint8(parse_fixed * 255), 'L'))
|
||||
agnostic.paste(img_black, None, Image.fromarray(np.uint8(parse_head * 255)))
|
||||
# agnostic.paste(img_black, None, Image.fromarray(np.uint8(parse_fixed * 255)))
|
||||
|
||||
mask_gray = agnostic.copy()
|
||||
mask = agnostic.point(lambda p: p == 128 and 255)
|
||||
@@ -180,33 +186,27 @@ def get_img_agnostic_lower_rectangle(im_parse, pose_data, offset_top, offset_bot
|
||||
except Exception:
|
||||
all_cloth_exist = False
|
||||
|
||||
x1 = [body[i, 0] for i in [8, 9, 10] if body[i, 0] != 0]+[foot2[i, 0] for i in [0, 1, 2] if foot2[i, 0] != 0]+\
|
||||
[body[i, 0] for i in [11, 12, 13] if body[i, 0] != 0]+[foot1[i, 0] for i in [0, 1, 2] if foot1[i, 0] != 0]
|
||||
if x1:
|
||||
x1 = np.min(x1)
|
||||
else:
|
||||
x1=0
|
||||
y1 = min(body[8, 1], body[11, 1])
|
||||
x2 = [body[i, 0] for i in [11, 12, 13] if body[i, 0] != 0]+[foot1[i, 0] for i in [0, 1, 2] if foot1[i, 0] != 0]+\
|
||||
[body[i, 0] for i in [8, 9, 10] if body[i, 0] != 0]+[foot2[i, 0] for i in [0, 1, 2] if foot2[i, 0] != 0]
|
||||
if x2:
|
||||
x2 = np.max(x2)
|
||||
else:
|
||||
x2 = parse_array.shape[1]
|
||||
y2 = [body[i, 1] for i in [10, 13] if body[i, 1] != 0]+[foot1[i, 1] for i in [2] if foot2[i, 1] != 0]+\
|
||||
[foot2[i, 1] for i in [2] if foot2[i, 1] != 0]
|
||||
if y2:
|
||||
y2 = np.max(y2)
|
||||
else:
|
||||
y2 = parse_array.shape[0]
|
||||
x1 = [body[i, 0] for i in [8, 9, 10] if body[i, 0] > 0]+[foot2[i, 0] for i in [0, 1, 2] if foot2[i, 0] > 0]+\
|
||||
[body[i, 0] for i in [11, 12, 13] if body[i, 0] > 0]+[foot1[i, 0] for i in [0, 1, 2] if foot1[i, 0] > 0]
|
||||
x1 = np.min(x1) if x1 else None
|
||||
# Keypoints that were not detected are negative (DWPose marks them -1): use only the detected hips.
|
||||
y1 = [body[i, 1] for i in [8, 11] if body[i, 1] > 0]
|
||||
y1 = np.min(y1) if y1 else None
|
||||
x2 = [body[i, 0] for i in [11, 12, 13] if body[i, 0] > 0]+[foot1[i, 0] for i in [0, 1, 2] if foot1[i, 0] > 0]+\
|
||||
[body[i, 0] for i in [8, 9, 10] if body[i, 0] > 0]+[foot2[i, 0] for i in [0, 1, 2] if foot2[i, 0] > 0]
|
||||
x2 = np.max(x2) if x2 else None
|
||||
y2 = [body[i, 1] for i in [10, 13] if body[i, 1] > 0]+[foot1[i, 1] for i in [2] if foot1[i, 1] > 0]+\
|
||||
[foot2[i, 1] for i in [2] if foot2[i, 1] > 0]
|
||||
y2 = np.max(y2) if y2 else None
|
||||
if cloth_exist:
|
||||
x1 = min(x1, _x1)
|
||||
x2 = max(x2, _x2)
|
||||
y1 = min(y1, _y1)
|
||||
y2 = max(y2, _y2)
|
||||
x1 = _merge_bound(x1, _x1, min)
|
||||
x2 = _merge_bound(x2, _x2, max)
|
||||
y1 = _merge_bound(y1, _y1, min)
|
||||
y2 = _merge_bound(y2, _y2, max)
|
||||
elif all_cloth_exist:
|
||||
x1 = min(x1, _x1)
|
||||
x2 = max(x2, _x2)
|
||||
x1 = _merge_bound(x1, _x1, min)
|
||||
x2 = _merge_bound(x2, _x2, max)
|
||||
x1, y1, x2, y2 = _fallback_box(x1, y1, x2, y2, parse_array.shape)
|
||||
|
||||
pad_y1 = 5 if cloth_exist else 50
|
||||
pad_y2 = 10
|
||||
@@ -214,8 +214,8 @@ def get_img_agnostic_lower_rectangle(im_parse, pose_data, offset_top, offset_bot
|
||||
|
||||
|
||||
agnostic_draw.rectangle((max(0, x1-pad_x1+offset_left), max(0, y1-pad_y1+offset_top), min(x2+pad_x2+offset_right, parse_array.shape[1]), min(y2+pad_y2+offset_bottom, parse_array.shape[0])), 'gray', 'gray')
|
||||
agnostic.paste(img_black, None, Image.fromarray(np.uint8(parse_head * 255), 'L'))
|
||||
agnostic.paste(img_black, None, Image.fromarray(np.uint8(parse_fixed * 255), 'L'))
|
||||
agnostic.paste(img_black, None, Image.fromarray(np.uint8(parse_head * 255)))
|
||||
agnostic.paste(img_black, None, Image.fromarray(np.uint8(parse_fixed * 255)))
|
||||
|
||||
mask_gray = agnostic.copy()
|
||||
mask = agnostic.point(lambda p: p == 128 and 255)
|
||||
@@ -266,48 +266,37 @@ def get_img_agnostic_dresses_rectangle(im_parse, pose_data, offset_top, offset_b
|
||||
except Exception:
|
||||
cloth_exist = False
|
||||
|
||||
x1 = [body[i, 0] for i in [2, 3, 4] if body[i, 0] != 0]+\
|
||||
[body[i, 0] for i in [8, 9, 10] if body[i, 0] != 0]+\
|
||||
[foot2[i, 0] for i in [0, 1, 2] if foot2[i, 0] != 0]+\
|
||||
[hands2[i, 0] for i in [5, 9, 13, 17] if hands2[i, 0] != 0]
|
||||
if x1:
|
||||
x1 = np.min(x1)
|
||||
else:
|
||||
x1 = 0
|
||||
y1 = [body[2, 1], body[5, 1]]+\
|
||||
[hands2[i, 1] for i in [5, 9, 13, 17] if hands2[i, 1] != 0]+\
|
||||
[hands1[i, 1] for i in [5, 9, 13, 17] if hands1[i, 1] != 0]
|
||||
if y1:
|
||||
y1 = np.min(y1)
|
||||
else:
|
||||
y1 = 0
|
||||
x2 = [body[i, 0] for i in [5, 6, 7] if body[i, 0] != 0]+\
|
||||
[body[i, 0] for i in [11, 12, 13] if body[i, 0] != 0]+\
|
||||
[foot1[i, 0] for i in [0, 1, 2] if foot1[i, 0] != 0]+\
|
||||
[hands1[i, 0] for i in [5, 9, 13, 17] if hands1[i, 0] != 0]
|
||||
if x2:
|
||||
x2 = np.max(x2)
|
||||
else:
|
||||
x2 = parse_array.shape[1]
|
||||
y2 = [body[i, 1] for i in [10, 13] if body[i, 1] != 0]+[foot1[i, 1] for i in [0, 1, 2] if foot1[i, 1] != 0]+\
|
||||
[foot2[i, 1] for i in [0, 1, 2] if foot2[i, 1] != 0]
|
||||
if y2:
|
||||
y2 = np.max(y2)
|
||||
else:
|
||||
y2 = parse_array.shape[0]
|
||||
x1 = [body[i, 0] for i in [2, 3, 4] if body[i, 0] > 0]+\
|
||||
[body[i, 0] for i in [8, 9, 10] if body[i, 0] > 0]+\
|
||||
[foot2[i, 0] for i in [0, 1, 2] if foot2[i, 0] > 0]+\
|
||||
[hands2[i, 0] for i in [5, 9, 13, 17] if hands2[i, 0] > 0]
|
||||
x1 = np.min(x1) if x1 else None
|
||||
y1 = [body[i, 1] for i in [2, 5] if body[i, 1] > 0]+\
|
||||
[hands2[i, 1] for i in [5, 9, 13, 17] if hands2[i, 1] > 0]+\
|
||||
[hands1[i, 1] for i in [5, 9, 13, 17] if hands1[i, 1] > 0]
|
||||
y1 = np.min(y1) if y1 else None
|
||||
x2 = [body[i, 0] for i in [5, 6, 7] if body[i, 0] > 0]+\
|
||||
[body[i, 0] for i in [11, 12, 13] if body[i, 0] > 0]+\
|
||||
[foot1[i, 0] for i in [0, 1, 2] if foot1[i, 0] > 0]+\
|
||||
[hands1[i, 0] for i in [5, 9, 13, 17] if hands1[i, 0] > 0]
|
||||
x2 = np.max(x2) if x2 else None
|
||||
y2 = [body[i, 1] for i in [10, 13] if body[i, 1] > 0]+[foot1[i, 1] for i in [0, 1, 2] if foot1[i, 1] > 0]+\
|
||||
[foot2[i, 1] for i in [0, 1, 2] if foot2[i, 1] > 0]
|
||||
y2 = np.max(y2) if y2 else None
|
||||
|
||||
if cloth_exist:
|
||||
x1 = min(x1, _x1)
|
||||
x2 = max(x2, _x2)
|
||||
y1 = min(y1, _y1)
|
||||
y2 = max(y2, _y2)
|
||||
x1 = _merge_bound(x1, _x1, min)
|
||||
x2 = _merge_bound(x2, _x2, max)
|
||||
y1 = _merge_bound(y1, _y1, min)
|
||||
y2 = _merge_bound(y2, _y2, max)
|
||||
x1, y1, x2, y2 = _fallback_box(x1, y1, x2, y2, parse_array.shape)
|
||||
|
||||
|
||||
pad_y1 = 20
|
||||
pad_y2 = 10
|
||||
pad_x1 = pad_x2 = 50
|
||||
|
||||
y_face = [faces[i, 1] for i in [5, 11] if faces[i, 1] != 0]
|
||||
y_face = [faces[i, 1] for i in [5, 11] if faces[i, 1] > 0]
|
||||
|
||||
if y_face:
|
||||
y_face = np.mean(y_face)
|
||||
@@ -317,9 +306,9 @@ def get_img_agnostic_dresses_rectangle(im_parse, pose_data, offset_top, offset_b
|
||||
|
||||
|
||||
agnostic_draw.rectangle((max(0, x1-pad_x1+offset_left), max(0, y1-pad_y1+offset_top), min(x2+pad_x2+offset_right, parse_array.shape[1]), min(y2+pad_y2+offset_bottom, parse_array.shape[0])), 'gray', 'gray')
|
||||
agnostic.paste(img_black, None, Image.fromarray(np.uint8(parse_head * 255), 'L'))
|
||||
agnostic.paste(img_black, None, Image.fromarray(np.uint8(parse_shoes * 255), 'L'))
|
||||
agnostic.paste(img_black, None, Image.fromarray(np.uint8(parse_fixed * 255), 'L'))
|
||||
agnostic.paste(img_black, None, Image.fromarray(np.uint8(parse_head * 255)))
|
||||
agnostic.paste(img_black, None, Image.fromarray(np.uint8(parse_shoes * 255)))
|
||||
agnostic.paste(img_black, None, Image.fromarray(np.uint8(parse_fixed * 255)))
|
||||
|
||||
mask_gray = agnostic.copy()
|
||||
mask = agnostic.point(lambda p: p == 128 and 255)
|
||||
|
||||
@@ -0,0 +1,80 @@
|
||||
# Changelog
|
||||
|
||||
All notable changes to NH-Nodes. Versions follow the Comfy Registry releases.
|
||||
"Behaviour change" marks fixes that can change the output of an existing workflow.
|
||||
|
||||
## 1.6.1 — Cleanup and hardening
|
||||
|
||||
- Internal restructuring with no change to node ids, inputs or outputs: the large image-tools module is split into one
|
||||
module per node family, shared helpers live in `core/` (`image_folders`, `image_ops`, `mask_ops`, `model_files`), and
|
||||
the frontend extensions share one widget helper module.
|
||||
- Widgets are hidden with the frontend's own `hidden` property when available (older frontends keep the previous
|
||||
method); values are still saved with the workflow.
|
||||
- `Load Images Folder`: the image count used for auto increment now honours `recursive`; a wired `folder_path` no
|
||||
longer fingerprints ComfyUI's input folder by mistake.
|
||||
- HTTP routes: the folder-count and large-preview routes no longer return absolute paths, answer with proper status
|
||||
codes (malformed JSON → 400) and do their file work off the server's event loop.
|
||||
- `NH Save Large Image` accepts single-channel images; bfloat16 IMAGE tensors convert correctly; image conversions no
|
||||
longer use Pillow's deprecated `fromarray(mode=…)`.
|
||||
|
||||
## 1.6.0 — Performance, caching and VRAM
|
||||
|
||||
### Faster graphs
|
||||
- **Lazy switches:** `If/Else (NH)`, `Switch N (NH)` and `Random Choice (NH)` now ask ComfyUI only for the branch they pick, so the other branches (for example a second sampler) are never executed. The choice itself is unchanged. A connected input whose upstream returns nothing still counts as connected.
|
||||
- `Layer Stack Composite (NH)` is cached like any pure node instead of re-running on every queue.
|
||||
- `Load images matching`: results are cached until the matched files change; the search stays inside the matching root (no more scanning parent folders, drive roots or the models folder). When nothing matches, `found` is `False` and `images` returns the new optional `fallback_images` input (evaluated only on a miss); without a fallback, `images` is blocked with a message and the rest of the workflow keeps running (behaviour change: it used to raise an error). Use `fallback_images` rather than routing `images` through a switch — a blocked input stops the switch too. A drive or network-share root cannot be searched recursively any more.
|
||||
- Indexed loaders (`Load Diffusion Model Index`, `Load VAE Index`, `Load Clip Index`) reuse the model ComfyUI still holds, so editing a slot you did not select no longer reloads it from disk; they keep only a weak reference, so ComfyUI's RAM-pressure cache can still free the model (`--cache-classic` keeps a normal reference). VAEs are loaded with their file metadata and validated like the core loader does.
|
||||
- `Mask Morphology (NH)`: rewritten on torch (CPU or GPU). Soft masks keep their gradient, masks touching the border no longer shrink from the edge, and large kernels are up to 100× faster. Even kernel lengths act as the next odd length (behaviour change).
|
||||
|
||||
### VRAM and models
|
||||
- `VTON Ultimate Processor (NH)`: new optional `keep_model_loaded` (default **off**). Off frees the ONNX Runtime GPU memory after each run that really used the GPU (about 1.5 s slower per run); on keeps today's speed and is much faster for image lists, which ComfyUI runs one item at a time. ComfyUI is asked to make room before the models are loaded on the GPU, except once ONNX Runtime has fallen back to the CPU.
|
||||
- SegFormer nodes now really make ComfyUI unload models when VRAM is tight (a device-index mismatch prevented it).
|
||||
- Model downloads are pinned to tested Hugging Face revisions and only fetch the needed files; a half-downloaded VTON folder is repaired instead of trusted.
|
||||
- `extra_model_paths.yaml` support: the folder names `nh_nodes` (SegFormer models, `.safetensors`) and `vton_mask` (VTON models, `.onnx`) can point to extra locations; existing installs keep working without moving anything.
|
||||
|
||||
### Viewer
|
||||
- `NH Large Image Preview` / `NH Large Image Compare`: browse every image of a batch with ◀ n / N ▶ or the ←/→ keys.
|
||||
|
||||
## 1.5.0 — Correctness fixes
|
||||
|
||||
### Image
|
||||
- **Batches everywhere:** `VTON Ultimate Processor`, `NH Simple Face Paste` and `NH Agnostic Image Generator` now process every image of a batch (smaller inputs with batch 1 are reused for every item). Behaviour change: Face Paste returns the full destination batch instead of one image. VTON stops with a clear error naming the batch index when no person is detected (it used to return a meaningless mask or crash).
|
||||
- **Values are clipped before 8-bit conversion**, so overshoot from resampling no longer wraps bright pixels to black (Face Paste, Mask-Aware Resize, shared helpers).
|
||||
- `NH Mask-Aware Resize`: an empty mask now respects the mode (pad = fit and pad with `pad_color_hex`, crop = cover and centre) instead of stretching.
|
||||
- `NH Agnostic Image Generator`: masks of a different size are resized to the image, 2-D masks are accepted, soft masks keep their gradient, and the `noise` fill is deterministic (behaviour change).
|
||||
- `Layer Layout Composite` / `Layer Stack Composite`: RGBA layers are pasted with their own alpha instead of as opaque rectangles. RGBA backgrounds give RGBA output.
|
||||
- Padding colours accept `#RRGGBBAA`; without an alpha part the padding is opaque (it used to copy the blue value into alpha).
|
||||
- `Save Image Path (NH)`: `filename_prefix` and `new_folder_name` can no longer write outside the chosen folder (on Windows a prefix may not contain `:`, which blocks drive prefixes such as `D:name`); a prefix such as `renders/img` creates the sub-folder. An absolute `folder_path` is still allowed.
|
||||
- `NH Save Large Image`: WEBP (16383 px) and JPEG (65500 px) size limits are checked before anything is written; `--disable-metadata` is respected.
|
||||
- `Image Label (NH)`: much faster font fitting (cached fonts + binary search), same sizes as before.
|
||||
- `Image Grid Composite (NH)`: refuses a canvas that would not fit in the available memory with a clear message instead of crashing.
|
||||
- `NH Large Preview Cache Manager`: actions run on every queue (it used to be cached).
|
||||
|
||||
### Resolution
|
||||
- `NH Remove Resize Padding`: images resized with `cover_crop` are no longer stretched back (that distorted them). With `original_size` the picture is scaled back uniformly into the original frame and the cropped-away borders are filled from the edges; use `NH Restore Image Geometry` to composite onto the original instead.
|
||||
- Resize nodes keep float precision for Lanczos/box/hamming (no 8-bit quantisation).
|
||||
- `NH Smart Resolution Picker`: a stale preset or model name is reported as a warning in `info`.
|
||||
|
||||
### Logic, text and lists
|
||||
- `Math Eval (NH)`: uses a shared safe evaluator. Behaviour change: an invalid expression now shows an error that names the expression instead of silently returning 0. `x / 0`, `x // 0` and `x % 0` give 0 with a warning. numpy/torch scalar inputs work.
|
||||
- `Random Choice (NH)`: the weight list follows the input slots (`weights` item *i* belongs to `input_i`). Behaviour change when some slots are empty.
|
||||
- Seeds are used literally in `Random Choice`, `Prompt Scheduler`, `Text Split Lines` and `Text Random Line`, so seed 0 is reproducible (behaviour change).
|
||||
- `Prompt Template (NH)`: an empty variable no longer glues the neighbouring words together.
|
||||
- `Regex Extract (NH)`: optional groups that did not match give an empty string instead of "no match".
|
||||
- `Pack Universal (NH)`: values are stored by slot, so an empty `input_0` no longer shifts the other indices (behaviour change).
|
||||
- `Compare (NH)` operator `in`: plain membership / substring test (no more fuzzy numeric matching).
|
||||
- `Batch Merge (NH)`: RGB and RGBA batches can be merged; letterbox bars are transparent for both; very thin inputs no longer crash `pad`.
|
||||
- `NH Mask Aspect Ratio Match`: `pad`/`crop` placement fixed (no crash at the image edge, requested size honoured) and the whole batch is processed.
|
||||
- VTON garment masks: undetected keypoints and missing keypoint groups no longer stretch the mask to the image edges; small parsing islands are removed without cutting rectangles out of the garment.
|
||||
- `Universal Slider Builder (NH)` is marked deprecated (it never produced output); existing workflows still load.
|
||||
|
||||
## 1.4.3 — Tooling, licensing and dependencies
|
||||
- CI runs ruff and the test suite on every pull request; publishing to the Comfy Registry only happens when it passes.
|
||||
- `onnxruntime` is no longer installed automatically (the CPU wheel overwrote `onnxruntime-gpu`). `VTON Ultimate Processor` explains how to install it and falls back to the CPU with a warning.
|
||||
- License notices: the VTON code and parsing models are CC BY-NC-SA 4.0 and the face-parsing model is non-commercial — see `THIRD_PARTY_NOTICES.md`.
|
||||
|
||||
## 1.4.2
|
||||
- `NH Smart Resolution Picker`: FLUX.2 latents are 128 channels at 1/16 (images came out twice the size).
|
||||
- `Load Clip Index (NH)`: works with current ComfyUI (`CLIPType`) and lists every CLIP type.
|
||||
- `Counter (NH)`: counts on every queue.
|
||||
- `NH Mask-Aware Resize`: `crop` keeps the whole mask.
|
||||
@@ -72,13 +72,15 @@ class NH_BatchMerge:
|
||||
CATEGORY = "NH-Nodes/Batch"
|
||||
|
||||
def merge(self, batch_a, batch_b=None, batch_c=None, resize_mode="stretch"):
|
||||
# Target size tu batch_a
|
||||
target_h, target_w = batch_a.shape[1], batch_a.shape[2]
|
||||
# Target size va so kenh tu batch_a
|
||||
target_h, target_w, channels = batch_a.shape[1], batch_a.shape[2], batch_a.shape[3]
|
||||
batches = [batch_a]
|
||||
|
||||
for batch in [batch_b, batch_c]:
|
||||
if batch is None:
|
||||
continue
|
||||
# Channels first, so pad bars of an RGB batch merged into RGBA get the same (transparent) alpha.
|
||||
batch = self._match_channels(batch, channels)
|
||||
if batch.shape[1] != target_h or batch.shape[2] != target_w:
|
||||
batch = self._resize_batch(batch, target_h, target_w, resize_mode)
|
||||
batches.append(batch)
|
||||
@@ -86,6 +88,19 @@ class NH_BatchMerge:
|
||||
result = torch.cat(batches, dim=0)
|
||||
return (result, result.shape[0])
|
||||
|
||||
@staticmethod
|
||||
def _match_channels(batch, channels):
|
||||
"""RGB <-> RGBA de torch.cat chay duoc: bo kenh alpha thua, thieu thi them alpha dac (1.0)."""
|
||||
have = batch.shape[-1]
|
||||
if have == channels:
|
||||
return batch
|
||||
if have not in (3, 4) or channels not in (3, 4):
|
||||
raise ValueError(f"Batch Merge: cannot merge a {have}-channel batch with the {channels}-channel batch_a")
|
||||
if have > channels:
|
||||
return batch[..., :channels]
|
||||
opaque = torch.ones((*batch.shape[:-1], 1), dtype=batch.dtype, device=batch.device)
|
||||
return torch.cat([batch, opaque], dim=-1)
|
||||
|
||||
@staticmethod
|
||||
def _resize_batch(batch, th, tw, mode):
|
||||
"""Resize batch ve target size."""
|
||||
@@ -99,7 +114,7 @@ class NH_BatchMerge:
|
||||
# Scale de fill, roi crop
|
||||
_, _, sh, sw = b.shape
|
||||
scale = max(th / sh, tw / sw)
|
||||
new_h, new_w = round(sh * scale), round(sw * scale)
|
||||
new_h, new_w = max(1, round(sh * scale)), max(1, round(sw * scale))
|
||||
b = F.interpolate(b, size=(new_h, new_w), mode='bilinear', align_corners=False)
|
||||
# Center crop
|
||||
y1 = (new_h - th) // 2
|
||||
@@ -110,7 +125,7 @@ class NH_BatchMerge:
|
||||
# Scale de fit, roi pad
|
||||
_, _, sh, sw = b.shape
|
||||
scale = min(th / sh, tw / sw)
|
||||
new_h, new_w = round(sh * scale), round(sw * scale)
|
||||
new_h, new_w = max(1, round(sh * scale)), max(1, round(sw * scale))
|
||||
b = F.interpolate(b, size=(new_h, new_w), mode='bilinear', align_corners=False)
|
||||
# Center pad
|
||||
pad_top = (th - new_h) // 2
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
@@ -10,6 +11,10 @@ class NH_AgnosticImageGenerator:
|
||||
- agnostic_img: ảnh đã fill vùng mask (dùng làm input cho inpaint/VTON)
|
||||
- masked_img: ảnh gốc với vùng mask bị xóa trắng (transparent)
|
||||
- composite: ảnh gốc với viền mask overlay để preview
|
||||
|
||||
Mask là alpha mềm: giá trị trung gian được blend một phần, feathering làm mềm thêm từ chính các giá trị đó.
|
||||
Mask 2-D [H,W] coi như [1,H,W]; mask khác kích thước ảnh được resize về kích thước ảnh.
|
||||
Fill "noise" tất định: ảnh thứ i trong batch dùng seed = i.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
@@ -36,19 +41,30 @@ class NH_AgnosticImageGenerator:
|
||||
results_masked = []
|
||||
results_composite = []
|
||||
|
||||
# Chuẩn hoá mask về [B,H,W] đúng kích thước ảnh (batch 1 được dùng lại cho mọi ảnh ở vòng lặp dưới).
|
||||
if mask.dim() == 2:
|
||||
mask = mask.unsqueeze(0)
|
||||
if mask.shape[-2:] != image.shape[1:3]:
|
||||
mask = F.interpolate(
|
||||
mask.unsqueeze(1).float(), size=tuple(image.shape[1:3]), mode="bilinear", align_corners=False, antialias=True,
|
||||
).squeeze(1)
|
||||
|
||||
for i in range(image.shape[0]):
|
||||
img_np = image[i].cpu().numpy() # (H, W, C) float32 [0,1]
|
||||
m_np = mask[i].cpu().numpy() if i < mask.shape[0] else mask[0].cpu().numpy() # (H, W)
|
||||
|
||||
h, w, c = img_np.shape
|
||||
|
||||
# Tạo alpha mask [0,1]
|
||||
alpha = (m_np > 0.5).astype(np.float32)
|
||||
# Alpha giữ nguyên giá trị mềm của mask (không ngưỡng hoá), kẹp về [0,1]
|
||||
alpha = np.clip(m_np, 0.0, 1.0).astype(np.float32)
|
||||
# Viền preview cần hình dạng rõ nét: lấy ở mức 0.5 của mask gốc (như trước đây), không dùng alpha mềm
|
||||
outline_mask = (m_np > 0.5).astype(np.float32)
|
||||
|
||||
# Feathering: làm mềm biên mask
|
||||
if feathering > 0:
|
||||
kernel_size = feathering * 2 + 1
|
||||
alpha = cv2.GaussianBlur(alpha, (kernel_size, kernel_size), 0)
|
||||
outline_mask = cv2.GaussianBlur(outline_mask, (kernel_size, kernel_size), 0)
|
||||
|
||||
alpha_3ch = alpha[:, :, np.newaxis] # (H, W, 1)
|
||||
|
||||
@@ -57,7 +73,8 @@ class NH_AgnosticImageGenerator:
|
||||
fill = np.full_like(img_np, gray_value)
|
||||
|
||||
elif fill_mode == "noise":
|
||||
noise = np.random.rand(h, w, c).astype(np.float32)
|
||||
# Seed cố định theo index trong batch: cùng input luôn ra cùng noise
|
||||
noise = np.random.default_rng(i).random((h, w, c), dtype=np.float32)
|
||||
# Blend noise với gray trung tâm để bớt harsh
|
||||
fill = noise * noise_strength + gray_value * (1.0 - noise_strength)
|
||||
|
||||
@@ -75,7 +92,7 @@ class NH_AgnosticImageGenerator:
|
||||
masked = np.clip(masked, 0.0, 1.0)
|
||||
|
||||
# --- composite: overlay viền mask lên ảnh gốc ---
|
||||
composite = self._create_composite(img_np, alpha, agnostic)
|
||||
composite = self._create_composite(img_np, outline_mask, agnostic)
|
||||
|
||||
results_agnostic.append(agnostic)
|
||||
results_masked.append(masked)
|
||||
@@ -87,12 +104,12 @@ class NH_AgnosticImageGenerator:
|
||||
return (out_agnostic, out_masked, out_composite)
|
||||
|
||||
@staticmethod
|
||||
def _create_composite(img_np, alpha, agnostic):
|
||||
def _create_composite(img_np, outline_mask, agnostic):
|
||||
"""Tạo preview: nửa trái = ảnh gốc với viền mask, nửa phải = agnostic."""
|
||||
h, w = img_np.shape[:2]
|
||||
|
||||
# Tìm contour của mask để vẽ viền
|
||||
mask_uint8 = np.ascontiguousarray((alpha * 255).astype(np.uint8))
|
||||
mask_uint8 = np.ascontiguousarray((outline_mask * 255).astype(np.uint8))
|
||||
contours, _ = cv2.findContours(mask_uint8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||||
|
||||
# Nửa trái: ảnh gốc + viền xanh
|
||||
|
||||
+40
-21
@@ -2,14 +2,15 @@ import torch
|
||||
import numpy as np
|
||||
import cv2 # Cần thư viện OpenCV
|
||||
|
||||
def get_bounding_box(mask_np):
|
||||
rows = np.any(mask_np, axis=1)
|
||||
cols = np.any(mask_np, axis=0)
|
||||
if not np.any(rows) or not np.any(cols):
|
||||
return None, None, None, None
|
||||
ymin, ymax = np.where(rows)[0][[0, -1]]
|
||||
xmin, xmax = np.where(cols)[0][[0, -1]]
|
||||
return xmin, ymin, xmax, ymax
|
||||
from ...core.mask_ops import mask_bbox
|
||||
|
||||
def _to_uint8(image):
|
||||
# Clip trước khi ép uint8: giá trị lệch khỏi [0,1] (vd. do resample) sẽ bị wrap về gần 0 nếu không clip.
|
||||
return (np.clip(image.cpu().numpy(), 0.0, 1.0) * 255).astype(np.uint8)
|
||||
|
||||
def _as_mask_batch(mask):
|
||||
# MASK 2-D [H,W] không có chiều batch: coi là [1,H,W].
|
||||
return mask.unsqueeze(0) if mask.dim() == 2 else mask
|
||||
|
||||
class NH_SimpleFacePaste:
|
||||
@classmethod
|
||||
@@ -29,20 +30,40 @@ class NH_SimpleFacePaste:
|
||||
CATEGORY = "NH-Nodes/Image"
|
||||
|
||||
def paste_face(self, dest_image, dest_mask, source_image, source_mask, feathering):
|
||||
src_img_np = (source_image[0].cpu().numpy() * 255).astype(np.uint8)
|
||||
src_mask_np = source_mask[0].cpu().numpy()
|
||||
|
||||
dest_img_np = (dest_image[0].cpu().numpy() * 255).astype(np.uint8)
|
||||
dest_mask_np = dest_mask[0].cpu().numpy()
|
||||
dest_mask = _as_mask_batch(dest_mask)
|
||||
source_mask = _as_mask_batch(source_mask)
|
||||
# Output theo batch của dest_image; input khác có batch 1 (hoặc ngắn hơn) được lặp lại theo vòng.
|
||||
output = dest_image.clone()
|
||||
for index in range(dest_image.shape[0]):
|
||||
pasted = self._paste_single(
|
||||
dest_image[index], dest_mask[index % dest_mask.shape[0]],
|
||||
source_image[index % source_image.shape[0]], source_mask[index % source_mask.shape[0]],
|
||||
feathering,
|
||||
)
|
||||
if pasted is not None: # mask rỗng: giữ nguyên ảnh đích của item này
|
||||
output[index] = pasted
|
||||
return (output,)
|
||||
|
||||
src_x1, src_y1, src_x2, src_y2 = get_bounding_box(src_mask_np > 0.5)
|
||||
dest_x1, dest_y1, dest_x2, dest_y2 = get_bounding_box(dest_mask_np > 0.5)
|
||||
@staticmethod
|
||||
def _paste_single(dest_image, dest_mask, source_image, source_mask, feathering):
|
||||
"""Dán mặt của một ảnh nguồn [H,W,C] lên một ảnh đích [H,W,C]; trả None nếu một trong hai mask rỗng."""
|
||||
src_img_np = _to_uint8(source_image)
|
||||
src_mask_np = source_mask.cpu().numpy()
|
||||
|
||||
if src_x1 is None or dest_x1 is None:
|
||||
return (dest_image,)
|
||||
dest_img_np = _to_uint8(dest_image)
|
||||
dest_mask_np = dest_mask.cpu().numpy()
|
||||
|
||||
src_box = mask_bbox(src_mask_np)
|
||||
dest_box = mask_bbox(dest_mask_np)
|
||||
|
||||
if src_box is None or dest_box is None:
|
||||
return None
|
||||
|
||||
src_x1, src_y1, src_x2, src_y2 = src_box
|
||||
dest_x1, dest_y1, dest_x2, dest_y2 = dest_box
|
||||
|
||||
src_face = src_img_np[src_y1:src_y2+1, src_x1:src_x2+1]
|
||||
src_face_mask = (src_mask_np[src_y1:src_y2+1, src_x1:src_x2+1] * 255).astype(np.uint8)
|
||||
src_face_mask = (np.clip(src_mask_np[src_y1:src_y2+1, src_x1:src_x2+1], 0.0, 1.0) * 255).astype(np.uint8)
|
||||
|
||||
dest_width = dest_x2 - dest_x1 + 1
|
||||
dest_height = dest_y2 - dest_y1 + 1
|
||||
@@ -61,10 +82,8 @@ class NH_SimpleFacePaste:
|
||||
|
||||
output_img_np = dest_img_np.copy()
|
||||
output_img_np[dest_y1:dest_y1+dest_height, dest_x1:dest_x1+dest_width] = blended_area.astype(np.uint8)
|
||||
|
||||
output_image = torch.from_numpy(output_img_np.astype(np.float32) / 255.0).unsqueeze(0)
|
||||
|
||||
return (output_image,)
|
||||
return torch.from_numpy(output_img_np.astype(np.float32) / 255.0)
|
||||
|
||||
# --- PHẦN BỊ THIẾU TRƯỚC ĐÂY ĐÃ ĐƯỢC THÊM VÀO ĐÂY ---
|
||||
# Ánh xạ để ComfyUI nhận diện class
|
||||
|
||||
@@ -0,0 +1,681 @@
|
||||
"""Load Images Folder (NH) and Load images matching: read the images of a folder, or of the folder that matches a name."""
|
||||
|
||||
import hashlib
|
||||
import os
|
||||
import re
|
||||
import unicodedata
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
from PIL import Image, ImageOps
|
||||
|
||||
import folder_paths
|
||||
import node_helpers
|
||||
from comfy_execution.graph_utils import ExecutionBlocker
|
||||
|
||||
from ...core.image_folders import IMAGE_EXTENSIONS, collect_image_files, natural_key, resolve_load_dir
|
||||
from ...core.image_ops import pil_to_image_tensor, resize_pil
|
||||
|
||||
|
||||
def _normalize_match_value(value):
|
||||
value = unicodedata.normalize("NFD", str(value or ""))
|
||||
value = "".join(char for char in value if unicodedata.category(char) != "Mn")
|
||||
value = re.sub(r"[_\-]+", " ", value)
|
||||
value = re.sub(r"\s+", " ", value)
|
||||
return value.strip().casefold()
|
||||
|
||||
|
||||
def _path_parts(value):
|
||||
return [part.strip() for part in re.split(r"[\\/]+", str(value or "")) if part.strip()]
|
||||
|
||||
|
||||
def _normalized_path_parts(value):
|
||||
return [part for part in (_normalize_match_value(item) for item in _path_parts(value)) if part]
|
||||
|
||||
|
||||
def _longest_common_path_run(candidate_parts, source_parts):
|
||||
if not candidate_parts or not source_parts:
|
||||
return 0
|
||||
|
||||
previous = [0] * (len(source_parts) + 1)
|
||||
best = 0
|
||||
for candidate_part in candidate_parts:
|
||||
current = [0] * (len(source_parts) + 1)
|
||||
for index, source_part in enumerate(source_parts, start=1):
|
||||
if candidate_part == source_part:
|
||||
current[index] = previous[index - 1] + 1
|
||||
best = max(best, current[index])
|
||||
previous = current
|
||||
return best
|
||||
|
||||
|
||||
def _ordered_path_overlap(candidate_parts, source_parts):
|
||||
if not candidate_parts or not source_parts:
|
||||
return 0
|
||||
|
||||
count = 0
|
||||
cursor = 0
|
||||
for candidate_part in candidate_parts:
|
||||
while cursor < len(source_parts) and source_parts[cursor] != candidate_part:
|
||||
cursor += 1
|
||||
if cursor >= len(source_parts):
|
||||
break
|
||||
count += 1
|
||||
cursor += 1
|
||||
return count
|
||||
|
||||
|
||||
def _candidate_context_score(candidate_path, source_context_paths):
|
||||
if not source_context_paths:
|
||||
return 0
|
||||
|
||||
candidate_parts = _normalized_path_parts(candidate_path)
|
||||
best_score = 0
|
||||
for source_path in source_context_paths:
|
||||
source_parts = _normalized_path_parts(source_path)
|
||||
run = _longest_common_path_run(candidate_parts, source_parts)
|
||||
ordered = _ordered_path_overlap(candidate_parts, source_parts)
|
||||
score = run * 100 + ordered
|
||||
if candidate_parts and len(source_parts) >= 2 and candidate_parts[-1] == source_parts[-2]:
|
||||
score += 10
|
||||
elif candidate_parts and source_parts and candidate_parts[-1] == source_parts[-1]:
|
||||
score += 5
|
||||
best_score = max(best_score, score)
|
||||
return best_score
|
||||
|
||||
|
||||
def _source_segment_set(source_context_paths):
|
||||
segments = set()
|
||||
for path_value in source_context_paths or []:
|
||||
segments.update(_normalized_path_parts(path_value))
|
||||
return segments
|
||||
|
||||
|
||||
def _add_unique_path(paths, path_value):
|
||||
if not path_value:
|
||||
return
|
||||
key = os.path.normcase(os.path.normpath(path_value))
|
||||
if key not in {os.path.normcase(os.path.normpath(path)) for path in paths}:
|
||||
paths.append(path_value)
|
||||
|
||||
|
||||
_SIGNATURE_ENTRY_LIMIT = 20000 # most folders + files Load Images Matching hashes for IS_CHANGED before it gives up
|
||||
|
||||
|
||||
class _NoMatchError(ValueError):
|
||||
"""The search itself came up empty: no matching folder, or no images in it. A bad setting is a plain ValueError."""
|
||||
|
||||
def __init__(self, message, match_text=""):
|
||||
super().__init__(message)
|
||||
self.match_text = match_text
|
||||
|
||||
|
||||
def _hash_parts(hasher, *parts):
|
||||
for part in parts:
|
||||
hasher.update(str(part).encode("utf-8", errors="replace"))
|
||||
hasher.update(b"\0")
|
||||
|
||||
|
||||
def _same_path(first, second):
|
||||
return os.path.normcase(os.path.abspath(first)) == os.path.normcase(os.path.abspath(second))
|
||||
|
||||
|
||||
def _is_inside(path, folder):
|
||||
"""True when ``path`` lies below ``folder`` (``folder`` itself does not count)."""
|
||||
path = os.path.normcase(os.path.abspath(path))
|
||||
folder = os.path.normcase(os.path.abspath(folder))
|
||||
return path != folder and path.startswith(folder.rstrip(os.sep) + os.sep)
|
||||
|
||||
|
||||
def _is_drive_root(path):
|
||||
path = os.path.abspath(path)
|
||||
return os.path.dirname(path) == path
|
||||
|
||||
|
||||
def _entry_is_dir(entry, follow_symlinks=True):
|
||||
try:
|
||||
return entry.is_dir(follow_symlinks=follow_symlinks)
|
||||
except OSError:
|
||||
return False
|
||||
|
||||
|
||||
def _absolute_root(root_path):
|
||||
"""The matching root as an absolute path; a relative one is taken relative to ComfyUI's input folder."""
|
||||
root_path = (root_path or "").strip()
|
||||
if root_path and not os.path.isabs(root_path):
|
||||
root_path = os.path.join(folder_paths.get_input_directory(), root_path)
|
||||
return root_path
|
||||
|
||||
|
||||
def _check_search_root(root_path, recursive):
|
||||
"""Refuse the roots a search must not scan: a whole drive (when it would be walked) and the models folder."""
|
||||
if recursive and _is_drive_root(root_path):
|
||||
raise ValueError(
|
||||
f"Matching root is a drive root ({root_path}); refusing to scan a whole drive. "
|
||||
"Point matching_root_path at the folder that holds the images, or turn recursive off."
|
||||
)
|
||||
if _same_path(root_path, folder_paths.models_dir):
|
||||
raise ValueError(f"Matching root is the ComfyUI models folder ({root_path}), which holds no images to match.")
|
||||
|
||||
|
||||
def _list_folder(folder_path):
|
||||
"""The os.DirEntry objects in ``folder_path`` in natural order, without the ComfyUI models folder; [] if unreadable."""
|
||||
try:
|
||||
with os.scandir(folder_path) as scan:
|
||||
entries = sorted(scan, key=lambda entry: natural_key(entry.name))
|
||||
except OSError:
|
||||
return []
|
||||
models_dir = os.path.normcase(os.path.abspath(folder_paths.models_dir))
|
||||
return [
|
||||
entry for entry in entries
|
||||
if not (_entry_is_dir(entry) and os.path.normcase(os.path.normpath(entry.path)) == models_dir)
|
||||
]
|
||||
|
||||
|
||||
def _child_folders(folder_path):
|
||||
"""Paths of the sub-folders directly inside ``folder_path``, in natural order."""
|
||||
return [entry.path for entry in _list_folder(folder_path) if _entry_is_dir(entry)]
|
||||
|
||||
|
||||
def _walk_folders(root_path, recursive=True):
|
||||
"""Yield ``(folder, entries)`` for ``root_path`` and, when ``recursive``, for every folder below it.
|
||||
|
||||
The matching search's stand-in for os.walk, with limits: it refuses the roots ``_check_search_root`` names,
|
||||
never enters the ComfyUI models folder and, like os.walk, does not follow symlinks into other trees.
|
||||
"""
|
||||
_check_search_root(root_path, recursive)
|
||||
pending = [root_path]
|
||||
while pending:
|
||||
folder = pending.pop()
|
||||
entries = _list_folder(folder)
|
||||
yield folder, entries
|
||||
if recursive:
|
||||
pending.extend(entry.path for entry in reversed(entries) if _entry_is_dir(entry, follow_symlinks=False))
|
||||
|
||||
|
||||
def _context_twins(root_path, context_names, wanted=None):
|
||||
"""Folders that stand in for ``root_path`` under another context folder.
|
||||
|
||||
When the root sits inside a context folder (``.../Nam/...``), the same place under a sibling context folder
|
||||
(``.../Nu/...``) is a twin, and so is the sibling context folder as a whole. ``wanted`` limits the siblings to
|
||||
those names (normalized); all other configured context names count when it is None.
|
||||
"""
|
||||
context = {_normalize_match_value(name) for name in context_names} - {""}
|
||||
if not context:
|
||||
return []
|
||||
|
||||
parts = list(Path(os.path.normpath(root_path)).parts)
|
||||
twins = []
|
||||
for position, part in enumerate(parts):
|
||||
own_name = _normalize_match_value(part)
|
||||
if position == 0 or own_name not in context:
|
||||
continue
|
||||
rest = parts[position + 1:]
|
||||
for sibling_path in _child_folders(str(Path(*parts[:position]))):
|
||||
sibling_name = _normalize_match_value(os.path.basename(sibling_path))
|
||||
if sibling_name == own_name or sibling_name not in context:
|
||||
continue
|
||||
if wanted is not None and sibling_name not in wanted:
|
||||
continue
|
||||
twins.append(sibling_path)
|
||||
if rest and os.path.isdir(os.path.join(sibling_path, *rest)):
|
||||
twins.append(os.path.join(sibling_path, *rest))
|
||||
return twins
|
||||
|
||||
|
||||
def _context_root_candidates(root_path, source_context_paths, context_names=()):
|
||||
"""Folders to search for the match, best context first.
|
||||
|
||||
That is the matching root, the sub-folders of it that the source path names (``Nu`` when the source lives in a
|
||||
``Nu`` folder) and the context twins of the root. Nothing else above the root is listed, let alone walked.
|
||||
"""
|
||||
root_path = _absolute_root(root_path)
|
||||
if not root_path:
|
||||
return []
|
||||
|
||||
candidates = []
|
||||
_add_unique_path(candidates, root_path)
|
||||
|
||||
source_segments = _source_segment_set(source_context_paths)
|
||||
if not source_segments:
|
||||
return candidates
|
||||
|
||||
for child_path in _child_folders(root_path):
|
||||
if _normalize_match_value(os.path.basename(child_path)) in source_segments:
|
||||
_add_unique_path(candidates, child_path)
|
||||
for twin_path in _context_twins(root_path, context_names, source_segments):
|
||||
_add_unique_path(candidates, twin_path)
|
||||
|
||||
candidates.sort(key=lambda item: (-_candidate_context_score(item, source_context_paths), natural_key(item)))
|
||||
return candidates
|
||||
|
||||
|
||||
def _resolve_matching_folder(root_path, match_text, match_mode, recursive, source_context_paths=None):
|
||||
root_path = _absolute_root(root_path)
|
||||
if not root_path:
|
||||
raise ValueError("Matching root path is empty")
|
||||
if not os.path.isdir(root_path):
|
||||
raise ValueError(f"Matching root folder does not exist: {root_path}")
|
||||
|
||||
normalized_match = _normalize_match_value(match_text)
|
||||
if not normalized_match:
|
||||
raise ValueError("Match text is empty")
|
||||
|
||||
candidates = [root_path]
|
||||
for _folder, entries in _walk_folders(root_path, recursive):
|
||||
candidates.extend(entry.path for entry in entries if _entry_is_dir(entry))
|
||||
|
||||
exact_matches = []
|
||||
contains_matches = []
|
||||
for candidate in candidates:
|
||||
normalized_name = _normalize_match_value(os.path.basename(candidate))
|
||||
if normalized_name == normalized_match:
|
||||
exact_matches.append(candidate)
|
||||
if normalized_match in normalized_name:
|
||||
contains_matches.append(candidate)
|
||||
|
||||
matches = exact_matches if match_mode == "exact" else contains_matches or exact_matches
|
||||
if not matches:
|
||||
raise _NoMatchError(f"No matching folder found for '{match_text}' under: {root_path}", match_text)
|
||||
|
||||
if source_context_paths:
|
||||
matches.sort(key=lambda item: (
|
||||
-_candidate_context_score(item, source_context_paths),
|
||||
natural_key(os.path.relpath(item, root_path)),
|
||||
))
|
||||
else:
|
||||
matches.sort(key=lambda item: natural_key(os.path.relpath(item, root_path)))
|
||||
return matches[0]
|
||||
|
||||
|
||||
def _resolve_best_matching_folder(root_path, match_text, match_mode, recursive, source_context_paths=None, context_names=()):
|
||||
root_path = _absolute_root(root_path)
|
||||
if not root_path:
|
||||
raise ValueError("Matching root path is empty")
|
||||
_check_search_root(root_path, recursive)
|
||||
|
||||
search_roots = _context_root_candidates(root_path, source_context_paths, context_names)
|
||||
if recursive:
|
||||
# Walking a folder covers everything below it, so a search root inside another one is not walked twice.
|
||||
search_roots = [path for path in search_roots if not any(_is_inside(path, other) for other in search_roots)]
|
||||
|
||||
matched_folders = []
|
||||
errors = []
|
||||
for search_root in search_roots:
|
||||
try:
|
||||
matched_folder = _resolve_matching_folder(
|
||||
search_root,
|
||||
match_text,
|
||||
match_mode,
|
||||
recursive,
|
||||
source_context_paths=source_context_paths,
|
||||
)
|
||||
_add_unique_path(matched_folders, matched_folder)
|
||||
except ValueError as exc:
|
||||
errors.append(exc)
|
||||
|
||||
if not matched_folders:
|
||||
# A setting that cannot work (missing root, empty text) is reported as such; otherwise nothing matched.
|
||||
for exc in errors:
|
||||
if not isinstance(exc, _NoMatchError):
|
||||
raise exc
|
||||
raise _NoMatchError(f"No matching folder found for '{match_text}' under: {root_path}", match_text)
|
||||
|
||||
matched_folders.sort(key=lambda item: (-_candidate_context_score(item, source_context_paths), natural_key(item)))
|
||||
return matched_folders[0]
|
||||
|
||||
|
||||
def _parse_context_folder_names(value):
|
||||
names = [item.strip() for item in re.split(r"[\n,;|]+", str(value or "")) if item.strip()]
|
||||
return names
|
||||
|
||||
|
||||
def _detect_context_folder(paths, context_folder_names):
|
||||
if not context_folder_names:
|
||||
return ""
|
||||
|
||||
normalized_context = {_normalize_match_value(name): name for name in context_folder_names}
|
||||
for path_value in paths:
|
||||
for part in re.split(r"[\\/]+", str(path_value or "")):
|
||||
match = normalized_context.get(_normalize_match_value(part))
|
||||
if match:
|
||||
return match
|
||||
return ""
|
||||
|
||||
|
||||
def _matching_selection(
|
||||
source_folder_path,
|
||||
matching_root_path,
|
||||
match_text,
|
||||
match_mode,
|
||||
recursive,
|
||||
index,
|
||||
step,
|
||||
context_folder_names,
|
||||
source_file_names,
|
||||
person_folder_path=None,
|
||||
):
|
||||
"""Work out what Load images matching reads: one ``(match_text, folder, [file names])`` per match item.
|
||||
|
||||
Raises _NoMatchError when an item finds no folder, or a folder without images, and ValueError for a setting
|
||||
that cannot work (missing root, empty match text, a root that must not be scanned).
|
||||
"""
|
||||
if not source_folder_path and person_folder_path:
|
||||
source_folder_path = person_folder_path
|
||||
|
||||
source_paths = [line.strip() for line in str(source_folder_path or "").splitlines() if line.strip()]
|
||||
first_source_path = source_paths[0] if source_paths else ""
|
||||
resolved_source_folder = resolve_load_dir(first_source_path)
|
||||
selected_match_text = (match_text or "").strip()
|
||||
source_names = [line.strip() for line in str(source_file_names or "").splitlines() if line.strip()]
|
||||
context_names = _parse_context_folder_names(context_folder_names)
|
||||
match_items = []
|
||||
|
||||
if selected_match_text:
|
||||
item_count = max(len(source_paths), len(source_names), 1)
|
||||
for item_index in range(item_count):
|
||||
match_items.append((
|
||||
selected_match_text,
|
||||
source_paths[item_index] if item_index < len(source_paths) else "",
|
||||
source_names[item_index] if item_index < len(source_names) else "",
|
||||
))
|
||||
elif source_names:
|
||||
for item_index, source_name in enumerate(source_names):
|
||||
source_dir = os.path.dirname(source_name)
|
||||
source_match_text = os.path.basename(os.path.normpath(source_dir)) if source_dir else os.path.basename(os.path.normpath(resolved_source_folder))
|
||||
match_items.append((
|
||||
source_match_text,
|
||||
source_paths[item_index] if item_index < len(source_paths) else "",
|
||||
source_name,
|
||||
))
|
||||
else:
|
||||
selected_match_text = os.path.basename(os.path.normpath(resolved_source_folder))
|
||||
match_items.append((selected_match_text, first_source_path, ""))
|
||||
|
||||
plan = []
|
||||
for current_match_text, current_source_path, current_source_name in match_items:
|
||||
context_paths = [current_source_path, current_source_name, resolved_source_folder]
|
||||
context_folder = _detect_context_folder(context_paths, context_names)
|
||||
if context_folder:
|
||||
context_paths.append(context_folder)
|
||||
matched_folder = _resolve_best_matching_folder(
|
||||
matching_root_path,
|
||||
current_match_text,
|
||||
match_mode,
|
||||
recursive,
|
||||
source_context_paths=context_paths,
|
||||
context_names=context_names,
|
||||
)
|
||||
try:
|
||||
image_files = collect_image_files(matched_folder)
|
||||
except ValueError as exc: # no images in the folder that matched
|
||||
raise _NoMatchError(str(exc), current_match_text) from exc
|
||||
start_index = index % len(image_files)
|
||||
selected_files = [image_files[(start_index + offset) % len(image_files)] for offset in range(step)]
|
||||
plan.append((current_match_text, matched_folder, selected_files))
|
||||
return plan
|
||||
|
||||
|
||||
def _selection_signature(plan):
|
||||
"""Digest of what a matching plan reads: the folders chosen, and the selected files with size and mtime."""
|
||||
hasher = hashlib.sha256()
|
||||
for current_match_text, matched_folder, selected_files in plan:
|
||||
_hash_parts(hasher, current_match_text, os.path.abspath(matched_folder))
|
||||
for file_name in selected_files:
|
||||
stat = os.stat(os.path.join(matched_folder, file_name))
|
||||
_hash_parts(hasher, file_name, stat.st_size, stat.st_mtime_ns)
|
||||
return hasher.hexdigest()
|
||||
|
||||
|
||||
def _tree_signature(roots, limit=None):
|
||||
"""Digest of every folder and image below ``roots`` (names, sizes, mtimes); None beyond ``limit`` entries."""
|
||||
limit = _SIGNATURE_ENTRY_LIMIT if limit is None else limit
|
||||
hasher = hashlib.sha256()
|
||||
seen = 0
|
||||
for root in roots:
|
||||
_hash_parts(hasher, "root", os.path.abspath(root))
|
||||
for folder, entries in _walk_folders(root, recursive=True):
|
||||
_hash_parts(hasher, "folder", os.path.relpath(folder, root))
|
||||
for entry in entries:
|
||||
seen += 1
|
||||
if seen > limit:
|
||||
return None
|
||||
if _entry_is_dir(entry):
|
||||
_hash_parts(hasher, "dir", entry.name)
|
||||
elif Path(entry.name).suffix.lower() in IMAGE_EXTENSIONS:
|
||||
try:
|
||||
stat = entry.stat()
|
||||
except OSError:
|
||||
_hash_parts(hasher, "image", entry.name)
|
||||
else:
|
||||
_hash_parts(hasher, "image", entry.name, stat.st_size, stat.st_mtime_ns)
|
||||
return hasher.hexdigest()
|
||||
|
||||
|
||||
def _load_image_batch(paths):
|
||||
"""The images at ``paths`` as one RGB IMAGE batch; every image after the first is scaled to the size of the first."""
|
||||
images = []
|
||||
target_size = None
|
||||
for path in paths:
|
||||
img = node_helpers.pillow(Image.open, path)
|
||||
img = node_helpers.pillow(ImageOps.exif_transpose, img).convert("RGB")
|
||||
if target_size is None:
|
||||
target_size = img.size
|
||||
elif img.size != target_size:
|
||||
img = resize_pil(img, target_size, "lanczos")
|
||||
images.append(pil_to_image_tensor(img))
|
||||
return torch.cat(images, dim=0)
|
||||
|
||||
|
||||
class NH_LoadImagesFromFolder:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"folder_path": ("STRING", {"default": "", "multiline": False}),
|
||||
"index": ("INT", {"default": 0, "min": 0, "max": 999999}),
|
||||
"step": ("INT", {"default": 1, "min": 1, "max": 9999}),
|
||||
"seed_mode": (["fixed", "increment", "decrement", "random"],),
|
||||
"recursive": ("BOOLEAN", {"default": False}),
|
||||
"filename_extension": ("BOOLEAN", {"default": True}),
|
||||
"path_output": (["file_path", "folder_path"], {"default": "file_path"}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "STRING", "STRING")
|
||||
RETURN_NAMES = ("images", "filename", "path")
|
||||
FUNCTION = "load_images"
|
||||
CATEGORY = "NH-Nodes/Image"
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, folder_path, index, step, seed_mode, recursive=False, filename_extension=True, path_output="file_path", **kwargs):
|
||||
# An input wired from another node arrives here as None, so the folder is unknown (resolving None would
|
||||
# fingerprint ComfyUI's input folder instead). The upstream value is already part of ComfyUI's cache key, so a
|
||||
# constant keeps an unchanged graph cached while a new path still runs the node again.
|
||||
if folder_path is None:
|
||||
return "wired-folder-path"
|
||||
try:
|
||||
resolved_folder = resolve_load_dir(folder_path)
|
||||
image_files = collect_image_files(resolved_folder, recursive=recursive)
|
||||
step = max(1, int(step))
|
||||
start_index = int(index) % len(image_files)
|
||||
|
||||
if seed_mode == "decrement":
|
||||
indices = [((start_index - offset) % len(image_files)) for offset in range(step)]
|
||||
else:
|
||||
indices = [((start_index + offset) % len(image_files)) for offset in range(step)]
|
||||
|
||||
hasher = hashlib.sha256()
|
||||
_hash_parts(
|
||||
hasher,
|
||||
os.path.abspath(resolved_folder),
|
||||
start_index,
|
||||
step,
|
||||
seed_mode,
|
||||
bool(recursive),
|
||||
bool(filename_extension),
|
||||
path_output,
|
||||
)
|
||||
for selected_index in indices:
|
||||
rel_path = image_files[selected_index]
|
||||
stat = os.stat(os.path.join(resolved_folder, rel_path))
|
||||
_hash_parts(hasher, rel_path, stat.st_size, stat.st_mtime_ns)
|
||||
|
||||
return hasher.hexdigest()
|
||||
except Exception as exc:
|
||||
return f"error:{type(exc).__name__}:{exc}"
|
||||
|
||||
def load_images(self, folder_path, index, step, seed_mode, recursive=False, filename_extension=True, path_output="file_path", **kwargs):
|
||||
resolved_folder = resolve_load_dir(folder_path)
|
||||
image_files = collect_image_files(resolved_folder, recursive=recursive)
|
||||
start_index = index % len(image_files)
|
||||
|
||||
if seed_mode == "decrement":
|
||||
indices = [((start_index - offset) % len(image_files)) for offset in range(step)]
|
||||
else:
|
||||
indices = [((start_index + offset) % len(image_files)) for offset in range(step)]
|
||||
|
||||
selected_files = [image_files[i] for i in indices]
|
||||
full_paths = [os.path.join(resolved_folder, file_name) for file_name in selected_files]
|
||||
if path_output == "folder_path":
|
||||
output_paths = [os.path.dirname(full_path) for full_path in full_paths]
|
||||
else:
|
||||
output_paths = full_paths
|
||||
|
||||
display_filenames = []
|
||||
for file_name in selected_files:
|
||||
base_name = os.path.basename(file_name)
|
||||
if not filename_extension:
|
||||
base_name = os.path.splitext(base_name)[0]
|
||||
display_filenames.append(base_name)
|
||||
|
||||
return (_load_image_batch(full_paths), "\n".join(display_filenames), "\n".join(output_paths))
|
||||
|
||||
|
||||
class NH_LoadImagesMatching:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"source_folder_path": ("STRING", {"default": "", "multiline": False}),
|
||||
"matching_root_path": ("STRING", {"default": "", "multiline": False}),
|
||||
"match_text": ("STRING", {"default": "", "multiline": False}),
|
||||
"match_mode": (["contains", "exact"],),
|
||||
"recursive": ("BOOLEAN", {"default": True}),
|
||||
"index": ("INT", {"default": 0, "min": 0, "max": 999999}),
|
||||
"step": ("INT", {"default": 1, "min": 1, "max": 9999}),
|
||||
"repeat_count": ("INT", {"default": 1, "min": 1, "max": 9999}),
|
||||
"context_folder_names": ("STRING", {"default": "Nam\nNu", "multiline": True}),
|
||||
},
|
||||
"optional": {
|
||||
"source_file_names": ("STRING", {"default": "", "multiline": True}),
|
||||
"fallback_images": ("IMAGE", {
|
||||
"lazy": True,
|
||||
"tooltip": (
|
||||
"Comes out as images when nothing matches, and is only evaluated then. Connect it here "
|
||||
"instead of routing images through a switch or If/Else: a node that receives the held-back "
|
||||
"images never runs."
|
||||
),
|
||||
}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "STRING", "STRING", "STRING", "BOOLEAN")
|
||||
RETURN_NAMES = ("images", "filename", "path", "matched_text", "found")
|
||||
FUNCTION = "load_matching"
|
||||
CATEGORY = "NH-Nodes/Image"
|
||||
DESCRIPTION = (
|
||||
"Load the images of the folder under matching_root_path whose name matches the source folder (or match_text). "
|
||||
"When nothing matches, found is False and images are held back: every node that receives them is skipped, "
|
||||
"a switch or If/Else included. To carry on with another image, connect it to fallback_images instead of "
|
||||
"routing images through a switch or If/Else: it is only evaluated when nothing matches, and then comes out "
|
||||
"as images (found stays False)."
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, source_folder_path, matching_root_path, match_text, match_mode, recursive, index, step, repeat_count, context_folder_names="Nam\nNu", source_file_names="", **kwargs):
|
||||
# ComfyUI's cache key already holds every setting and every upstream result; this adds what is on disk:
|
||||
# the folder(s) the search picks and the selected files with their size and time, so an edit re-runs the node.
|
||||
# An input wired from another node arrives here as None, which leaves the folder unknown: the whole matching
|
||||
# tree is fingerprinted instead. A search that finds nothing, or a tree too big to hash cheaply, gives NaN:
|
||||
# that result is recomputed every time (so a "not found" is never remembered).
|
||||
unknown = float("nan")
|
||||
try:
|
||||
settings = (source_folder_path, matching_root_path, match_text, match_mode, recursive, index, step, context_folder_names, source_file_names)
|
||||
if None not in settings:
|
||||
plan = _matching_selection(
|
||||
source_folder_path, matching_root_path, match_text, match_mode, recursive, index, step,
|
||||
context_folder_names, source_file_names, kwargs.get("person_folder_path"),
|
||||
)
|
||||
return _selection_signature(plan)
|
||||
root_path = _absolute_root(matching_root_path)
|
||||
if not root_path or context_folder_names is None:
|
||||
return unknown
|
||||
roots = [root_path] + _context_twins(root_path, _parse_context_folder_names(context_folder_names))
|
||||
signature = _tree_signature(roots)
|
||||
return unknown if signature is None else signature
|
||||
except (ValueError, OSError):
|
||||
return unknown
|
||||
|
||||
def check_lazy_status(self, source_folder_path, matching_root_path, match_text, match_mode, recursive, index, step, repeat_count, context_folder_names="Nam\nNu", source_file_names="", **kwargs):
|
||||
# fallback_images is asked for only when it is wired (ComfyUI sends its key, with None until the input has
|
||||
# run) and the search finds nothing; on a match its upstream nodes never run. The search is the one
|
||||
# load_matching makes. A setting that cannot work asks for nothing: load_matching raises it right after.
|
||||
if "fallback_images" not in kwargs or kwargs["fallback_images"] is not None:
|
||||
return []
|
||||
try:
|
||||
_matching_selection(
|
||||
source_folder_path, matching_root_path, match_text, match_mode, recursive, index, step,
|
||||
context_folder_names, source_file_names, kwargs.get("person_folder_path"),
|
||||
)
|
||||
except _NoMatchError:
|
||||
return ["fallback_images"]
|
||||
except (ValueError, OSError):
|
||||
pass
|
||||
return []
|
||||
|
||||
def load_matching(self, source_folder_path, matching_root_path, match_text, match_mode, recursive, index, step, repeat_count, context_folder_names="Nam\nNu", source_file_names="", fallback_images=None, **kwargs):
|
||||
try:
|
||||
plan = _matching_selection(
|
||||
source_folder_path, matching_root_path, match_text, match_mode, recursive, index, step,
|
||||
context_folder_names, source_file_names, kwargs.get("person_folder_path"),
|
||||
)
|
||||
except _NoMatchError as exc:
|
||||
# Nothing to load for this request. `found` lets a workflow carry on with something else, and so does
|
||||
# fallback_images, which comes out as the images. Without it the images are held back (whatever needs
|
||||
# them is told why). A setting that cannot work still raises.
|
||||
message = str(exc)
|
||||
if fallback_images is not None:
|
||||
print(f"[NH-Nodes] Load images matching: {message} (using fallback_images)")
|
||||
return (fallback_images, "", "", exc.match_text, False)
|
||||
print(f"[NH-Nodes] Load images matching: {message}")
|
||||
return (ExecutionBlocker(message), "", "", exc.match_text, False)
|
||||
|
||||
selected_files = []
|
||||
matched_folders = []
|
||||
matched_texts = []
|
||||
|
||||
for current_match_text, matched_folder, file_names in plan:
|
||||
for file_name in file_names:
|
||||
selected_files.append(file_name)
|
||||
matched_folders.append(matched_folder)
|
||||
matched_texts.append(current_match_text)
|
||||
|
||||
images = _load_image_batch(
|
||||
[os.path.join(folder, file_name) for folder, file_name in zip(matched_folders, selected_files)]
|
||||
)
|
||||
if repeat_count > 1:
|
||||
images = images.repeat(repeat_count, 1, 1, 1)
|
||||
|
||||
return (images, "\n".join(selected_files), "\n".join(matched_folders), "\n".join(matched_texts), True)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"NH_LoadImagesFromFolder": NH_LoadImagesFromFolder,
|
||||
"NH_LoadImagesMatching": NH_LoadImagesMatching,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"NH_LoadImagesFromFolder": "Load Images Folder (NH)",
|
||||
"NH_LoadImagesMatching": "Load images matching",
|
||||
}
|
||||
@@ -0,0 +1,95 @@
|
||||
"""Image Compare (NH): two images side by side, split, blended or as a difference."""
|
||||
|
||||
import torch
|
||||
from PIL import Image, ImageChops
|
||||
|
||||
from ...core.utils import tensor_to_pil
|
||||
from ...core.image_ops import pil_to_image_tensor, resize_pil
|
||||
|
||||
|
||||
def _match_batch_sizes(batch_a, batch_b):
|
||||
count_a = batch_a.shape[0]
|
||||
count_b = batch_b.shape[0]
|
||||
|
||||
if count_a == count_b:
|
||||
return batch_a, batch_b
|
||||
if count_a == 1:
|
||||
return batch_a.repeat(count_b, 1, 1, 1), batch_b
|
||||
if count_b == 1:
|
||||
return batch_a, batch_b.repeat(count_a, 1, 1, 1)
|
||||
|
||||
shared = min(count_a, count_b)
|
||||
print(f"[NH-Nodes] Image Compare: batch mismatch ({count_a} vs {count_b}), truncating to {shared}")
|
||||
return batch_a[:shared], batch_b[:shared]
|
||||
|
||||
|
||||
class NH_ImageCompare:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"image_a": ("IMAGE",),
|
||||
"image_b": ("IMAGE",),
|
||||
"compare_mode": (
|
||||
[
|
||||
"side_by_side",
|
||||
"top_bottom",
|
||||
"split_vertical",
|
||||
"split_horizontal",
|
||||
"difference",
|
||||
"overlay",
|
||||
],
|
||||
),
|
||||
"split_position": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
"overlay_opacity": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
"resize_method": (["lanczos", "bicubic", "bilinear", "nearest"],),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("image",)
|
||||
FUNCTION = "compare"
|
||||
CATEGORY = "NH-Nodes/Image"
|
||||
|
||||
def compare(self, image_a, image_b, compare_mode, split_position, overlay_opacity, resize_method):
|
||||
batch_a, batch_b = _match_batch_sizes(image_a, image_b)
|
||||
compared = []
|
||||
|
||||
for tensor_a, tensor_b in zip(batch_a, batch_b):
|
||||
pil_a = tensor_to_pil(tensor_a, rgb=True)
|
||||
pil_b = resize_pil(tensor_to_pil(tensor_b, rgb=True), pil_a.size, resize_method)
|
||||
width, height = pil_a.size
|
||||
|
||||
if compare_mode == "side_by_side":
|
||||
canvas = Image.new("RGB", (width * 2, height))
|
||||
canvas.paste(pil_a, (0, 0))
|
||||
canvas.paste(pil_b, (width, 0))
|
||||
elif compare_mode == "top_bottom":
|
||||
canvas = Image.new("RGB", (width, height * 2))
|
||||
canvas.paste(pil_a, (0, 0))
|
||||
canvas.paste(pil_b, (0, height))
|
||||
elif compare_mode == "split_vertical":
|
||||
split_x = max(0, min(width, int(round(width * split_position))))
|
||||
canvas = pil_b.copy()
|
||||
canvas.paste(pil_a.crop((0, 0, split_x, height)), (0, 0))
|
||||
elif compare_mode == "split_horizontal":
|
||||
split_y = max(0, min(height, int(round(height * split_position))))
|
||||
canvas = pil_b.copy()
|
||||
canvas.paste(pil_a.crop((0, 0, width, split_y)), (0, 0))
|
||||
elif compare_mode == "difference":
|
||||
canvas = ImageChops.difference(pil_a, pil_b)
|
||||
else:
|
||||
canvas = Image.blend(pil_a, pil_b, overlay_opacity)
|
||||
|
||||
compared.append(pil_to_image_tensor(canvas))
|
||||
|
||||
return (torch.cat(compared, dim=0),)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"NH_ImageCompare": NH_ImageCompare,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"NH_ImageCompare": "Image Compare (NH)",
|
||||
}
|
||||
@@ -1,32 +1,85 @@
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image, ImageColor
|
||||
|
||||
from ...core.utils import pil_to_tensor, tensor_to_pil
|
||||
|
||||
try:
|
||||
import psutil
|
||||
except ImportError: # ComfyUI installs psutil; without it the memory guard below is simply skipped
|
||||
psutil = None
|
||||
|
||||
|
||||
_RESAMPLING = getattr(Image, "Resampling", Image)
|
||||
|
||||
# Size guard for the output canvas. It is about RAM, not about Pillow's decompression-bomb limit
|
||||
# (Image.MAX_IMAGE_PIXELS): that limit is process-wide, other packs raise or disable it, and it only applies to
|
||||
# Image.open, never to an IMAGE that stays in memory.
|
||||
#
|
||||
# Peak RAM per canvas pixel, reached inside pil_to_tensor (measured 28.1-28.2 B/px on Windows, Pillow 11, NumPy 2):
|
||||
# 4 B the PIL canvas (Pillow stores an RGB image as 4 bytes per pixel)
|
||||
# 12 B the float32 array from astype(np.float32)
|
||||
# 12 B the float32 result of "/ 255.0" (NumPy only reuses that temporary on Linux)
|
||||
# By then the fitted cells and the uint8 array from np.array() are freed, and torch.from_numpy() shares the float32
|
||||
# result, so the IMAGE tensor adds nothing (its 12 B/px are what stays alive after the node returns). The input
|
||||
# images were converted before the check, so psutil's "available" already leaves their memory out.
|
||||
_PEAK_BYTES_PER_PIXEL = 4 + 12 + 12
|
||||
# Share of the available RAM the node may plan to use; the rest is left to the OS and the other programs.
|
||||
_USABLE_MEMORY_FRACTION = 0.85
|
||||
# Image.new raises a bare MemoryError for an image wider than INT_MAX // 4 - 1 px, whatever the memory; the node's
|
||||
# own limits never get near it. The same ceiling is applied to the height to keep one rule.
|
||||
_PIL_MAX_SIDE = (2**31 - 1) // 4 - 1
|
||||
|
||||
|
||||
def _estimate_canvas_bytes(canvas_width, canvas_height):
|
||||
"""Peak RAM in bytes the node needs to turn a canvas of this size into an IMAGE tensor."""
|
||||
return _PEAK_BYTES_PER_PIXEL * canvas_width * canvas_height
|
||||
|
||||
|
||||
def _available_memory_bytes():
|
||||
"""Free physical RAM in bytes, or None when it cannot be read (the guard is advisory, so it then stays quiet)."""
|
||||
if psutil is None:
|
||||
return None
|
||||
try:
|
||||
return int(psutil.virtual_memory().available)
|
||||
except Exception: # whatever the platform backend raises must not stop a canvas that may well fit
|
||||
return None
|
||||
|
||||
|
||||
def _format_size(nbytes):
|
||||
return f"{nbytes / 2**30:.1f} GiB" if nbytes >= 2**30 else f"{nbytes / 2**20:.0f} MiB"
|
||||
|
||||
|
||||
def _check_canvas_size(canvas_width, canvas_height, columns, rows, cell_width, cell_height):
|
||||
"""Raise ValueError, before anything is allocated, for a canvas the node cannot build in the free memory."""
|
||||
layout = f"{canvas_width}x{canvas_height} px canvas ({columns}x{rows} cells of {cell_width}x{cell_height} px)"
|
||||
hint = "Lower columns/rows, the cell size or the gaps."
|
||||
|
||||
if max(canvas_width, canvas_height) > _PIL_MAX_SIDE:
|
||||
raise ValueError(
|
||||
f"Image Grid Composite: the {layout} is too large: Pillow cannot create an image wider or taller "
|
||||
f"than {_PIL_MAX_SIDE:,} px. {hint}"
|
||||
)
|
||||
|
||||
available = _available_memory_bytes()
|
||||
if available is None:
|
||||
return
|
||||
needed = _estimate_canvas_bytes(canvas_width, canvas_height)
|
||||
if needed > available * _USABLE_MEMORY_FRACTION:
|
||||
raise ValueError(
|
||||
f"Image Grid Composite: the {layout} is too large for the free memory: it needs about "
|
||||
f"{_format_size(needed)} of RAM and {_format_size(available)} is available. {hint}"
|
||||
)
|
||||
|
||||
|
||||
def _parse_hex_color(value, default):
|
||||
# Pillow's own parser, not core.utils.parse_hex_color: it also reads CSS names and rgb(...), and it needs the '#'.
|
||||
# Kept so that every background colour the node accepted before still gives the same fill.
|
||||
try:
|
||||
return ImageColor.getrgb((value or "").strip())
|
||||
except Exception:
|
||||
return default
|
||||
|
||||
|
||||
def _tensor_to_pil(image_tensor):
|
||||
array = np.clip(255.0 * image_tensor.cpu().numpy(), 0, 255).astype(np.uint8)
|
||||
if array.ndim == 2:
|
||||
return Image.fromarray(array).convert("RGB")
|
||||
if array.ndim == 3 and array.shape[-1] == 1:
|
||||
return Image.fromarray(array[:, :, 0]).convert("RGB")
|
||||
return Image.fromarray(array).convert("RGB")
|
||||
|
||||
|
||||
def _pil_to_tensor(image):
|
||||
array = np.array(image).astype(np.float32) / 255.0
|
||||
return torch.from_numpy(array).unsqueeze(0)
|
||||
|
||||
|
||||
def _first_value(value):
|
||||
if isinstance(value, (list, tuple)):
|
||||
if not value:
|
||||
@@ -46,10 +99,10 @@ def _flatten_images(images):
|
||||
raise TypeError(f"Expected IMAGE tensor input, got {type(image_item).__name__}")
|
||||
|
||||
if image_item.ndim == 3:
|
||||
pil_images.append(_tensor_to_pil(image_item))
|
||||
pil_images.append(tensor_to_pil(image_item, rgb=True))
|
||||
elif image_item.ndim == 4:
|
||||
for batch_index in range(int(image_item.shape[0])):
|
||||
pil_images.append(_tensor_to_pil(image_item[batch_index]))
|
||||
pil_images.append(tensor_to_pil(image_item[batch_index], rgb=True))
|
||||
else:
|
||||
raise ValueError(f"Unsupported IMAGE tensor shape: {tuple(image_item.shape)}")
|
||||
|
||||
@@ -157,6 +210,7 @@ class NH_ImageGridComposite:
|
||||
|
||||
canvas_width = (columns * cell_width) + ((columns - 1) * gap_x)
|
||||
canvas_height = (rows * cell_height) + ((rows - 1) * gap_y)
|
||||
_check_canvas_size(canvas_width, canvas_height, columns, rows, cell_width, cell_height)
|
||||
fill_rgb = _parse_hex_color(background_color_hex, (0, 0, 0))
|
||||
canvas = Image.new("RGB", (canvas_width, canvas_height), fill_rgb)
|
||||
|
||||
@@ -173,13 +227,13 @@ class NH_ImageGridComposite:
|
||||
x = column * (cell_width + gap_x)
|
||||
y = row * (cell_height + gap_y)
|
||||
source = source_images[image_index]
|
||||
fitted = _fit_image_to_cell(source, cell_width, cell_height, resize_mode, fill_rgb)
|
||||
canvas.paste(fitted, (x, y))
|
||||
# No local name for the fitted cell: it must be freed here, not kept until the tensor conversion.
|
||||
canvas.paste(_fit_image_to_cell(source, cell_width, cell_height, resize_mode, fill_rgb), (x, y))
|
||||
info_lines.append(
|
||||
f"{image_index}: x={x}, y={y}, w={cell_width}, h={cell_height}, source={source.width}x{source.height}"
|
||||
)
|
||||
|
||||
return (_pil_to_tensor(canvas), slot_count, used_count, "\n".join(info_lines))
|
||||
return (pil_to_tensor(canvas), slot_count, used_count, "\n".join(info_lines))
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
|
||||
@@ -0,0 +1,214 @@
|
||||
"""Image Label (NH): a text band above, below or over an image."""
|
||||
|
||||
import functools
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
|
||||
from ...core.utils import tensor_to_pil
|
||||
from ...core.image_ops import RESAMPLING, parse_color, pil_to_image_tensor
|
||||
|
||||
|
||||
_FONT_EXTENSIONS = {".ttf", ".otf", ".ttc"}
|
||||
|
||||
|
||||
def _font_choices():
|
||||
fonts = ["default"]
|
||||
seen = set(fonts)
|
||||
font_dir = Path(os.environ.get("WINDIR", "C:\\Windows")) / "Fonts"
|
||||
if font_dir.exists():
|
||||
for path in sorted(font_dir.iterdir(), key=lambda item: item.name.lower()):
|
||||
if path.is_file() and path.suffix.lower() in _FONT_EXTENSIONS and path.name not in seen:
|
||||
fonts.append(path.name)
|
||||
seen.add(path.name)
|
||||
return fonts
|
||||
|
||||
|
||||
_FONT_CHOICES = _font_choices()
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=64)
|
||||
def _cached_truetype(path, size):
|
||||
"""Open a TrueType font once per (path, size), as the Label node asks for many sizes. None if it cannot be opened."""
|
||||
try:
|
||||
return ImageFont.truetype(path, size)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=1)
|
||||
def _default_font():
|
||||
return ImageFont.load_default()
|
||||
|
||||
|
||||
def _load_font(font_name, font_size):
|
||||
if font_name and font_name != "default":
|
||||
font_path = Path(os.environ.get("WINDIR", "C:\\Windows")) / "Fonts" / font_name
|
||||
if font_path.exists():
|
||||
font = _cached_truetype(str(font_path), font_size)
|
||||
if font is not None:
|
||||
return font
|
||||
font = _cached_truetype("arial.ttf", font_size)
|
||||
return font if font is not None else _default_font()
|
||||
|
||||
|
||||
_MIN_FIT_FONT_SIZE = 8
|
||||
|
||||
|
||||
def _fit_font_size(font_size, fits):
|
||||
"""Largest size in [8, font_size] for which ``fits(size)`` is true; 8 when none does.
|
||||
|
||||
A label gets wider as the font grows, so a binary search lands on the same size as shrinking one point at a
|
||||
time, after about log2(font_size) measurements instead of up to font_size. Hinting can make a text a pixel or
|
||||
two narrower at a few isolated sizes; then the size returned still fits but may sit a point or two below what
|
||||
the one-point-at-a-time loop would have found.
|
||||
"""
|
||||
if font_size <= _MIN_FIT_FONT_SIZE or fits(font_size):
|
||||
return font_size
|
||||
low, high = _MIN_FIT_FONT_SIZE, font_size # `low` fits (or is the floor), `high` does not
|
||||
while high - low > 1:
|
||||
middle = (low + high) // 2
|
||||
if fits(middle):
|
||||
low = middle
|
||||
else:
|
||||
high = middle
|
||||
return low
|
||||
|
||||
|
||||
def _label_text_width(draw, text, font_name, font_size, stroke_width):
|
||||
left, _top, right, _bottom = draw.multiline_textbbox(
|
||||
(0, 0), text, font=_load_font(font_name, font_size), stroke_width=stroke_width
|
||||
)
|
||||
return right - left
|
||||
|
||||
|
||||
class NH_ImageLabel:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE",),
|
||||
"label_text": ("STRING", {"default": "NH Label", "multiline": True}),
|
||||
"position": (["header", "bottom"],),
|
||||
"frame_height": ("INT", {"default": 80, "min": 1, "max": 4096, "step": 1}),
|
||||
"fill_enabled": ("BOOLEAN", {"default": True}),
|
||||
"frame_color_hex": ("STRING", {"default": "#111111", "multiline": False}),
|
||||
"text_color_hex": ("STRING", {"default": "#FFFFFF", "multiline": False}),
|
||||
"font_name": (_FONT_CHOICES,),
|
||||
"font_size": ("INT", {"default": 32, "min": 6, "max": 512, "step": 1}),
|
||||
"fit_text": ("BOOLEAN", {"default": True}),
|
||||
"overlay_on_image": ("BOOLEAN", {"default": False}),
|
||||
"align": (["left", "center", "right"],),
|
||||
"padding_x": ("INT", {"default": 20, "min": 0, "max": 2048, "step": 1}),
|
||||
"outline_width": ("INT", {"default": 0, "min": 0, "max": 32, "step": 1}),
|
||||
"outline_color_hex": ("STRING", {"default": "#000000", "multiline": False}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("image",)
|
||||
FUNCTION = "add_label"
|
||||
CATEGORY = "NH-Nodes/Image"
|
||||
|
||||
def _band_image(self, image, position, frame_height, fill_enabled, frame_color):
|
||||
width, height = image.size
|
||||
if fill_enabled:
|
||||
return Image.new("RGB", (width, frame_height), frame_color)
|
||||
|
||||
sample_box = (0, 0, width, 1) if position == "header" else (0, max(0, height - 1), width, height)
|
||||
return image.crop(sample_box).resize((width, frame_height), RESAMPLING.BILINEAR).convert("RGB")
|
||||
|
||||
def add_label(
|
||||
self,
|
||||
image,
|
||||
label_text,
|
||||
position,
|
||||
frame_height,
|
||||
fill_enabled,
|
||||
frame_color_hex,
|
||||
text_color_hex,
|
||||
font_name,
|
||||
font_size,
|
||||
fit_text,
|
||||
overlay_on_image,
|
||||
align,
|
||||
padding_x,
|
||||
outline_width,
|
||||
outline_color_hex,
|
||||
):
|
||||
frame_color = parse_color(frame_color_hex, (17, 17, 17))
|
||||
text_color = parse_color(text_color_hex, (255, 255, 255))
|
||||
outline_color = parse_color(outline_color_hex, (0, 0, 0))
|
||||
|
||||
labeled = []
|
||||
for item in image:
|
||||
pil_image = tensor_to_pil(item, rgb=True)
|
||||
width, height = pil_image.size
|
||||
band = self._band_image(pil_image, position, frame_height, fill_enabled, frame_color)
|
||||
|
||||
if overlay_on_image:
|
||||
canvas = pil_image.copy()
|
||||
max_band_height = min(frame_height, height)
|
||||
if band.height != max_band_height:
|
||||
band = band.resize((width, max_band_height), RESAMPLING.BILINEAR)
|
||||
if position == "header":
|
||||
band_top = 0
|
||||
else:
|
||||
band_top = height - max_band_height
|
||||
canvas.paste(band, (0, band_top))
|
||||
else:
|
||||
canvas = Image.new("RGB", (width, height + frame_height))
|
||||
if position == "header":
|
||||
canvas.paste(band, (0, 0))
|
||||
canvas.paste(pil_image, (0, frame_height))
|
||||
band_top = 0
|
||||
else:
|
||||
canvas.paste(pil_image, (0, 0))
|
||||
canvas.paste(band, (0, height))
|
||||
band_top = height
|
||||
|
||||
draw = ImageDraw.Draw(canvas)
|
||||
current_size = font_size
|
||||
|
||||
if fit_text:
|
||||
max_width = max(1, width - (padding_x * 2))
|
||||
current_size = _fit_font_size(
|
||||
font_size,
|
||||
lambda size: _label_text_width(draw, label_text, font_name, size, outline_width) <= max_width,
|
||||
)
|
||||
|
||||
font = _load_font(font_name, current_size)
|
||||
bbox = draw.multiline_textbbox((0, 0), label_text, font=font, stroke_width=outline_width)
|
||||
text_width = bbox[2] - bbox[0]
|
||||
text_height = bbox[3] - bbox[1]
|
||||
if align == "left":
|
||||
text_x = padding_x
|
||||
elif align == "right":
|
||||
text_x = max(padding_x, width - padding_x - text_width)
|
||||
else:
|
||||
text_x = max(0, (width - text_width) // 2)
|
||||
|
||||
text_y = band_top + max(0, (frame_height - text_height) // 2) - bbox[1]
|
||||
draw.multiline_text(
|
||||
(text_x, text_y),
|
||||
label_text,
|
||||
font=font,
|
||||
fill=text_color,
|
||||
align=align,
|
||||
stroke_width=outline_width,
|
||||
stroke_fill=outline_color,
|
||||
)
|
||||
labeled.append(pil_to_image_tensor(canvas))
|
||||
|
||||
return (torch.cat(labeled, dim=0),)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"NH_ImageLabel": NH_ImageLabel,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"NH_ImageLabel": "Image Label (NH)",
|
||||
}
|
||||
@@ -0,0 +1,123 @@
|
||||
"""Image Resize Unit (NH): resize to a size given in pixels, millimetres or centimetres at a DPI."""
|
||||
|
||||
import torch
|
||||
|
||||
import comfy.utils
|
||||
|
||||
from ...core.image_ops import parse_color
|
||||
|
||||
|
||||
def _image_resize_methods():
|
||||
return ["nearest", "bilinear", "bicubic", "area", "nearest-exact", "lanczos"]
|
||||
|
||||
|
||||
def _resize_value_to_pixels(value, unit, dpi):
|
||||
if unit == "pixel":
|
||||
return max(0, int(round(value)))
|
||||
divisor = 25.4 if unit == "mm" else 2.54
|
||||
return max(0, int(round((value / divisor) * dpi)))
|
||||
|
||||
|
||||
def _resolve_target_size(orig_width, orig_height, width_value, height_value, unit, dpi):
|
||||
width_px = _resize_value_to_pixels(width_value, unit, dpi)
|
||||
height_px = _resize_value_to_pixels(height_value, unit, dpi)
|
||||
|
||||
if width_px == 0 and height_px == 0:
|
||||
return orig_width, orig_height
|
||||
if width_px == 0:
|
||||
width_px = max(1, round(orig_width * height_px / orig_height))
|
||||
elif height_px == 0:
|
||||
height_px = max(1, round(orig_height * width_px / orig_width))
|
||||
return width_px, height_px
|
||||
|
||||
|
||||
def _common_upscale_method(resize_method):
|
||||
return "nearest-exact" if resize_method == "nearest" else resize_method
|
||||
|
||||
|
||||
def _pad_samples(samples, target_width, target_height, fill_rgb):
|
||||
"""Centre ``samples`` ([B, C, H, W]) on a canvas of the padding colour.
|
||||
|
||||
``fill_rgb`` is (r, g, b) or (r, g, b, a) in 0-255. The fourth channel of RGBA images takes the alpha value, which
|
||||
is opaque when the colour has none.
|
||||
"""
|
||||
batch, channels, resized_h, resized_w = samples.shape
|
||||
canvas = torch.empty((batch, channels, target_height, target_width), dtype=samples.dtype, device=samples.device)
|
||||
fill_values = [channel / 255.0 for channel in fill_rgb[:3]] + [fill_rgb[3] / 255.0 if len(fill_rgb) > 3 else 1.0]
|
||||
for channel in range(channels):
|
||||
canvas[:, channel, :, :] = fill_values[min(channel, 3)]
|
||||
|
||||
offset_x = max(0, (target_width - resized_w) // 2)
|
||||
offset_y = max(0, (target_height - resized_h) // 2)
|
||||
canvas[:, :, offset_y:offset_y + resized_h, offset_x:offset_x + resized_w] = samples
|
||||
return canvas
|
||||
|
||||
|
||||
class NH_ImageResizeByUnit:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE",),
|
||||
"unit": (["pixel", "mm", "cm"],),
|
||||
"width_value": ("FLOAT", {"default": 1024.0, "min": 0.0, "max": 100000.0, "step": 0.1}),
|
||||
"height_value": ("FLOAT", {"default": 1024.0, "min": 0.0, "max": 100000.0, "step": 0.1}),
|
||||
"dpi": ("INT", {"default": 300, "min": 1, "max": 2400, "step": 1}),
|
||||
"resize_mode": (["stretch", "pad", "crop", "lock_ratio"],),
|
||||
"resize_method": (_image_resize_methods(),),
|
||||
"pad_color_hex": ("STRING", {"default": "#000000", "multiline": False}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "INT", "INT")
|
||||
RETURN_NAMES = ("image", "width_px", "height_px")
|
||||
FUNCTION = "resize"
|
||||
CATEGORY = "NH-Nodes/Image"
|
||||
|
||||
def resize(self, image, unit, width_value, height_value, dpi, resize_mode, resize_method, pad_color_hex):
|
||||
samples = image.movedim(-1, 1)
|
||||
orig_height = samples.shape[-2]
|
||||
orig_width = samples.shape[-1]
|
||||
|
||||
target_width, target_height = _resolve_target_size(
|
||||
orig_width, orig_height, width_value, height_value, unit, dpi
|
||||
)
|
||||
upscale_method = _common_upscale_method(resize_method)
|
||||
|
||||
if resize_mode == "lock_ratio":
|
||||
if width_value <= 0 and height_value <= 0:
|
||||
result = samples
|
||||
actual_width, actual_height = orig_width, orig_height
|
||||
elif width_value <= 0 or height_value <= 0:
|
||||
actual_width, actual_height = target_width, target_height
|
||||
result = comfy.utils.common_upscale(samples, actual_width, actual_height, upscale_method, "disabled")
|
||||
else:
|
||||
scale = min(target_width / orig_width, target_height / orig_height)
|
||||
actual_width = max(1, round(orig_width * scale))
|
||||
actual_height = max(1, round(orig_height * scale))
|
||||
result = comfy.utils.common_upscale(samples, actual_width, actual_height, upscale_method, "disabled")
|
||||
elif resize_mode == "crop":
|
||||
actual_width, actual_height = target_width, target_height
|
||||
result = comfy.utils.common_upscale(samples, actual_width, actual_height, upscale_method, "center")
|
||||
elif resize_mode == "pad":
|
||||
actual_width, actual_height = target_width, target_height
|
||||
scale = min(target_width / orig_width, target_height / orig_height)
|
||||
scaled_width = max(1, round(orig_width * scale))
|
||||
scaled_height = max(1, round(orig_height * scale))
|
||||
resized = comfy.utils.common_upscale(samples, scaled_width, scaled_height, upscale_method, "disabled")
|
||||
pad_color = parse_color(pad_color_hex, (0, 0, 0))
|
||||
result = _pad_samples(resized, target_width, target_height, pad_color)
|
||||
else:
|
||||
actual_width, actual_height = target_width, target_height
|
||||
result = comfy.utils.common_upscale(samples, actual_width, actual_height, upscale_method, "disabled")
|
||||
|
||||
return (result.movedim(1, -1), actual_width, actual_height)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"NH_ImageResizeByUnit": NH_ImageResizeByUnit,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"NH_ImageResizeByUnit": "Image Resize Unit (NH)",
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -4,10 +4,24 @@ import json
|
||||
import os
|
||||
|
||||
import folder_paths
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from comfy.cli_args import args
|
||||
from PIL.PngImagePlugin import PngInfo
|
||||
|
||||
from ...core.utils import tensor_to_pil
|
||||
|
||||
# Cạnh dài nhất (px) mà từng định dạng ghi được; PNG/TIFF không bị giới hạn trong thực tế.
|
||||
# JPEG: spec cho tới 65535 nhưng libjpeg (Pillow) từ chối > 65500 (JPEG_MAX_DIMENSION).
|
||||
_MAX_SIDE_BY_EXTENSION = {"webp": 16383, "jpg": 65500}
|
||||
|
||||
|
||||
def _check_size_limit(extension, width, height):
|
||||
limit = _MAX_SIDE_BY_EXTENSION.get(extension)
|
||||
if limit is not None and max(width, height) > limit:
|
||||
raise ValueError(
|
||||
f"NH Save Large Image: {width}x{height} px is too large for {extension.upper()} "
|
||||
f"(max {limit} px per side). Use png or tiff for images this large."
|
||||
)
|
||||
|
||||
|
||||
def _extension_for_save_format(file_format):
|
||||
file_format = str(file_format or "png").lower()
|
||||
@@ -20,14 +34,6 @@ def _extension_for_save_format(file_format):
|
||||
return "png"
|
||||
|
||||
|
||||
def _tensor_to_pil(image_tensor):
|
||||
array = image_tensor.detach().cpu().numpy()
|
||||
array = np.clip(array * 255.0, 0, 255).astype(np.uint8)
|
||||
if array.ndim == 3 and array.shape[-1] == 4:
|
||||
return Image.fromarray(array).convert("RGBA")
|
||||
return Image.fromarray(array).convert("RGB")
|
||||
|
||||
|
||||
def _metadata_pnginfo(prompt=None, extra_pnginfo=None):
|
||||
metadata = PngInfo()
|
||||
if prompt is not None:
|
||||
@@ -85,6 +91,10 @@ class NH_SaveLargeImage:
|
||||
|
||||
def save(self, image, filename_prefix="NH_Large", file_format="png", quality=95, embed_metadata=True, prompt=None, extra_pnginfo=None):
|
||||
extension = _extension_for_save_format(file_format)
|
||||
# Kiểm tra giới hạn trước khi ghi bất kỳ file nào (mọi ảnh trong batch cùng kích thước), kể cả tạo thư mục.
|
||||
_check_size_limit(extension, int(image.shape[2]), int(image.shape[1]))
|
||||
# --disable-metadata của ComfyUI thắng ô embed_metadata: không nhúng prompt/workflow.
|
||||
embed_metadata = bool(embed_metadata) and not args.disable_metadata
|
||||
first = image[0]
|
||||
full_output_folder, filename, counter, _subfolder, _filename_prefix = folder_paths.get_save_image_path(
|
||||
filename_prefix,
|
||||
@@ -98,8 +108,8 @@ class NH_SaveLargeImage:
|
||||
filename_with_batch_num = filename.replace("%batch_num%", str(batch_number))
|
||||
file_name = f"{filename_with_batch_num}_{counter:05}_.{extension}"
|
||||
file_path = os.path.join(full_output_folder, file_name)
|
||||
pil_image = _tensor_to_pil(image_tensor)
|
||||
_save_image(pil_image, file_path, extension, quality, bool(embed_metadata), prompt, extra_pnginfo)
|
||||
pil_image = tensor_to_pil(image_tensor, rgb=True, keep_alpha=True)
|
||||
_save_image(pil_image, file_path, extension, quality, embed_metadata, prompt, extra_pnginfo)
|
||||
saved_paths.append(file_path)
|
||||
counter += 1
|
||||
|
||||
|
||||
@@ -24,6 +24,12 @@ class NH_LargePreviewCacheManager:
|
||||
FUNCTION = "run"
|
||||
CATEGORY = "NH-Nodes/Image"
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, **kwargs):
|
||||
# Info/cleanup đọc và xoá file trên đĩa, input không đổi vẫn phải chạy lại mỗi lần queue;
|
||||
# NaN không bằng chính nó nên ComfyUI không bao giờ dùng lại output cache.
|
||||
return float("nan")
|
||||
|
||||
def run(self, action="info", max_age_hours=24.0, max_total_mb=2048.0, cache_id=""):
|
||||
action = str(action or "info")
|
||||
before = cache_info()
|
||||
|
||||
@@ -0,0 +1,426 @@
|
||||
"""Layer Layout Composite (NH) and Layer Stack Composite (NH): paste images onto a canvas."""
|
||||
|
||||
import re
|
||||
|
||||
from PIL import Image
|
||||
|
||||
from ...core.image_ops import parse_color, pil_to_image_tensor, resize_pil
|
||||
from ...core.safe_eval import evaluate
|
||||
from ...core.utils import tensor_to_pil
|
||||
|
||||
|
||||
_MAX_LAYER_STACK_INPUTS = 64
|
||||
|
||||
|
||||
def _resize_cover(image, target_size):
|
||||
target_width, target_height = target_size
|
||||
width, height = image.size
|
||||
scale = max(target_width / max(width, 1), target_height / max(height, 1))
|
||||
resized = resize_pil(image, (max(1, round(width * scale)), max(1, round(height * scale))), "lanczos")
|
||||
left = max(0, (resized.width - target_width) // 2)
|
||||
top = max(0, (resized.height - target_height) // 2)
|
||||
return resized.crop((left, top, left + target_width, top + target_height))
|
||||
|
||||
|
||||
def _resize_contain(image, target_size):
|
||||
target_width, target_height = target_size
|
||||
width, height = image.size
|
||||
scale = min(target_width / max(width, 1), target_height / max(height, 1))
|
||||
return resize_pil(image, (max(1, round(width * scale)), max(1, round(height * scale))), "lanczos")
|
||||
|
||||
|
||||
def _paste_layer(canvas, layer, position):
|
||||
"""Paste ``layer`` onto ``canvas`` through its own alpha, so transparent areas leave the canvas untouched."""
|
||||
if layer.mode != "RGBA":
|
||||
canvas.paste(layer, position) # opaque, whatever the canvas mode
|
||||
elif canvas.mode == "RGBA":
|
||||
# Both have alpha: a straight paste would blend the alpha channels as if they were colours.
|
||||
box = (position[0], position[1], position[0] + layer.width, position[1] + layer.height)
|
||||
canvas.paste(Image.alpha_composite(canvas.crop(box), layer), position)
|
||||
else:
|
||||
canvas.paste(layer, position, layer)
|
||||
|
||||
|
||||
def _fit_layer_image(image, target_size, resize_mode):
|
||||
target_width, target_height = target_size
|
||||
if target_width <= 0 or target_height <= 0:
|
||||
raise ValueError("Layer slot width and height must be greater than 0")
|
||||
if resize_mode == "stretch":
|
||||
return resize_pil(image, (target_width, target_height), "lanczos"), 0, 0
|
||||
if resize_mode == "contain":
|
||||
fitted = _resize_contain(image, (target_width, target_height))
|
||||
return fitted, (target_width - fitted.width) // 2, (target_height - fitted.height) // 2
|
||||
fitted = _resize_cover(image, (target_width, target_height))
|
||||
return fitted, 0, 0
|
||||
|
||||
|
||||
def _grid_slots(canvas_width, canvas_height, columns, rows, margin_x, margin_y, gap_x, gap_y):
|
||||
columns = max(1, int(columns))
|
||||
rows = max(1, int(rows))
|
||||
content_width = max(1, canvas_width - (2 * margin_x) - (gap_x * (columns - 1)))
|
||||
content_height = max(1, canvas_height - (2 * margin_y) - (gap_y * (rows - 1)))
|
||||
cell_width = max(1, content_width // columns)
|
||||
cell_height = max(1, content_height // rows)
|
||||
slots = []
|
||||
right_edge = canvas_width - margin_x
|
||||
bottom_edge = canvas_height - margin_y
|
||||
|
||||
for row in range(rows):
|
||||
for column in range(columns):
|
||||
x = margin_x + column * (cell_width + gap_x)
|
||||
y = margin_y + row * (cell_height + gap_y)
|
||||
width = max(1, (right_edge - x) if column == columns - 1 else cell_width)
|
||||
height = max(1, (bottom_edge - y) if row == rows - 1 else cell_height)
|
||||
slots.append((x, y, width, height))
|
||||
return slots
|
||||
|
||||
|
||||
def _relative_slot(canvas_width, canvas_height, x, y, width, height):
|
||||
return (
|
||||
int(round(canvas_width * x)),
|
||||
int(round(canvas_height * y)),
|
||||
int(round(canvas_width * width)),
|
||||
int(round(canvas_height * height)),
|
||||
)
|
||||
|
||||
|
||||
def _preset_layout_slots(preset, canvas_width, canvas_height, margin_x, margin_y, gap_x, gap_y):
|
||||
content_x = margin_x
|
||||
content_y = margin_y
|
||||
content_width = max(1, canvas_width - (2 * margin_x))
|
||||
content_height = max(1, canvas_height - (2 * margin_y))
|
||||
|
||||
if preset == "1_full":
|
||||
return [(content_x, content_y, content_width, content_height)]
|
||||
if preset == "2_vertical":
|
||||
return _grid_slots(canvas_width, canvas_height, 1, 2, margin_x, margin_y, gap_x, gap_y)
|
||||
if preset == "2_horizontal":
|
||||
return _grid_slots(canvas_width, canvas_height, 2, 1, margin_x, margin_y, gap_x, gap_y)
|
||||
if preset == "3_vertical":
|
||||
return _grid_slots(canvas_width, canvas_height, 1, 3, margin_x, margin_y, gap_x, gap_y)
|
||||
if preset == "4_grid":
|
||||
return _grid_slots(canvas_width, canvas_height, 2, 2, margin_x, margin_y, gap_x, gap_y)
|
||||
if preset == "2_top_1_bottom":
|
||||
top_height = max(1, round((content_height - gap_y) * 0.42))
|
||||
bottom_y = content_y + top_height + gap_y
|
||||
bottom_height = max(1, content_y + content_height - bottom_y)
|
||||
half_width = max(1, (content_width - gap_x) // 2)
|
||||
return [
|
||||
(content_x, content_y, half_width, top_height),
|
||||
(content_x + half_width + gap_x, content_y, max(1, content_width - half_width - gap_x), top_height),
|
||||
(content_x, bottom_y, content_width, bottom_height),
|
||||
]
|
||||
if preset == "2x3_mixed":
|
||||
row_heights = [0.20, 0.20, 0.44]
|
||||
usable_height = max(1, content_height - (2 * gap_y))
|
||||
heights = [max(1, round(usable_height * ratio / sum(row_heights))) for ratio in row_heights]
|
||||
heights[-1] = max(1, content_height - (2 * gap_y) - heights[0] - heights[1])
|
||||
half_width = max(1, (content_width - gap_x) // 2)
|
||||
slots = []
|
||||
y = content_y
|
||||
for height in heights:
|
||||
slots.append((content_x, y, half_width, height))
|
||||
slots.append((content_x + half_width + gap_x, y, max(1, content_width - half_width - gap_x), height))
|
||||
y += height + gap_y
|
||||
return slots
|
||||
return _grid_slots(canvas_width, canvas_height, 2, 3, margin_x, margin_y, gap_x, gap_y)
|
||||
|
||||
|
||||
def _parse_custom_layout(custom_layout, canvas_width, canvas_height, custom_units):
|
||||
slots = []
|
||||
for raw_line in str(custom_layout or "").splitlines():
|
||||
line = raw_line.split("#", 1)[0].strip()
|
||||
if not line:
|
||||
continue
|
||||
values = [float(value) for value in re.split(r"[\s,;]+", line) if value.strip()]
|
||||
if len(values) < 4:
|
||||
raise ValueError(f"Invalid custom layout line: {raw_line}")
|
||||
if len(values) >= 5:
|
||||
values = values[-4:]
|
||||
x, y, width, height = values[:4]
|
||||
|
||||
units = custom_units
|
||||
if units == "auto":
|
||||
units = "relative" if max(abs(x), abs(y), abs(width), abs(height)) <= 1.0 else "pixels"
|
||||
if units == "relative":
|
||||
slot = _relative_slot(canvas_width, canvas_height, x, y, width, height)
|
||||
elif units == "percent":
|
||||
slot = _relative_slot(canvas_width, canvas_height, x / 100.0, y / 100.0, width / 100.0, height / 100.0)
|
||||
else:
|
||||
slot = (int(round(x)), int(round(y)), int(round(width)), int(round(height)))
|
||||
|
||||
if slot[2] <= 0 or slot[3] <= 0:
|
||||
raise ValueError(f"Custom layout slot size must be greater than 0: {raw_line}")
|
||||
slots.append(slot)
|
||||
return slots
|
||||
|
||||
|
||||
def _layout_slots(layout_mode, layout_preset, canvas_width, canvas_height, columns, rows, margin_x, margin_y, gap_x, gap_y, custom_layout, custom_units):
|
||||
if layout_mode == "custom":
|
||||
slots = _parse_custom_layout(custom_layout, canvas_width, canvas_height, custom_units)
|
||||
if not slots:
|
||||
raise ValueError("Custom layout is empty")
|
||||
return slots
|
||||
if layout_mode == "grid":
|
||||
return _grid_slots(canvas_width, canvas_height, columns, rows, margin_x, margin_y, gap_x, gap_y)
|
||||
return _preset_layout_slots(layout_preset, canvas_width, canvas_height, margin_x, margin_y, gap_x, gap_y)
|
||||
|
||||
|
||||
def _eval_position_expr(expression, variables):
|
||||
# Variable keys are already casefolded. core.safe_eval works in floating point, so 9**9**9 overflows into an
|
||||
# error instead of building a huge integer and freezing the worker.
|
||||
try:
|
||||
return evaluate(expression, variables, zero_division="error", case_insensitive=True)
|
||||
except ValueError as exc: # core.safe_eval.ExpressionError is a ValueError
|
||||
raise ValueError(f"Invalid position expression '{expression}': {exc}") from exc
|
||||
|
||||
|
||||
def _split_position_line(line):
|
||||
parts = [part.strip() for part in re.split(r"[,;]+", line, maxsplit=1)]
|
||||
if len(parts) < 2 or not parts[0] or not parts[1]:
|
||||
raise ValueError(f"Invalid position line: {line}")
|
||||
return parts[0], parts[1]
|
||||
|
||||
|
||||
def _parse_layer_positions(positions, count, background_width=0, background_height=0, layer_sizes=None):
|
||||
variables = {
|
||||
"a": int(background_width),
|
||||
"b": int(background_height),
|
||||
"bg_w": int(background_width),
|
||||
"bg_h": int(background_height),
|
||||
"width": int(background_width),
|
||||
"height": int(background_height),
|
||||
}
|
||||
for index, (layer_width, layer_height) in enumerate(layer_sizes or [], start=1):
|
||||
variables[f"w{index}"] = int(layer_width)
|
||||
variables[f"h{index}"] = int(layer_height)
|
||||
|
||||
parsed = []
|
||||
for raw_line in str(positions or "").splitlines():
|
||||
line = raw_line.split("#", 1)[0].strip()
|
||||
if not line:
|
||||
continue
|
||||
x_expr, y_expr = _split_position_line(line)
|
||||
x = int(round(_eval_position_expr(x_expr, variables)))
|
||||
y = int(round(_eval_position_expr(y_expr, variables)))
|
||||
parsed.append((x, y))
|
||||
|
||||
index = len(parsed)
|
||||
variables[f"x{index}"] = x
|
||||
variables[f"y{index}"] = y
|
||||
|
||||
while len(parsed) < count:
|
||||
parsed.append((0, 0))
|
||||
return parsed[:count]
|
||||
|
||||
|
||||
class NH_LayerLayoutComposite:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE",),
|
||||
"canvas_width": ("INT", {"default": 2480, "min": 1, "max": 20000, "step": 1}),
|
||||
"canvas_height": ("INT", {"default": 3508, "min": 1, "max": 20000, "step": 1}),
|
||||
"background_color": ("STRING", {"default": "#FFFFFF", "multiline": False}),
|
||||
"layout_mode": (["preset", "grid", "custom"],),
|
||||
"layout_preset": (["2x3", "2x3_mixed", "2_top_1_bottom", "1_full", "2_vertical", "2_horizontal", "3_vertical", "4_grid"],),
|
||||
"resize_mode": (["cover", "contain", "stretch"],),
|
||||
"columns": ("INT", {"default": 2, "min": 1, "max": 20, "step": 1}),
|
||||
"rows": ("INT", {"default": 3, "min": 1, "max": 20, "step": 1}),
|
||||
"margin_x": ("INT", {"default": 150, "min": 0, "max": 10000, "step": 1}),
|
||||
"margin_y": ("INT", {"default": 150, "min": 0, "max": 10000, "step": 1}),
|
||||
"gap_x": ("INT", {"default": 150, "min": 0, "max": 10000, "step": 1}),
|
||||
"gap_y": ("INT", {"default": 150, "min": 0, "max": 10000, "step": 1}),
|
||||
"custom_units": (["auto", "relative", "percent", "pixels"],),
|
||||
"custom_layout": ("STRING", {
|
||||
"default": "0.06,0.04,0.39,0.20\n0.55,0.04,0.39,0.20\n0.06,0.28,0.39,0.20\n0.55,0.28,0.39,0.20\n0.06,0.58,0.39,0.36\n0.55,0.58,0.39,0.36",
|
||||
"multiline": True,
|
||||
}),
|
||||
"background_fit": (["stretch", "cover", "contain"],),
|
||||
},
|
||||
"optional": {
|
||||
"background_image": ("IMAGE",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "INT", "INT", "STRING")
|
||||
RETURN_NAMES = ("image", "slot_count", "used_count", "layout_info")
|
||||
FUNCTION = "composite"
|
||||
CATEGORY = "NH-Nodes/Image"
|
||||
|
||||
@staticmethod
|
||||
def _paste_clipped(canvas, layer, x, y):
|
||||
left = max(0, int(x))
|
||||
top = max(0, int(y))
|
||||
right = min(canvas.width, int(x) + layer.width)
|
||||
bottom = min(canvas.height, int(y) + layer.height)
|
||||
if right <= left or bottom <= top:
|
||||
return
|
||||
crop_box = (left - int(x), top - int(y), right - int(x), bottom - int(y))
|
||||
_paste_layer(canvas, layer.crop(crop_box), (left, top))
|
||||
|
||||
def _create_canvas(self, canvas_width, canvas_height, background_color, background_fit, background_image=None):
|
||||
fill = parse_color(background_color, (255, 255, 255))
|
||||
# An opaque colour gives the RGB canvas (and RGB output) this node always had; a colour with transparency
|
||||
# ('#RRGGBBAA' with AA below FF) gives an RGBA canvas so that the transparency survives.
|
||||
canvas = Image.new("RGBA" if fill[3] < 255 else "RGB", (canvas_width, canvas_height), fill)
|
||||
|
||||
if background_image is None:
|
||||
return canvas
|
||||
|
||||
background = tensor_to_pil(background_image[0], rgb=True, keep_alpha=True)
|
||||
fitted, offset_x, offset_y = _fit_layer_image(background, (canvas_width, canvas_height), background_fit)
|
||||
self._paste_clipped(canvas, fitted, offset_x, offset_y)
|
||||
return canvas
|
||||
|
||||
def composite(
|
||||
self,
|
||||
images,
|
||||
canvas_width,
|
||||
canvas_height,
|
||||
background_color,
|
||||
layout_mode,
|
||||
layout_preset,
|
||||
resize_mode,
|
||||
columns,
|
||||
rows,
|
||||
margin_x,
|
||||
margin_y,
|
||||
gap_x,
|
||||
gap_y,
|
||||
custom_units,
|
||||
custom_layout,
|
||||
background_fit,
|
||||
background_image=None,
|
||||
):
|
||||
canvas_width = int(canvas_width)
|
||||
canvas_height = int(canvas_height)
|
||||
slots = _layout_slots(
|
||||
layout_mode,
|
||||
layout_preset,
|
||||
canvas_width,
|
||||
canvas_height,
|
||||
int(columns),
|
||||
int(rows),
|
||||
int(margin_x),
|
||||
int(margin_y),
|
||||
int(gap_x),
|
||||
int(gap_y),
|
||||
custom_layout,
|
||||
custom_units,
|
||||
)
|
||||
|
||||
canvas = self._create_canvas(canvas_width, canvas_height, background_color, background_fit, background_image)
|
||||
used_count = min(int(images.shape[0]), len(slots))
|
||||
info_lines = []
|
||||
|
||||
for image_index in range(used_count):
|
||||
x, y, width, height = slots[image_index]
|
||||
layer = tensor_to_pil(images[image_index], rgb=True, keep_alpha=True)
|
||||
fitted, offset_x, offset_y = _fit_layer_image(layer, (int(width), int(height)), resize_mode)
|
||||
paste_x = int(x) + offset_x
|
||||
paste_y = int(y) + offset_y
|
||||
self._paste_clipped(canvas, fitted, paste_x, paste_y)
|
||||
info_lines.append(
|
||||
f"{image_index}: x={int(x)}, y={int(y)}, w={int(width)}, h={int(height)}, mode={resize_mode}"
|
||||
)
|
||||
|
||||
return (pil_to_image_tensor(canvas), len(slots), used_count, "\n".join(info_lines))
|
||||
|
||||
|
||||
class NH_LayerStackComposite:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
optional_inputs = {
|
||||
f"layer_{index}": ("IMAGE",)
|
||||
for index in range(1, _MAX_LAYER_STACK_INPUTS + 1)
|
||||
}
|
||||
return {
|
||||
"required": {
|
||||
"background_image": ("IMAGE",),
|
||||
"layer_count": ("INT", {"default": 1, "min": 0, "max": _MAX_LAYER_STACK_INPUTS, "step": 1}),
|
||||
"positions": ("STRING", {
|
||||
"default": "0,0",
|
||||
"multiline": True,
|
||||
}),
|
||||
},
|
||||
"optional": optional_inputs,
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "INT", "STRING")
|
||||
RETURN_NAMES = ("image", "used_count", "layer_info")
|
||||
FUNCTION = "composite"
|
||||
CATEGORY = "NH-Nodes/Image"
|
||||
|
||||
@staticmethod
|
||||
def _paste_clipped(canvas, layer, x, y):
|
||||
x = int(x)
|
||||
y = int(y)
|
||||
left = max(0, x)
|
||||
top = max(0, y)
|
||||
right = min(canvas.width, x + layer.width)
|
||||
bottom = min(canvas.height, y + layer.height)
|
||||
if right <= left or bottom <= top:
|
||||
return False
|
||||
|
||||
crop_box = (left - x, top - y, right - x, bottom - y)
|
||||
_paste_layer(canvas, layer.crop(crop_box), (left, top))
|
||||
return True
|
||||
|
||||
def composite(self, background_image, layer_count, positions, **kwargs):
|
||||
# The output has the channels of the background: RGB stays RGB, an RGBA background gives RGBA.
|
||||
background = tensor_to_pil(background_image[0], rgb=True, keep_alpha=True)
|
||||
canvas = background.copy()
|
||||
requested_count = max(0, min(int(layer_count), _MAX_LAYER_STACK_INPUTS))
|
||||
layer_images = {}
|
||||
layer_sizes = []
|
||||
for index in range(1, requested_count + 1):
|
||||
layer = kwargs.get(f"layer_{index}")
|
||||
if layer is None:
|
||||
layer_sizes.append((0, 0))
|
||||
continue
|
||||
layer_image = tensor_to_pil(layer[0], rgb=True, keep_alpha=True)
|
||||
layer_images[index] = layer_image
|
||||
layer_sizes.append((layer_image.width, layer_image.height))
|
||||
|
||||
layer_positions = _parse_layer_positions(
|
||||
positions,
|
||||
requested_count,
|
||||
background.width,
|
||||
background.height,
|
||||
layer_sizes=layer_sizes,
|
||||
)
|
||||
used_count = 0
|
||||
info_lines = []
|
||||
|
||||
for index in range(1, requested_count + 1):
|
||||
layer_image = layer_images.get(index)
|
||||
if layer_image is None:
|
||||
info_lines.append(f"layer_{index}: missing")
|
||||
continue
|
||||
|
||||
x, y = layer_positions[index - 1]
|
||||
pasted = self._paste_clipped(canvas, layer_image, x, y)
|
||||
if pasted:
|
||||
used_count += 1
|
||||
info_lines.append(
|
||||
f"layer_{index}: x={x}, y={y}, w={layer_image.width}, h={layer_image.height}"
|
||||
)
|
||||
else:
|
||||
info_lines.append(
|
||||
f"layer_{index}: outside canvas x={x}, y={y}, w={layer_image.width}, h={layer_image.height}"
|
||||
)
|
||||
|
||||
return (pil_to_image_tensor(canvas), used_count, "\n".join(info_lines))
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"NH_LayerLayoutComposite": NH_LayerLayoutComposite,
|
||||
"NH_LayerStackComposite": NH_LayerStackComposite,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"NH_LayerLayoutComposite": "Layer Layout Composite (NH)",
|
||||
"NH_LayerStackComposite": "Layer Stack Composite (NH)",
|
||||
}
|
||||
@@ -2,12 +2,15 @@ import torch
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
from ...core.mask_ops import mask_bbox
|
||||
|
||||
|
||||
class NH_MaskResizeImage:
|
||||
"""Resize ảnh theo kích thước target, đảm bảo vùng mask không bị mất.
|
||||
|
||||
- Crop: scale để mask vừa khít target, crop phần ảnh thừa (pad nếu ảnh nhỏ hơn target).
|
||||
- Pad: scale vừa đủ để mask nằm trong target, pad phần thừa bằng màu tùy chọn.
|
||||
- Mask rỗng: coi như mask phủ cả ảnh (crop: lấp đầy target rồi cắt giữa; pad: cả ảnh nằm gọn trong target).
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
@@ -44,7 +47,8 @@ class NH_MaskResizeImage:
|
||||
results_img.append(out_img)
|
||||
results_mask.append(out_mask)
|
||||
|
||||
out_images = torch.from_numpy(np.stack(results_img, axis=0)).float()
|
||||
# Lanczos có thể vọt ra ngoài [0,1] (ringing) quanh cạnh tương phản cao -> kẹp lại
|
||||
out_images = torch.from_numpy(np.clip(np.stack(results_img, axis=0), 0.0, 1.0)).float()
|
||||
out_masks = torch.from_numpy(np.stack(results_mask, axis=0)).float()
|
||||
return (out_images, out_masks)
|
||||
|
||||
@@ -52,13 +56,12 @@ class NH_MaskResizeImage:
|
||||
sh, sw = img_np.shape[:2]
|
||||
|
||||
# Tìm bounding box của mask
|
||||
rows, cols = np.where(mask_np > 0.5)
|
||||
if len(rows) == 0:
|
||||
bbox = mask_bbox(mask_np)
|
||||
if bbox is None:
|
||||
# Mask rỗng: resize bình thường
|
||||
return self._simple_resize(img_np, mask_np, tw, th, mode, pad_color)
|
||||
|
||||
my1, my2 = int(np.min(rows)), int(np.max(rows))
|
||||
mx1, mx2 = int(np.min(cols)), int(np.max(cols))
|
||||
mx1, my1, mx2, my2 = bbox
|
||||
mw = mx2 - mx1 + 1
|
||||
mh = my2 - my1 + 1
|
||||
|
||||
@@ -144,10 +147,23 @@ class NH_MaskResizeImage:
|
||||
return canvas_img, canvas_mask
|
||||
|
||||
def _simple_resize(self, img_np, mask_np, tw, th, mode, pad_color):
|
||||
"""Fallback khi mask rỗng."""
|
||||
img_resized = cv2.resize(img_np, (tw, th), interpolation=cv2.INTER_LANCZOS4)
|
||||
mask_resized = cv2.resize(mask_np, (tw, th), interpolation=cv2.INTER_NEAREST)
|
||||
return img_resized, mask_resized
|
||||
"""Fallback khi mask rỗng: coi như mask phủ cả ảnh, không méo ảnh.
|
||||
|
||||
crop = scale lớn nhất (lấp đầy target) rồi cắt giữa; pad = scale nhỏ nhất (cả ảnh nằm trong target), phần thừa pad_color.
|
||||
"""
|
||||
sh, sw = img_np.shape[:2]
|
||||
scale = max(tw / sw, th / sh) if mode == "crop" else min(tw / sw, th / sh)
|
||||
new_sw = max(1, round(sw * scale))
|
||||
new_sh = max(1, round(sh * scale))
|
||||
img_scaled = cv2.resize(img_np, (new_sw, new_sh), interpolation=cv2.INTER_LANCZOS4)
|
||||
mask_scaled = cv2.resize(mask_np, (new_sw, new_sh), interpolation=cv2.INTER_NEAREST)
|
||||
|
||||
# Tâm ảnh sau khi scale
|
||||
cx = round(sw / 2.0 * scale)
|
||||
cy = round(sh / 2.0 * scale)
|
||||
if mode == "crop":
|
||||
return self._crop_around_center(img_scaled, mask_scaled, tw, th, cx, cy, pad_color)
|
||||
return self._pad_around_center(img_scaled, mask_scaled, tw, th, cx, cy, pad_color)
|
||||
|
||||
@staticmethod
|
||||
def _hex_to_rgb(hex_str):
|
||||
|
||||
@@ -0,0 +1,197 @@
|
||||
"""Save Image Path (NH): numbered image files in any folder, with an optional preview."""
|
||||
|
||||
import os
|
||||
import re
|
||||
import uuid
|
||||
|
||||
import folder_paths
|
||||
|
||||
from ...core.utils import tensor_to_pil
|
||||
|
||||
_WINDOWS = os.name == "nt"
|
||||
|
||||
|
||||
def _save_dir_from_text(folder_path):
|
||||
folder_path = (folder_path or "").strip()
|
||||
if not folder_path:
|
||||
return folder_paths.get_output_directory()
|
||||
if os.path.isfile(folder_path):
|
||||
return os.path.dirname(folder_path)
|
||||
if os.path.isabs(folder_path):
|
||||
return folder_path
|
||||
return os.path.join(folder_paths.get_output_directory(), folder_path)
|
||||
|
||||
|
||||
def _contained_path(base_dir, relative, what):
|
||||
"""Absolute path of ``relative`` below ``base_dir``; ValueError (naming ``what``) if it would leave ``base_dir``.
|
||||
|
||||
Both slash styles count as separators on every OS, so a workflow written on Windows is judged the same elsewhere.
|
||||
"""
|
||||
base_dir = os.path.abspath(base_dir)
|
||||
path = os.path.abspath(os.path.join(base_dir, str(relative).replace("\\", "/")))
|
||||
try:
|
||||
inside = os.path.commonpath([base_dir, path]) == base_dir
|
||||
except ValueError: # another drive or a UNC share
|
||||
inside = False
|
||||
if not inside:
|
||||
raise ValueError(f"{what} must stay inside the target folder: {relative!r}")
|
||||
return path
|
||||
|
||||
|
||||
def _save_child_dir(save_dir, folder_name):
|
||||
folder_name = str(folder_name or "").strip()
|
||||
if not folder_name:
|
||||
return save_dir
|
||||
return _contained_path(save_dir, folder_name, "new_folder_name")
|
||||
|
||||
|
||||
def _split_filename_prefix(save_dir, prefix):
|
||||
"""('sub/dir/', 'name', directory) for a filename_prefix such as 'sub/dir/name'.
|
||||
|
||||
Subfolders are allowed, but the folder part must stay inside ``save_dir`` (no '..', no absolute path) and the
|
||||
file name must be a plain name. Without a separator the name is not looked at by the folder check, and on
|
||||
Windows a colon turns it into a path: os.path.join(save_dir, 'D:x0001.png') is a file on drive D:, 'C:x' is
|
||||
relative to the current folder of drive C: (join drops the 'C:', so the reported name is wrong and the
|
||||
numbering never finds the file again) and 'name:stream' writes an NTFS alternate data stream. A Windows file
|
||||
name cannot contain a colon anyway, so such a prefix is refused whole, before anything is created.
|
||||
"""
|
||||
if _WINDOWS and ":" in prefix:
|
||||
raise ValueError(
|
||||
"filename_prefix cannot contain ':' on Windows (it selects a drive or an NTFS stream, "
|
||||
f"not a file name): {prefix!r}"
|
||||
)
|
||||
folder_part, name = re.match(r"(.*[\\/])?([^\\/]*)\Z", prefix, re.DOTALL).groups()
|
||||
if not folder_part:
|
||||
return "", name, save_dir
|
||||
return folder_part, name, _contained_path(save_dir, folder_part, "filename_prefix")
|
||||
|
||||
|
||||
def _save_pil_image(image, save_path, file_format, dpi):
|
||||
image = image.convert("RGB")
|
||||
save_kwargs = {"dpi": (dpi, dpi)}
|
||||
|
||||
if file_format == "png":
|
||||
pil_format = "PNG"
|
||||
save_kwargs["compress_level"] = 4
|
||||
elif file_format == "jpg":
|
||||
pil_format = "JPEG"
|
||||
save_kwargs["quality"] = 95
|
||||
save_kwargs["subsampling"] = 0
|
||||
elif file_format == "webp":
|
||||
pil_format = "WEBP"
|
||||
save_kwargs["quality"] = 95
|
||||
save_kwargs["method"] = 6
|
||||
elif file_format == "avif":
|
||||
pil_format = "AVIF"
|
||||
save_kwargs["quality"] = 90
|
||||
else:
|
||||
raise ValueError(f"Unsupported format: {file_format}")
|
||||
|
||||
try:
|
||||
image.save(save_path, format=pil_format, **save_kwargs)
|
||||
except TypeError:
|
||||
if file_format == "avif" and "dpi" in save_kwargs:
|
||||
save_kwargs.pop("dpi", None)
|
||||
image.save(save_path, format=pil_format, **save_kwargs)
|
||||
print("[NH-Nodes] Save Image: AVIF DPI metadata is not supported by the current Pillow build.")
|
||||
else:
|
||||
raise
|
||||
|
||||
|
||||
class NH_SaveImagePath:
|
||||
OUTPUT_NODE = True
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE",),
|
||||
"folder_path": ("STRING", {"default": "", "multiline": False}),
|
||||
"create_new_folder": ("BOOLEAN", {"default": False}),
|
||||
"new_folder_name": ("STRING", {"default": "", "multiline": False}),
|
||||
"filename_prefix": ("STRING", {"default": "NH", "multiline": False}),
|
||||
"number_suffix": ("BOOLEAN", {"default": True}),
|
||||
"preview_image": ("BOOLEAN", {"default": False}),
|
||||
"file_format": (["png", "jpg", "webp", "avif"],),
|
||||
"dpi": ("INT", {"default": 300, "min": 1, "max": 2400, "step": 1}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING", "STRING", "INT")
|
||||
RETURN_NAMES = ("filenames", "path", "saved_count")
|
||||
FUNCTION = "save_images"
|
||||
CATEGORY = "NH-Nodes/Image"
|
||||
|
||||
def save_images(self, images, folder_path, filename_prefix, number_suffix=True, preview_image=False, file_format="png", dpi=300, create_new_folder=False, new_folder_name=""):
|
||||
save_dir = _save_dir_from_text(folder_path)
|
||||
if create_new_folder:
|
||||
save_dir = _save_child_dir(save_dir, new_folder_name)
|
||||
|
||||
# A prefix may name subfolders ("renders/img"); they live below save_dir and are where files are counted.
|
||||
prefix = (filename_prefix or "NH").strip() or "NH"
|
||||
prefix_folder, prefix_name, target_dir = _split_filename_prefix(save_dir, prefix)
|
||||
os.makedirs(target_dir, exist_ok=True)
|
||||
|
||||
numbered_pattern = re.compile(rf"^{re.escape(prefix_name)}(\d+)\.(?:png|jpe?g|webp|avif)$", re.IGNORECASE)
|
||||
legacy_numbered_pattern = re.compile(rf"^(?:{re.escape(prefix_name)})?(\d+)\.(?:png|jpe?g|webp|avif)$", re.IGNORECASE)
|
||||
bare_pattern = re.compile(rf"^{re.escape(prefix_name)}\.(?:png|jpe?g|webp|avif)$", re.IGNORECASE)
|
||||
|
||||
used_numbers = set()
|
||||
bare_exists = False
|
||||
for name in os.listdir(target_dir):
|
||||
if bare_pattern.match(name):
|
||||
bare_exists = True
|
||||
continue
|
||||
match = numbered_pattern.match(name)
|
||||
if match is None and number_suffix:
|
||||
match = legacy_numbered_pattern.match(name)
|
||||
if match:
|
||||
used_numbers.add(int(match.group(1)))
|
||||
|
||||
saved_files = []
|
||||
preview_results = []
|
||||
preview_dir = folder_paths.get_temp_directory()
|
||||
next_number = 1
|
||||
for batch_index, image in enumerate(images):
|
||||
if number_suffix:
|
||||
while next_number in used_numbers:
|
||||
next_number += 1
|
||||
filename = f"{prefix_name}{next_number:04d}.{file_format}"
|
||||
used_numbers.add(next_number)
|
||||
next_number += 1
|
||||
elif not bare_exists and not used_numbers:
|
||||
filename = f"{prefix_name}.{file_format}"
|
||||
bare_exists = True
|
||||
else:
|
||||
while next_number in used_numbers:
|
||||
next_number += 1
|
||||
filename = f"{prefix_name}{next_number}.{file_format}"
|
||||
used_numbers.add(next_number)
|
||||
next_number += 1
|
||||
|
||||
pil_image = tensor_to_pil(image, rgb=True)
|
||||
_save_pil_image(pil_image, os.path.join(target_dir, filename), file_format, int(dpi))
|
||||
saved_files.append(prefix_folder + filename)
|
||||
|
||||
if preview_image:
|
||||
preview_filename = f"NH_preview_{uuid.uuid4().hex}_{batch_index}.png"
|
||||
_save_pil_image(pil_image, os.path.join(preview_dir, preview_filename), "png", int(dpi))
|
||||
preview_results.append({
|
||||
"filename": preview_filename,
|
||||
"subfolder": "",
|
||||
"type": "temp",
|
||||
})
|
||||
|
||||
result = ("\n".join(saved_files), save_dir, len(saved_files))
|
||||
if preview_image:
|
||||
return {"ui": {"images": preview_results}, "result": result}
|
||||
return result
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"NH_SaveImagePath": NH_SaveImagePath,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"NH_SaveImagePath": "Save Image Path (NH)",
|
||||
}
|
||||
@@ -1,8 +1,12 @@
|
||||
"""Indexed loader nodes for LoRA and diffusion models."""
|
||||
|
||||
import os
|
||||
import weakref
|
||||
|
||||
import torch
|
||||
|
||||
import folder_paths
|
||||
from comfy.cli_args import args
|
||||
from comfy_execution.graph_utils import ExecutionBlocker
|
||||
|
||||
|
||||
@@ -41,6 +45,62 @@ def _selected_name(kwargs, prefix, count, index):
|
||||
return selected or None
|
||||
|
||||
|
||||
def _file_stamp(path):
|
||||
"""(size, mtime) of a model file, so a file replaced under the same name is not taken for the one already loaded."""
|
||||
try:
|
||||
stat = os.stat(path)
|
||||
except OSError:
|
||||
return None
|
||||
return (stat.st_size, stat.st_mtime_ns)
|
||||
|
||||
|
||||
def _keeps_strong_reference():
|
||||
"""True when ComfyUI runs with --cache-classic (--high-ram implies it)."""
|
||||
return bool(getattr(args, "cache_classic", False))
|
||||
|
||||
|
||||
class _LoadedObjectCache:
|
||||
"""Remembers the last object a loader node built, under a key of everything that decides what gets loaded.
|
||||
|
||||
ComfyUI keeps one instance per graph node, but the node's cache key holds every one of its 64 slot widgets, so
|
||||
editing a slot nobody selected runs the node again. Handing back the object that is already in memory then saves
|
||||
reading it from disk again.
|
||||
|
||||
The object is held by a weak reference. While ComfyUI's own output cache holds it, the next run gets it back; when
|
||||
that cache evicts it (the default RAM-pressure cache does, under memory pressure) nothing here keeps it alive, so
|
||||
it can be freed. A strong reference would defeat that eviction. Only with --cache-classic, which drops an output as
|
||||
soon as the prompt stops using it, the node holds the object itself: otherwise there would be nothing to reuse.
|
||||
An object that cannot be weakly referenced is not remembered at all and is loaded again each time.
|
||||
|
||||
One object only: choosing another file drops the old one before the new one loads, so two big models never sit in
|
||||
memory together, and "Free model and node cache" (which throws the instance away) frees it as well.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self._loaded = None # None, or (key, a callable that returns the object - None once it has been freed)
|
||||
|
||||
def _load_cached(self, key, load):
|
||||
remembered, self._loaded = self._loaded, None # whatever happens below, the node lets go of the old object
|
||||
if remembered is not None and remembered[0] == key:
|
||||
obj = remembered[1]()
|
||||
if obj is not None:
|
||||
self._loaded = remembered
|
||||
return obj
|
||||
del remembered # the old object must be gone before the next one loads
|
||||
obj = load()
|
||||
self._loaded = self._remember(key, obj)
|
||||
return obj
|
||||
|
||||
@staticmethod
|
||||
def _remember(key, obj):
|
||||
if _keeps_strong_reference():
|
||||
return (key, lambda: obj)
|
||||
try:
|
||||
return (key, weakref.ref(obj))
|
||||
except TypeError: # e.g. a __slots__ class without __weakref__: loading it again beats keeping it for good
|
||||
return None
|
||||
|
||||
|
||||
class _IndexedLoraBase:
|
||||
def __init__(self):
|
||||
self.loaded_lora = None
|
||||
@@ -180,7 +240,7 @@ class NH_LoraClipIndex(_IndexedLoraBase):
|
||||
return (model_lora, clip_lora, lora_name, final_index)
|
||||
|
||||
|
||||
class NH_DiffusionModelIndex:
|
||||
class NH_DiffusionModelIndex(_LoadedObjectCache):
|
||||
"""Load one diffusion model by 1-based index."""
|
||||
|
||||
@classmethod
|
||||
@@ -240,11 +300,14 @@ class NH_DiffusionModelIndex:
|
||||
model_options["dtype"] = torch.float8_e5m2
|
||||
|
||||
model_path = folder_paths.get_full_path_or_raise("diffusion_models", model_name)
|
||||
model = comfy.sd.load_diffusion_model(model_path, model_options=model_options)
|
||||
model = self._load_cached(
|
||||
(model_path, weight_dtype, _file_stamp(model_path)),
|
||||
lambda: comfy.sd.load_diffusion_model(model_path, model_options=model_options),
|
||||
)
|
||||
return (model, model_name, final_index)
|
||||
|
||||
|
||||
class NH_VAEIndex:
|
||||
class NH_VAEIndex(_LoadedObjectCache):
|
||||
"""Load one VAE by 1-based index."""
|
||||
|
||||
@classmethod
|
||||
@@ -288,16 +351,26 @@ class NH_VAEIndex:
|
||||
if vae_name is None:
|
||||
return (_block(), "", -1)
|
||||
|
||||
import folder_paths
|
||||
import comfy.sd
|
||||
import comfy.utils
|
||||
|
||||
vae_path = folder_paths.get_full_path_or_raise("vae", vae_name)
|
||||
vae = comfy.sd.VAE(sd=comfy.utils.load_torch_file(vae_path))
|
||||
|
||||
def load():
|
||||
# Like the core VAE loader: the file's metadata goes to the VAE as well (it can hold the VAE's config).
|
||||
sd, metadata = comfy.utils.load_torch_file(vae_path, return_metadata=True)
|
||||
vae = comfy.sd.VAE(sd=sd, metadata=metadata)
|
||||
# The core VAE loader does this too: a file with no usable VAE fails here, not in the sampler's decode.
|
||||
check = getattr(vae, "throw_exception_if_invalid", None)
|
||||
if check is not None:
|
||||
check()
|
||||
return vae
|
||||
|
||||
vae = self._load_cached((vae_path, _file_stamp(vae_path)), load)
|
||||
return (vae, vae_name, final_index)
|
||||
|
||||
|
||||
class NH_ClipIndex:
|
||||
class NH_ClipIndex(_LoadedObjectCache):
|
||||
"""Load one CLIP by 1-based index."""
|
||||
|
||||
@classmethod
|
||||
@@ -342,11 +415,13 @@ class NH_ClipIndex:
|
||||
if clip_name is None:
|
||||
return (_block(), "", -1)
|
||||
|
||||
import folder_paths
|
||||
import comfy.sd
|
||||
clip_path = folder_paths.get_full_path_or_raise("clip", clip_name)
|
||||
clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION)
|
||||
clip = comfy.sd.load_clip(ckpt_paths=[clip_path], embedding_directory=folder_paths.get_folder_paths("embeddings"), clip_type=clip_type)
|
||||
clip = self._load_cached(
|
||||
(clip_path, clip_type, _file_stamp(clip_path)),
|
||||
lambda: comfy.sd.load_clip(ckpt_paths=[clip_path], embedding_directory=folder_paths.get_folder_paths("embeddings"), clip_type=clip_type),
|
||||
)
|
||||
return (clip, clip_name, final_index)
|
||||
|
||||
|
||||
|
||||
+66
-47
@@ -1,11 +1,10 @@
|
||||
"""Phase 1: Logic Core — compare, branching, and gating nodes."""
|
||||
|
||||
import re
|
||||
|
||||
from comfy_execution.graph_utils import ExecutionBlocker
|
||||
|
||||
|
||||
_MAX_VALUE_MATCH_ITEMS = 64
|
||||
_SWITCH_SLOTS = 10 # input_0 .. input_9 of Switch N
|
||||
|
||||
|
||||
class NH_Compare:
|
||||
@@ -28,6 +27,11 @@ class NH_Compare:
|
||||
RETURN_NAMES = ("result", "a_passthrough", "b_passthrough")
|
||||
FUNCTION = "compare"
|
||||
CATEGORY = "NH-Nodes/Logic"
|
||||
DESCRIPTION = (
|
||||
'Compare a with b. "in" is plain membership: when b is a list, tuple, set or dict, a must be one of '
|
||||
"its items or keys; otherwise str(a) must be a substring of str(b), case-sensitive "
|
||||
'(13 is in "a13b" but not in "1024x768 3").'
|
||||
)
|
||||
|
||||
def compare(self, a, b, op, type_cast="auto"):
|
||||
if a is None or b is None:
|
||||
@@ -73,29 +77,19 @@ class NH_Compare:
|
||||
|
||||
@staticmethod
|
||||
def _contains(a, b):
|
||||
"""`a in b` as plain membership / substring, nothing fuzzier:
|
||||
|
||||
- b is a dict: a is one of its keys.
|
||||
- b is a list, tuple or set: a equals one of its items.
|
||||
- anything else: str(a) is a substring of str(b), case-sensitive, so 13 in "a13b" is True
|
||||
and 13 in "1024x768 3" is False.
|
||||
Values also match by their text (7 matches "7"), because widgets hand over strings.
|
||||
"""
|
||||
if isinstance(b, dict):
|
||||
return a in b or str(a) in [str(x) for x in b.keys()]
|
||||
if isinstance(b, (list, tuple, set)):
|
||||
return a in b or str(a) in [str(x) for x in b]
|
||||
needle = str(a)
|
||||
haystack = str(b)
|
||||
if needle in haystack:
|
||||
return True
|
||||
|
||||
if not NH_Compare._supports_subsequence(a, needle):
|
||||
return False
|
||||
|
||||
idx = 0
|
||||
for char in haystack:
|
||||
if idx < len(needle) and char == needle[idx]:
|
||||
idx += 1
|
||||
return idx == len(needle)
|
||||
|
||||
@staticmethod
|
||||
def _supports_subsequence(a, needle):
|
||||
if isinstance(a, (int, float)):
|
||||
return True
|
||||
return bool(re.fullmatch(r"[+-]?\d+(?:\.\d+)?", needle.strip()))
|
||||
return str(a) in str(b)
|
||||
|
||||
|
||||
class NH_LogicGate:
|
||||
@@ -148,8 +142,8 @@ class NH_IfElse:
|
||||
return {
|
||||
"required": {
|
||||
"condition": ("BOOLEAN", {"default": True}),
|
||||
"if_true": ("*",),
|
||||
"if_false": ("*",),
|
||||
"if_true": ("*", {"lazy": True}),
|
||||
"if_false": ("*", {"lazy": True}),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -157,8 +151,17 @@ class NH_IfElse:
|
||||
RETURN_NAMES = ("result", "condition_out")
|
||||
FUNCTION = "switch"
|
||||
CATEGORY = "NH-Nodes/Logic"
|
||||
DESCRIPTION = "Pick if_true or if_false by the condition. Only the chosen branch is executed; the other one never runs."
|
||||
|
||||
def switch(self, condition, if_true, if_false):
|
||||
def check_lazy_status(self, condition, if_true=None, if_false=None):
|
||||
# Chi xin nhanh duoc chon; nhanh con lai khong bao gio chay
|
||||
if condition and if_true is None:
|
||||
return ["if_true"]
|
||||
if (not condition) and if_false is None:
|
||||
return ["if_false"]
|
||||
return []
|
||||
|
||||
def switch(self, condition, if_true=None, if_false=None):
|
||||
return (if_true if condition else if_false, condition)
|
||||
|
||||
|
||||
@@ -172,16 +175,16 @@ class NH_SwitchN:
|
||||
"index": ("INT", {"default": 0, "min": 0, "max": 9}),
|
||||
},
|
||||
"optional": {
|
||||
"input_0": ("*",),
|
||||
"input_1": ("*",),
|
||||
"input_2": ("*",),
|
||||
"input_3": ("*",),
|
||||
"input_4": ("*",),
|
||||
"input_5": ("*",),
|
||||
"input_6": ("*",),
|
||||
"input_7": ("*",),
|
||||
"input_8": ("*",),
|
||||
"input_9": ("*",),
|
||||
"input_0": ("*", {"lazy": True}),
|
||||
"input_1": ("*", {"lazy": True}),
|
||||
"input_2": ("*", {"lazy": True}),
|
||||
"input_3": ("*", {"lazy": True}),
|
||||
"input_4": ("*", {"lazy": True}),
|
||||
"input_5": ("*", {"lazy": True}),
|
||||
"input_6": ("*", {"lazy": True}),
|
||||
"input_7": ("*", {"lazy": True}),
|
||||
"input_8": ("*", {"lazy": True}),
|
||||
"input_9": ("*", {"lazy": True}),
|
||||
},
|
||||
}
|
||||
|
||||
@@ -189,25 +192,41 @@ class NH_SwitchN:
|
||||
RETURN_NAMES = ("result", "count")
|
||||
FUNCTION = "switch"
|
||||
CATEGORY = "NH-Nodes/Logic"
|
||||
DESCRIPTION = (
|
||||
"Pick input_<index>. Only the chosen input is executed. If that slot is not connected, the lowest connected "
|
||||
"slot is used instead. count = number of connected slots. Connected inputs count even if they produce no "
|
||||
"value (an upstream node that returns nothing still takes its slot)."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _connected_slots(kwargs):
|
||||
# "Noi" = key co mat trong kwargs (slot chua noi thi ComfyUI khong gui key). Gia tri None chi co nghia
|
||||
# la nhanh lazy chua chay, nen khong dung gia tri de biet slot nao dang noi. Tang dan theo slot.
|
||||
return [i for i in range(_SWITCH_SLOTS) if f"input_{i}" in kwargs]
|
||||
|
||||
@staticmethod
|
||||
def _chosen_slot(index, slots):
|
||||
# Slot duoc chon; neu slot do khong noi thi roi ve slot noi thap nhat (nhu truoc day)
|
||||
return next((slot for slot in slots if slot == index), slots[0])
|
||||
|
||||
def check_lazy_status(self, index, **kwargs):
|
||||
slots = self._connected_slots(kwargs)
|
||||
if not slots:
|
||||
return []
|
||||
key = f"input_{self._chosen_slot(index, slots)}"
|
||||
return [key] if kwargs[key] is None else []
|
||||
|
||||
def switch(self, index, **kwargs):
|
||||
connected = {}
|
||||
for i in range(10):
|
||||
key = f"input_{i}"
|
||||
if key in kwargs and kwargs[key] is not None:
|
||||
connected[i] = kwargs[key]
|
||||
|
||||
count = len(connected)
|
||||
slots = self._connected_slots(kwargs)
|
||||
count = len(slots)
|
||||
if count == 0:
|
||||
print("[NH-Nodes] Warning: SwitchN has no connected inputs")
|
||||
return (None, 0)
|
||||
|
||||
if index in connected:
|
||||
return (connected[index], count)
|
||||
|
||||
first_key = sorted(connected.keys())[0]
|
||||
print(f"[NH-Nodes] SwitchN: index {index} not connected, falling back to input_{first_key}")
|
||||
return (connected[first_key], count)
|
||||
chosen = self._chosen_slot(index, slots)
|
||||
if chosen != index:
|
||||
print(f"[NH-Nodes] SwitchN: index {index} not connected, falling back to input_{chosen}")
|
||||
return (kwargs[f"input_{chosen}"], count)
|
||||
|
||||
|
||||
class NH_AnySwitchBoolean:
|
||||
|
||||
+60
-99
@@ -1,74 +1,15 @@
|
||||
"""Phase 2: Math & Random — NH_MathEval, NH_RandomChoice"""
|
||||
|
||||
import ast
|
||||
import operator
|
||||
import math
|
||||
import random
|
||||
|
||||
from ...core import safe_eval
|
||||
|
||||
# --- Safe eval engine ---
|
||||
_SAFE_OPS = {
|
||||
ast.Add: operator.add,
|
||||
ast.Sub: operator.sub,
|
||||
ast.Mult: operator.mul,
|
||||
ast.Div: operator.truediv,
|
||||
ast.FloorDiv: operator.floordiv,
|
||||
ast.Mod: operator.mod,
|
||||
ast.Pow: operator.pow,
|
||||
ast.USub: operator.neg,
|
||||
}
|
||||
|
||||
_SAFE_FUNCS = {
|
||||
"min": min, "max": max, "abs": abs,
|
||||
"round": round,
|
||||
"sqrt": math.sqrt, "floor": math.floor, "ceil": math.ceil,
|
||||
"clamp": lambda v, lo, hi: max(lo, min(hi, v)),
|
||||
}
|
||||
|
||||
|
||||
def _safe_eval(node, variables):
|
||||
"""De quy eval AST node, chi cho phep operators va functions an toan."""
|
||||
if isinstance(node, ast.Constant):
|
||||
return node.value
|
||||
elif isinstance(node, ast.Name):
|
||||
if node.id in variables:
|
||||
return variables[node.id]
|
||||
if node.id in _SAFE_FUNCS:
|
||||
return _SAFE_FUNCS[node.id]
|
||||
raise ValueError(f"Unknown variable: {node.id}")
|
||||
elif isinstance(node, ast.BinOp):
|
||||
left = _safe_eval(node.left, variables)
|
||||
right = _safe_eval(node.right, variables)
|
||||
op = _SAFE_OPS.get(type(node.op))
|
||||
if op is None:
|
||||
raise ValueError(f"Unsupported operator: {type(node.op).__name__}")
|
||||
if isinstance(node.op, ast.Div) and right == 0:
|
||||
print("[NH-Nodes] Warning: Division by zero, returning 0")
|
||||
return 0.0
|
||||
if isinstance(node.op, ast.Pow):
|
||||
# Gioi han power de tranh treo
|
||||
if abs(right) > 100:
|
||||
print("[NH-Nodes] Warning: Power too large, clamped to 100")
|
||||
right = 100 if right > 0 else -100
|
||||
return op(left, right)
|
||||
elif isinstance(node, ast.UnaryOp):
|
||||
operand = _safe_eval(node.operand, variables)
|
||||
op = _SAFE_OPS.get(type(node.op))
|
||||
if op is None:
|
||||
raise ValueError(f"Unsupported unary op: {type(node.op).__name__}")
|
||||
return op(operand)
|
||||
elif isinstance(node, ast.Call):
|
||||
func_name = node.func.id if isinstance(node.func, ast.Name) else None
|
||||
if func_name not in _SAFE_FUNCS:
|
||||
raise ValueError(f"Unknown function: {func_name}")
|
||||
args = [_safe_eval(arg, variables) for arg in node.args]
|
||||
return _SAFE_FUNCS[func_name](*args)
|
||||
else:
|
||||
raise ValueError(f"Unsupported expression type: {type(node).__name__}")
|
||||
# Random Choice co 6 input: input_0 .. input_5
|
||||
_CHOICE_SLOTS = 6
|
||||
|
||||
|
||||
class NH_MathEval:
|
||||
"""Tinh toan bieu thuc toan hoc an toan bang AST parsing."""
|
||||
"""Tinh toan bieu thuc toan hoc an toan bang AST parsing (core/safe_eval.py)."""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
@@ -94,15 +35,18 @@ class NH_MathEval:
|
||||
CATEGORY = "NH-Nodes/Logic"
|
||||
|
||||
def evaluate(self, expression, a=0.0, b=0.0, c=0.0, d=0.0, round_to=-1):
|
||||
variables = {"a": a, "b": b, "c": c, "d": d}
|
||||
# Bieu thuc sai hoac ket qua khong huu han -> ValueError (ComfyUI hien loi, khong tra 0 im lang).
|
||||
# x / 0, x // 0, x % 0 -> 0 kem canh bao.
|
||||
try:
|
||||
tree = ast.parse(expression, mode='eval')
|
||||
result = _safe_eval(tree.body, variables)
|
||||
except Exception as e:
|
||||
print(f"[NH-Nodes] MathEval error: {e}")
|
||||
result = 0.0
|
||||
|
||||
result = float(result)
|
||||
result = safe_eval.evaluate(
|
||||
expression,
|
||||
{"a": a, "b": b, "c": c, "d": d},
|
||||
safe_eval.MATH_FUNCTIONS,
|
||||
zero_division="zero",
|
||||
on_warning=lambda message: print(f"[NH-Nodes] MathEval warning: {message}"),
|
||||
)
|
||||
except safe_eval.ExpressionError as exc:
|
||||
raise ValueError(f"Math Eval: {exc} (expression: {expression!r})") from exc
|
||||
if round_to >= 0:
|
||||
result = round(result, round_to)
|
||||
|
||||
@@ -120,9 +64,9 @@ class NH_RandomChoice:
|
||||
"weights": ("STRING", {"default": "1,1,1,1,1,1"}),
|
||||
},
|
||||
"optional": {
|
||||
"input_0": ("*",), "input_1": ("*",),
|
||||
"input_2": ("*",), "input_3": ("*",),
|
||||
"input_4": ("*",), "input_5": ("*",),
|
||||
"input_0": ("*", {"lazy": True}), "input_1": ("*", {"lazy": True}),
|
||||
"input_2": ("*", {"lazy": True}), "input_3": ("*", {"lazy": True}),
|
||||
"input_4": ("*", {"lazy": True}), "input_5": ("*", {"lazy": True}),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -130,48 +74,65 @@ class NH_RandomChoice:
|
||||
RETURN_NAMES = ("result", "picked_index", "probabilities")
|
||||
FUNCTION = "choose"
|
||||
CATEGORY = "NH-Nodes/Logic"
|
||||
DESCRIPTION = (
|
||||
"Pick one connected input at random (weight N belongs to input_N, the seed is used as given). "
|
||||
"Only the picked input is executed; the other ones never run. Connected inputs count even if they produce "
|
||||
"no value (an upstream node that returns nothing still takes its slot)."
|
||||
)
|
||||
|
||||
def choose(self, seed, weights, **kwargs):
|
||||
connected = []
|
||||
for i in range(6):
|
||||
key = f"input_{i}"
|
||||
if key in kwargs and kwargs[key] is not None:
|
||||
connected.append((i, kwargs[key]))
|
||||
@staticmethod
|
||||
def _connected_slots(kwargs):
|
||||
# "Noi" = key co mat trong kwargs (slot chua noi thi ComfyUI khong gui key). Gia tri None chi co nghia
|
||||
# la nhanh lazy chua chay, nen khong dung gia tri de biet slot nao dang noi. Tang dan theo slot.
|
||||
return [i for i in range(_CHOICE_SLOTS) if f"input_{i}" in kwargs]
|
||||
|
||||
if not connected:
|
||||
return (None, -1, "")
|
||||
|
||||
# Parse weights
|
||||
w = []
|
||||
for s in weights.split(","):
|
||||
@staticmethod
|
||||
def _pick(seed, weights, slots):
|
||||
"""Tra ve (slot duoc chon, xac suat cua cac slot dang noi). Chi phu thuoc seed, weights va tap slot
|
||||
dang noi, khong dung gia tri input: nho vay check_lazy_status chon dung slot ma choose se tra ve."""
|
||||
# Weight thu i thuoc input_i (khong phai input thu i dang noi), slot thieu weight = 1
|
||||
slot_weights = []
|
||||
for s in weights.split(",")[:_CHOICE_SLOTS]:
|
||||
try:
|
||||
w.append(max(0.0, float(s.strip())))
|
||||
slot_weights.append(max(0.0, float(s.strip())))
|
||||
except (ValueError, TypeError):
|
||||
w.append(1.0)
|
||||
while len(w) < len(connected):
|
||||
w.append(1.0)
|
||||
w = w[:len(connected)]
|
||||
slot_weights.append(1.0)
|
||||
slot_weights += [1.0] * (_CHOICE_SLOTS - len(slot_weights))
|
||||
w = [slot_weights[i] for i in slots]
|
||||
|
||||
# Normalize
|
||||
total = sum(w) or 1.0
|
||||
probs = [x / total for x in w]
|
||||
|
||||
# Seed
|
||||
rng = random.Random(seed if seed > 0 else None)
|
||||
# Seed dung dung gia tri (0 cung tat dinh): ComfyUI cache output theo input
|
||||
rng = random.Random(seed)
|
||||
|
||||
# Weighted pick
|
||||
r = rng.random()
|
||||
cumsum = 0.0
|
||||
for idx, (orig_i, item) in enumerate(connected):
|
||||
cumsum += probs[idx]
|
||||
for slot, prob in zip(slots, probs):
|
||||
cumsum += prob
|
||||
if r <= cumsum:
|
||||
prob_str = ", ".join(f"{p:.1%}" for p in probs)
|
||||
return (item, orig_i, prob_str)
|
||||
return slot, probs
|
||||
|
||||
# Fallback
|
||||
orig_i, item = connected[-1]
|
||||
return slots[-1], probs
|
||||
|
||||
def check_lazy_status(self, seed, weights, **kwargs):
|
||||
slots = self._connected_slots(kwargs)
|
||||
if not slots:
|
||||
return []
|
||||
key = f"input_{self._pick(seed, weights, slots)[0]}"
|
||||
return [key] if kwargs[key] is None else []
|
||||
|
||||
def choose(self, seed, weights, **kwargs):
|
||||
slots = self._connected_slots(kwargs)
|
||||
if not slots:
|
||||
return (None, -1, "")
|
||||
|
||||
slot, probs = self._pick(seed, weights, slots)
|
||||
prob_str = ", ".join(f"{p:.1%}" for p in probs)
|
||||
return (item, orig_i, prob_str)
|
||||
return (kwargs[f"input_{slot}"], slot, prob_str)
|
||||
|
||||
|
||||
# --- Dang ky ---
|
||||
|
||||
@@ -1,24 +1,65 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
def get_bounding_box_and_ratio(mask_tensor):
|
||||
"""Tính toán bounding box và tỷ lệ (width / height) từ một mask tensor."""
|
||||
m_np = mask_tensor.cpu().numpy()
|
||||
rows, cols = np.where(m_np > 0.5)
|
||||
|
||||
if len(rows) == 0:
|
||||
return None, 1.0 # Trả về None và tỷ lệ 1:1 nếu mask rỗng
|
||||
from ...core.mask_ops import mask_bbox
|
||||
|
||||
|
||||
def _fit_window(start, length, limit):
|
||||
"""Cửa sổ [start, start + length) trượt vào trong [0, limit); lớn hơn limit thì cắt ở mép cuối."""
|
||||
start = max(0, min(start, limit - length))
|
||||
return start, min(limit, start + length)
|
||||
|
||||
|
||||
def _adjust_mask(mask_np, bbox, target_size, mode):
|
||||
"""Chỉnh hộp (bbox) của một mask [H, W] về đúng tỷ lệ target_size = (rộng, cao).
|
||||
Kích thước hộp mới tính bằng số nguyên nên không bị lệch 1px do số thực. Ngoài hộp luôn là 0."""
|
||||
height, width = mask_np.shape
|
||||
xmin, ymin, xmax, ymax = (int(v) for v in bbox)
|
||||
cur_w, cur_h = xmax - xmin + 1, ymax - ymin + 1
|
||||
target_w, target_h = target_size
|
||||
target_is_wider = target_w * cur_h > cur_w * target_h
|
||||
output = np.zeros_like(mask_np)
|
||||
|
||||
if mode in ("stretch", "pad"):
|
||||
# Nới một chiều: target rộng hơn thì nới ngang, ngược lại nới dọc
|
||||
new_w, new_h = cur_w, cur_h
|
||||
if target_is_wider:
|
||||
new_w = max(1, cur_h * target_w // target_h)
|
||||
else:
|
||||
new_h = max(1, cur_w * target_h // target_w)
|
||||
|
||||
# Hộp mới nới đối xứng quanh hộp cũ. Chạm mép ảnh thì trượt vào trong để vẫn đủ kích thước.
|
||||
x0, x1 = _fit_window(xmin - (new_w - cur_w) // 2, new_w, width)
|
||||
y0, y1 = _fit_window(ymin - (new_h - cur_h) // 2, new_h, height)
|
||||
|
||||
if mode == "stretch":
|
||||
# Kéo dãn nội dung (nearest) cho vừa hộp mới. Hộp lớn hơn cả ảnh thì chỉ hiện phần giữa.
|
||||
off_x = (new_w - (x1 - x0)) // 2
|
||||
off_y = (new_h - (y1 - y0)) // 2
|
||||
cols = (np.arange(x1 - x0) + off_x) * cur_w // new_w
|
||||
rows = (np.arange(y1 - y0) + off_y) * cur_h // new_h
|
||||
output[y0:y1, x0:x1] = mask_np[ymin:ymax + 1, xmin:xmax + 1][np.ix_(rows, cols)]
|
||||
else:
|
||||
# Pad: phần đệm thuộc vùng mask (1.0), nội dung gốc giữ nguyên vị trí cũ
|
||||
output[y0:y1, x0:x1] = 1.0
|
||||
output[ymin:ymax + 1, xmin:xmax + 1] = mask_np[ymin:ymax + 1, xmin:xmax + 1]
|
||||
|
||||
elif mode == "crop":
|
||||
# Cắt bớt một chiều quanh tâm; phần giữ lại không bị dịch chỗ
|
||||
new_w, new_h = cur_w, cur_h
|
||||
if target_is_wider:
|
||||
new_h = max(1, cur_w * target_h // target_w)
|
||||
else:
|
||||
new_w = max(1, cur_h * target_w // target_h)
|
||||
x0 = xmin + (cur_w - new_w) // 2
|
||||
y0 = ymin + (cur_h - new_h) // 2
|
||||
output[y0:y0 + new_h, x0:x0 + new_w] = mask_np[y0:y0 + new_h, x0:x0 + new_w]
|
||||
|
||||
else:
|
||||
raise ValueError(f"NH Mask Aspect Ratio Match: unknown mode '{mode}'")
|
||||
|
||||
return output
|
||||
|
||||
ymin, ymax = np.min(rows), np.max(rows)
|
||||
xmin, xmax = np.min(cols), np.max(cols)
|
||||
|
||||
bbox = (xmin, ymin, xmax, ymax)
|
||||
width = xmax - xmin + 1
|
||||
height = ymax - ymin + 1
|
||||
|
||||
aspect_ratio = width / height if height > 0 else 1.0
|
||||
return bbox, aspect_ratio
|
||||
|
||||
class NH_MaskAspectRatioMatch:
|
||||
@classmethod
|
||||
@@ -35,88 +76,38 @@ class NH_MaskAspectRatioMatch:
|
||||
RETURN_TYPES = ("MASK",)
|
||||
FUNCTION = "match_aspect_ratio"
|
||||
CATEGORY = "NH-Nodes/Mask"
|
||||
DESCRIPTION = (
|
||||
"Resizes the bounding box of mask_to_adjust to the aspect ratio of target_ratio_mask, around the same centre. "
|
||||
"stretch: scale the content to the new box. "
|
||||
"pad: keep the content and add mask pixels (1.0) around it until the box has the ratio. "
|
||||
"crop: keep the middle part of the content. "
|
||||
"A box that would leave the image slides back inside it. Every mask of the batch is processed; "
|
||||
"a shorter target batch repeats."
|
||||
)
|
||||
|
||||
def match_aspect_ratio(self, target_ratio_mask, mask_to_adjust, mode):
|
||||
mask_a = target_ratio_mask[0]
|
||||
mask_b = mask_to_adjust[0]
|
||||
|
||||
_ , ratio_a = get_bounding_box_and_ratio(mask_a)
|
||||
bbox_b, ratio_b = get_bounding_box_and_ratio(mask_b)
|
||||
|
||||
if bbox_b is None:
|
||||
# Chấp nhận mask [H, W] hoặc [B, H, W]; xử lý từng mask trong batch
|
||||
targets = target_ratio_mask.reshape(-1, *target_ratio_mask.shape[-2:])
|
||||
masks = mask_to_adjust.reshape(-1, *mask_to_adjust.shape[-2:])
|
||||
if targets.shape[0] == 0 or masks.shape[0] == 0:
|
||||
return (mask_to_adjust,)
|
||||
|
||||
xmin, ymin, xmax, ymax = bbox_b
|
||||
current_width = xmax - xmin + 1
|
||||
current_height = ymax - ymin + 1
|
||||
center_x = (xmin + xmax) // 2
|
||||
center_y = (ymin + ymax) // 2
|
||||
adjusted = []
|
||||
for index in range(masks.shape[0]):
|
||||
mask_np = masks[index].cpu().numpy()
|
||||
bbox_b = mask_bbox(mask_np)
|
||||
if bbox_b is None:
|
||||
adjusted.append(masks[index]) # Mask B rỗng: giữ nguyên
|
||||
continue
|
||||
|
||||
mask_b_np = mask_b.cpu().numpy()
|
||||
output_mask_np = np.zeros_like(mask_b_np)
|
||||
# Batch của mask A ngắn hơn thì lặp lại vòng tròn
|
||||
bbox_a = mask_bbox(targets[index % targets.shape[0]].cpu().numpy())
|
||||
# Mask A rỗng thì coi như tỷ lệ 1:1
|
||||
target_size = (1, 1) if bbox_a is None else (bbox_a[2] - bbox_a[0] + 1, bbox_a[3] - bbox_a[1] + 1)
|
||||
output_np = _adjust_mask(mask_np, bbox_b, target_size, mode)
|
||||
adjusted.append(torch.from_numpy(output_np).to(masks.device))
|
||||
|
||||
# Cắt nội dung thực sự của Mask B
|
||||
cropped_content_b = mask_b_np[ymin:ymax+1, xmin:xmax+1]
|
||||
|
||||
# ==========================================================
|
||||
# LOGIC CHO TỪNG CHẾ ĐỘ
|
||||
# ==========================================================
|
||||
|
||||
if mode == "stretch":
|
||||
new_height, new_width = current_height, current_width
|
||||
if ratio_a > ratio_b: # Target rộng hơn
|
||||
new_width = int(current_height * ratio_a)
|
||||
else: # Target cao hơn
|
||||
new_height = int(current_width / ratio_a)
|
||||
|
||||
# Kéo dãn nội dung đã cắt
|
||||
stretched_content = cv2.resize(cropped_content_b, (new_width, new_height), interpolation=cv2.INTER_NEAREST)
|
||||
|
||||
# Dán vào canvas mới tại vị trí trung tâm
|
||||
paste_x = center_x - new_width // 2
|
||||
paste_y = center_y - new_height // 2
|
||||
|
||||
# Giới hạn tọa độ để không bị tràn
|
||||
p_x_start, p_y_start = max(0, paste_x), max(0, paste_y)
|
||||
p_x_end, p_y_end = min(mask_b_np.shape[1], paste_x + new_width), min(mask_b_np.shape[0], paste_y + new_height)
|
||||
|
||||
s_x_start, s_y_start = p_x_start - paste_x, p_y_start - paste_y
|
||||
s_x_end, s_y_end = s_x_start + (p_x_end - p_x_start), s_y_start + (p_y_end - p_y_start)
|
||||
|
||||
output_mask_np[p_y_start:p_y_end, p_x_start:p_x_end] = stretched_content[s_y_start:s_y_end, s_x_start:s_x_end]
|
||||
|
||||
elif mode == "pad":
|
||||
new_height, new_width = current_height, current_width
|
||||
if ratio_a > ratio_b: # Target rộng hơn -> Thêm đệm ngang
|
||||
new_width = int(current_height * ratio_a)
|
||||
else: # Target cao hơn -> Thêm đệm dọc
|
||||
new_height = int(current_width / ratio_a)
|
||||
|
||||
# Dán nội dung gốc vào trung tâm của hộp mới lớn hơn
|
||||
paste_x = center_x - current_width // 2
|
||||
paste_y = center_y - current_height // 2
|
||||
output_mask_np[paste_y:paste_y + current_height, paste_x:paste_x + current_width] = cropped_content_b
|
||||
|
||||
elif mode == "crop":
|
||||
new_height, new_width = current_height, current_width
|
||||
if ratio_a > ratio_b: # Target rộng hơn -> Cắt bớt chiều dọc
|
||||
new_height = int(current_width / ratio_a)
|
||||
else: # Target cao hơn -> Cắt bớt chiều ngang
|
||||
new_width = int(current_height * ratio_a)
|
||||
|
||||
# Cắt bớt nội dung
|
||||
crop_y = (current_height - new_height) // 2
|
||||
crop_x = (current_width - new_width) // 2
|
||||
cropped_content = cropped_content_b[crop_y:crop_y + new_height, crop_x:crop_x + new_width]
|
||||
|
||||
# Dán nội dung đã cắt vào vị trí trung tâm
|
||||
paste_x = center_x - new_width // 2
|
||||
paste_y = center_y - new_height // 2
|
||||
output_mask_np[paste_y:paste_y + new_height, paste_x:paste_x + new_width] = cropped_content
|
||||
|
||||
# Chuyển lại sang định dạng tensor
|
||||
output_mask = torch.from_numpy(output_mask_np).unsqueeze(0)
|
||||
return (output_mask,)
|
||||
return (torch.stack(adjusted),)
|
||||
|
||||
|
||||
# Ánh xạ để ComfyUI nhận diện class
|
||||
@@ -127,4 +118,4 @@ NODE_CLASS_MAPPINGS = {
|
||||
# Tên sẽ hiển thị trong menu của ComfyUI
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"NH_MaskAspectRatioMatch": "NH Mask Aspect Ratio Match"
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,10 +1,50 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
from scipy.ndimage import binary_fill_holes, gaussian_filter, binary_dilation, binary_erosion
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from scipy.ndimage import binary_fill_holes, gaussian_filter
|
||||
|
||||
# Moi lan chi xu ly toi da ~16 trieu diem anh (64 MB moi buffer float32): batch dai (mask video) khong lam bo nho tang vot
|
||||
_CHUNK_PIXELS = 1 << 24
|
||||
|
||||
|
||||
def _running_extreme(x, radius, dim, reduce_fn):
|
||||
"""Max (reduce_fn = torch.maximum) hoac min (torch.minimum) trong cua so 2*radius+1 diem tam la moi diem,
|
||||
truot doc theo `dim` cua x [B, 1, H, W] (-1 = chieu rong, -2 = chieu cao).
|
||||
|
||||
Bien anh duoc nhan ban (replicate): ngoai mep coi nhu bang diem mep, nen mask cham mep khong bi an tu mep.
|
||||
Khong quet tung diem cua cua so (O(radius), max_pool2d tren CPU mat vai giay o kernel 255) ma nhan doi cua so:
|
||||
moi buoc la mot reduce_fn giua hai lat cat lech nhau `step` diem, cua so dai 1, 2, 4, ... -> O(log radius) buoc,
|
||||
chay tren thiet bi cua tensor. Ket qua chinh xac (max/min khong lam tron).
|
||||
"""
|
||||
size = x.shape[dim]
|
||||
radius = min(radius, size - 1) # cua so rong hon anh thi cung chi thay het truc do
|
||||
if radius <= 0:
|
||||
return x
|
||||
window = 2 * radius + 1
|
||||
pad = (radius, radius, 0, 0) if dim == -1 else (0, 0, radius, radius)
|
||||
run = F.pad(x, pad, mode="replicate")
|
||||
step = 1
|
||||
while 2 * step <= window:
|
||||
length = run.shape[dim] - step
|
||||
run = reduce_fn(run.narrow(dim, 0, length), run.narrow(dim, step, length))
|
||||
step *= 2
|
||||
# run[i] = reduce(padded[i : i + step]); hai doan dai `step` chong nhau phu dung cua so `window`
|
||||
return reduce_fn(run.narrow(dim, 0, size), run.narrow(dim, window - step, size))
|
||||
|
||||
|
||||
def _expand(x, amount, dim):
|
||||
"""amount > 0: no ra (dilation, max); amount < 0: co lai (erosion, min); 0: giu nguyen (tra ve chinh x).
|
||||
|amount| la do dai kernel theo truc do: le thi dung nguyen (5 -> moi ben 2 diem). Chan thi lam tron len so
|
||||
le ke tiep (10 -> 11, moi ben 5 diem) de ket qua doi xung; ban cu kernel chan lech nua diem ve mot phia."""
|
||||
if amount == 0:
|
||||
return x
|
||||
reduce_fn = torch.maximum if amount > 0 else torch.minimum
|
||||
return _running_extreme(x, abs(amount) // 2, dim, reduce_fn)
|
||||
|
||||
|
||||
class NH_MaskMorphology:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"mask": ("MASK",),
|
||||
@@ -18,48 +58,63 @@ class NH_MaskMorphology:
|
||||
RETURN_TYPES = ("MASK",)
|
||||
FUNCTION = "process"
|
||||
CATEGORY = "NH-Nodes/Mask"
|
||||
DESCRIPTION = (
|
||||
"Grow (positive) or shrink (negative) the mask along each axis; the value is the kernel length in pixels, "
|
||||
"so 5 moves each edge by 2 px. An even length is rounded up to the next odd one (10 acts like 11) so the "
|
||||
"result stays centred. Soft masks keep their gradient and the image border never erodes the mask. "
|
||||
"Order: horizontal, vertical, fill_holes (mask > 0.5), then Gaussian blur_radius."
|
||||
)
|
||||
|
||||
def process(self, mask, horizontal_expand, vertical_expand, fill_holes, blur_radius):
|
||||
masks_out = []
|
||||
|
||||
for m in mask:
|
||||
m_np = m.cpu().numpy()
|
||||
|
||||
# 1. Co giãn theo chiều ngang
|
||||
if horizontal_expand != 0:
|
||||
kernel_width = abs(horizontal_expand)
|
||||
# Bỏ qua nếu kernel = 0
|
||||
if kernel_width > 0:
|
||||
structure_h = np.ones((1, kernel_width))
|
||||
if horizontal_expand > 0:
|
||||
m_np = binary_dilation(m_np, structure=structure_h).astype(m_np.dtype)
|
||||
else:
|
||||
m_np = binary_erosion(m_np, structure=structure_h).astype(m_np.dtype)
|
||||
if mask.dim() == 2:
|
||||
mask = mask.unsqueeze(0)
|
||||
if mask.dim() != 3:
|
||||
raise ValueError(f"Mask Morphology expects a MASK [batch, height, width], got shape {tuple(mask.shape)}")
|
||||
|
||||
# 2. Co giãn theo chiều dọc
|
||||
if vertical_expand != 0:
|
||||
kernel_height = abs(vertical_expand)
|
||||
# Bỏ qua nếu kernel = 0
|
||||
if kernel_height > 0:
|
||||
structure_v = np.ones((kernel_height, 1))
|
||||
if vertical_expand > 0:
|
||||
m_np = binary_dilation(m_np, structure=structure_v).astype(m_np.dtype)
|
||||
else:
|
||||
m_np = binary_erosion(m_np, structure=structure_v).astype(m_np.dtype)
|
||||
# Theo lo tren thiet bi cua tensor (CPU hoac GPU): anh nho gop chung mot lan chay, anh lon tung cai
|
||||
source = (mask if mask.is_floating_point() else mask.float()).unsqueeze(1) # [B, 1, H, W]
|
||||
count, _, height, width = source.shape
|
||||
if height == 0 or width == 0:
|
||||
# Khong co diem anh de no / co / lap / lam mo: tra ve nguyen ven (ban sao) nhu ban scipy cu. F.pad(replicate)
|
||||
# khong pad duoc tensor rong nen khong de no chay.
|
||||
return (source.squeeze(1).clone(),)
|
||||
per_chunk = max(1, _CHUNK_PIXELS // max(1, height * width))
|
||||
if count <= per_chunk:
|
||||
grown = self._grow(source, horizontal_expand, vertical_expand)
|
||||
else:
|
||||
grown = torch.empty_like(source)
|
||||
for start in range(0, count, per_chunk):
|
||||
chunk = source[start : start + per_chunk]
|
||||
grown[start : start + per_chunk] = self._grow(chunk, horizontal_expand, vertical_expand)
|
||||
result = grown.squeeze(1)
|
||||
|
||||
# 3. Lấp đầy lỗ trống (nếu cần)
|
||||
# 3. Lap day lo trong, 4. lam mo: scipy, tung anh tren CPU (giu nguyen ket qua cu)
|
||||
if fill_holes or blur_radius > 0:
|
||||
result = self._fill_and_blur(result, fill_holes, blur_radius)
|
||||
elif grown is source:
|
||||
result = result.clone() # khong doi gi: van tra ban sao, node phia sau co the sua tai cho
|
||||
return (result,)
|
||||
|
||||
@staticmethod
|
||||
def _grow(x, horizontal_expand, vertical_expand):
|
||||
x = _expand(x, horizontal_expand, -1) # 1. chieu ngang truoc
|
||||
return _expand(x, vertical_expand, -2) # 2. roi chieu doc
|
||||
|
||||
@staticmethod
|
||||
def _fill_and_blur(x, fill_holes, blur_radius):
|
||||
arrays = x.detach().cpu()
|
||||
if arrays.dtype not in (torch.float32, torch.float64):
|
||||
arrays = arrays.float() # numpy/scipy khong nhan half/bfloat16
|
||||
arrays = arrays.numpy()
|
||||
out = np.empty_like(arrays)
|
||||
for i, m_np in enumerate(arrays):
|
||||
if fill_holes:
|
||||
m_bool = m_np > 0.5
|
||||
m_filled_bool = binary_fill_holes(m_bool)
|
||||
m_np = m_filled_bool.astype(m_np.dtype)
|
||||
|
||||
# 4. Làm mờ
|
||||
m_np = binary_fill_holes(m_np > 0.5).astype(m_np.dtype)
|
||||
if blur_radius > 0:
|
||||
m_np = gaussian_filter(m_np, sigma=blur_radius)
|
||||
|
||||
masks_out.append(torch.from_numpy(m_np))
|
||||
out[i] = m_np
|
||||
return torch.from_numpy(out).to(device=x.device, dtype=x.dtype)
|
||||
|
||||
return (torch.stack(masks_out),)
|
||||
|
||||
NODE_CLASS_MAPPINGS = { "NH_MaskAdvancedMorphology": NH_MaskMorphology }
|
||||
NODE_DISPLAY_NAME_MAPPINGS = { "NH_MaskAdvancedMorphology": "Mask Morphology (NH)" }
|
||||
NODE_DISPLAY_NAME_MAPPINGS = { "NH_MaskAdvancedMorphology": "Mask Morphology (NH)" }
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
import numpy as np
|
||||
from ...core.mask_ops import mask_bbox
|
||||
|
||||
# File này chứa logic cho node Mask Properties (phiên bản gốc).
|
||||
|
||||
@@ -17,32 +17,19 @@ class NH_MaskProperties:
|
||||
CATEGORY = "NH-Nodes/Mask"
|
||||
|
||||
def get_properties(self, mask):
|
||||
# Chỉ xử lý mask đầu tiên trong batch
|
||||
m = mask[0]
|
||||
|
||||
# Chuyển tensor sang numpy array
|
||||
m_np = m.cpu().numpy()
|
||||
|
||||
# Lấy tọa độ của tất cả các pixel không phải màu đen (giá trị > 0.5)
|
||||
rows, cols = np.where(m_np > 0.5)
|
||||
|
||||
# Chỉ xử lý mask đầu tiên trong batch; hộp bao quanh các pixel > 0.5
|
||||
bbox = mask_bbox(mask[0].cpu().numpy())
|
||||
|
||||
# Xử lý trường hợp mask rỗng (toàn màu đen)
|
||||
if len(rows) == 0:
|
||||
if bbox is None:
|
||||
return (0, 0, 0, 0, 0, 0, (0, 0, 0, 0))
|
||||
|
||||
# Tìm các giá trị min và max để xác định hộp bao quanh
|
||||
x1 = int(np.min(cols))
|
||||
y1 = int(np.min(rows))
|
||||
x2 = int(np.max(cols))
|
||||
y2 = int(np.max(rows))
|
||||
|
||||
|
||||
x1, y1, x2, y2 = bbox
|
||||
|
||||
# Tính toán chiều rộng và chiều cao
|
||||
width = x2 - x1 + 1
|
||||
height = y2 - y1 + 1
|
||||
|
||||
# Tạo bounding box tuple để trả về
|
||||
bbox = (x1, y1, x2, y2)
|
||||
|
||||
|
||||
# Trả về tất cả các giá trị theo thứ tự đã định nghĩa
|
||||
return (width, height, x1, y1, x2, y2, bbox)
|
||||
|
||||
@@ -54,4 +41,4 @@ NODE_CLASS_MAPPINGS = {
|
||||
# Tên sẽ hiển thị trong menu của ComfyUI
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"NH_MaskProperties": "Mask Properties (NH)"
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
from ...core.mask_ops import mask_bbox
|
||||
|
||||
class NH_MaskToBBox:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
@@ -22,21 +24,19 @@ class NH_MaskToBBox:
|
||||
|
||||
for i in range(batch_size):
|
||||
single_mask_np = mask[i].cpu().numpy()
|
||||
|
||||
rows, cols = np.where(single_mask_np > 0.5)
|
||||
|
||||
|
||||
bbox_mask_np = np.zeros_like(single_mask_np)
|
||||
|
||||
if len(rows) > 0 and len(cols) > 0:
|
||||
ymin, ymax = np.min(rows), np.max(rows)
|
||||
xmin, xmax = np.min(cols), np.max(cols)
|
||||
|
||||
bbox = mask_bbox(single_mask_np)
|
||||
if bbox is not None:
|
||||
xmin, ymin, xmax, ymax = bbox
|
||||
|
||||
bbox_mask_np[ymin:ymax+1, xmin:xmax+1] = 1.0
|
||||
|
||||
output_masks.append(torch.from_numpy(bbox_mask_np))
|
||||
|
||||
final_output = torch.stack(output_masks)
|
||||
|
||||
|
||||
return (final_output,)
|
||||
|
||||
# Ánh xạ class không đổi
|
||||
@@ -47,4 +47,4 @@ NODE_CLASS_MAPPINGS = {
|
||||
# THAY ĐỔI Ở ĐÂY: Đổi tên hiển thị của node
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"NH_MaskToBBox": "NH Create Box Mask"
|
||||
}
|
||||
}
|
||||
|
||||
@@ -19,6 +19,7 @@ from ...core.resolution_data import (
|
||||
PRESET_LABELS_BY_MODEL,
|
||||
PRESET_LOOKUP_BY_MODEL,
|
||||
)
|
||||
from ...core.utils import parse_hex_color
|
||||
|
||||
|
||||
TARGET_RESOLUTION_LEVELS = [f"{mp} MP" for mp in range(1, 13)]
|
||||
@@ -46,6 +47,9 @@ _PIL_RESAMPLE_METHODS = {
|
||||
"box": _RESAMPLING.BOX,
|
||||
"hamming": _RESAMPLING.HAMMING,
|
||||
}
|
||||
# Lanczos has negative lobes, so it overshoots beside hard edges. Box and Hamming are plain weighted averages (their
|
||||
# weights are all positive), so they stay inside [0, 1] and need no clipping between the width and the height pass.
|
||||
_RINGING_METHODS = {"lanczos"}
|
||||
|
||||
|
||||
def _clamp(value, lower, upper):
|
||||
@@ -68,18 +72,6 @@ def _ratio_label(width, height):
|
||||
return f"{ratio_w}:{ratio_h}"
|
||||
|
||||
|
||||
def _parse_hex_color(value, default=(255, 255, 255)):
|
||||
text = str(value or "").strip().lstrip("#")
|
||||
try:
|
||||
if len(text) == 3:
|
||||
text = "".join(char * 2 for char in text)
|
||||
if len(text) != 6:
|
||||
raise ValueError
|
||||
return tuple(int(text[index:index + 2], 16) for index in (0, 2, 4))
|
||||
except Exception:
|
||||
return default
|
||||
|
||||
|
||||
def _needs_warning(entry, extra_warnings):
|
||||
if extra_warnings:
|
||||
return True
|
||||
@@ -89,14 +81,21 @@ def _needs_warning(entry, extra_warnings):
|
||||
return any(marker in text for marker in ("community", "extrapolation", "rủi ro", "risk"))
|
||||
|
||||
|
||||
def _build_info(model_spec, tier, ratio, width, height, entry=None, extra_warnings=None):
|
||||
def _build_info(model_spec, tier, ratio, width, height, entry=None, extra_warnings=None, warnings=None):
|
||||
"""Info text of the resolution nodes. ``warnings`` become 'Warning:' lines.
|
||||
|
||||
Presets carry their own reliability warning; max_res only bounds custom sizes (see resolution_data).
|
||||
"""
|
||||
extra_warnings = extra_warnings or []
|
||||
warnings = list(warnings or [])
|
||||
mp = (int(width) * int(height)) / 1_000_000
|
||||
prefix = "⚠️ " if _needs_warning(entry, extra_warnings) else ""
|
||||
prefix = "⚠️ " if warnings or _needs_warning(entry, extra_warnings) else ""
|
||||
info = f"{prefix}{model_spec['name']} | {tier} | {ratio} | {int(width)}×{int(height)} | {mp:.2f}MP"
|
||||
warning = entry.get("warning") if entry else ""
|
||||
if warning:
|
||||
info += f"\nWarning: {warning}"
|
||||
for message in warnings:
|
||||
info += f"\nWarning: {message}"
|
||||
note = model_spec.get("note")
|
||||
if note:
|
||||
info += f"\nNote: {note}"
|
||||
@@ -335,25 +334,57 @@ def _build_resize_data(source_width, source_height, target_width, target_height,
|
||||
return data
|
||||
|
||||
|
||||
def _resize_planes_pil(planes, size, resample, premultiplied):
|
||||
"""Resize [C, H, W] float planes with one Pillow call and clip the result to [0, 1].
|
||||
|
||||
Pillow's 8-bit path clips after each of its two passes (width, then height); on float planes it clips nothing.
|
||||
Lanczos overshoots beside hard edges, so the caller resizes one direction per call to clip in between: otherwise
|
||||
the second pass amplifies the overshoot of the first (15-33/255 off on text and line art). Premultiplied colour
|
||||
cannot be brighter than its alpha, so it is clipped to [0, alpha].
|
||||
"""
|
||||
resized = np.empty((planes.shape[0], size[1], size[0]), dtype=np.float32)
|
||||
for index in (3, 0, 1, 2) if premultiplied else range(planes.shape[0]): # alpha first, it bounds the colour
|
||||
upper = resized[3] if premultiplied and index != 3 else 1.0
|
||||
np.clip(np.asarray(Image.fromarray(planes[index]).resize(size, resample)), 0.0, upper, out=resized[index])
|
||||
return resized
|
||||
|
||||
|
||||
def _resize_bchw_pil(samples, height, width, resize_method):
|
||||
"""Resize [B, C, H, W] with Pillow's Lanczos/box/Hamming filters, in float.
|
||||
|
||||
Every channel goes through a 32-bit float ("F") image, so the result keeps its precision instead of being rounded
|
||||
to 8 bits. Four-channel images are resized with premultiplied alpha, as Pillow does for RGBA, so colour does not
|
||||
bleed in from transparent pixels. Lanczos resizes the width and the height in two passes, as Pillow does inside one
|
||||
resize, but clips each pass to [0, 1] (colour to [0, alpha] when premultiplied), like Pillow's 8-bit path does.
|
||||
"""
|
||||
original_dtype = samples.dtype
|
||||
original_device = samples.device
|
||||
resample = _PIL_RESAMPLE_METHODS[resize_method]
|
||||
size = (int(width), int(height))
|
||||
resized_items = []
|
||||
|
||||
for sample in samples:
|
||||
array = sample.movedim(0, -1).detach().cpu().numpy()
|
||||
array = np.clip(array * 255.0, 0, 255).astype(np.uint8)
|
||||
if array.shape[-1] == 1:
|
||||
pil_image = Image.fromarray(array[:, :, 0], mode="L")
|
||||
else:
|
||||
pil_image = Image.fromarray(array)
|
||||
|
||||
pil_image = pil_image.resize((int(width), int(height)), resample)
|
||||
resized_array = np.array(pil_image).astype(np.float32) / 255.0
|
||||
if resized_array.ndim == 2:
|
||||
resized_array = resized_array[:, :, None]
|
||||
resized_items.append(torch.from_numpy(resized_array).movedim(-1, 0))
|
||||
# [C, H, W] float32 in 0-1 (NaN and out-of-range values are cleaned up first, as the uint8 round trip did).
|
||||
# Done in torch, which is several times faster than numpy on large images.
|
||||
cleaned = torch.nan_to_num(sample.detach().cpu().float(), nan=0.0, posinf=1.0, neginf=0.0).clamp_(0.0, 1.0)
|
||||
if (cleaned.shape[2], cleaned.shape[1]) == size:
|
||||
# Nothing to resample; skipping the premultiply round trip keeps the colour under transparent pixels.
|
||||
resized_items.append(cleaned)
|
||||
continue
|
||||
planes = cleaned.contiguous().numpy()
|
||||
premultiply = planes.shape[0] == 4
|
||||
if premultiply:
|
||||
planes = np.concatenate((planes[:3] * planes[3:], planes[3:]))
|
||||
# Lanczos: the width first, then the height (Pillow's own order), clipped after each. The others: one call.
|
||||
passes = ((size[0], planes.shape[1]), size) if resize_method in _RINGING_METHODS else (size,)
|
||||
for pass_size in passes:
|
||||
if pass_size != (planes.shape[2], planes.shape[1]):
|
||||
planes = _resize_planes_pil(planes, pass_size, resample, premultiply)
|
||||
if premultiply:
|
||||
alpha = planes[3:]
|
||||
colour = np.divide(planes[:3], alpha, out=np.zeros_like(planes[:3]), where=alpha > 1e-6)
|
||||
planes = np.concatenate((colour, alpha))
|
||||
resized_items.append(torch.from_numpy(planes))
|
||||
|
||||
return torch.stack(resized_items, dim=0).to(device=original_device, dtype=original_dtype)
|
||||
|
||||
@@ -401,14 +432,15 @@ def _resize_image_batch_with_data(images, target_width, target_height, resize_mo
|
||||
cropped = resized[:, :, top:top + int(target_height), left:left + int(target_width)]
|
||||
return cropped.movedim(1, -1), resize_data
|
||||
|
||||
fill_rgb = _parse_hex_color(pad_color_hex, (255, 255, 255))
|
||||
fill_rgba = parse_hex_color(pad_color_hex, (255, 255, 255, 255))
|
||||
canvas = torch.empty(
|
||||
(batch, channels, int(target_height), int(target_width)),
|
||||
dtype=images.dtype,
|
||||
device=images.device,
|
||||
)
|
||||
for channel in range(channels):
|
||||
canvas[:, channel, :, :] = fill_rgb[min(channel, 2)] / 255.0
|
||||
# r, g, b, then alpha for the fourth channel of RGBA images (opaque unless the colour is #RRGGBBAA).
|
||||
canvas[:, channel, :, :] = fill_rgba[min(channel, 3)] / 255.0
|
||||
|
||||
left = int(resize_data["content_left"])
|
||||
top = int(resize_data["content_top"])
|
||||
@@ -450,6 +482,76 @@ def _has_padding(resize_data):
|
||||
return any(int(resize_data.get(key, 0)) > 0 for key in ("pad_left", "pad_top", "pad_right", "pad_bottom"))
|
||||
|
||||
|
||||
def _has_crop(resize_data):
|
||||
return any(int(resize_data.get(key, 0)) > 0 for key in ("crop_left", "crop_top", "crop_right", "crop_bottom"))
|
||||
|
||||
|
||||
def _cropped_footprint(resize_data):
|
||||
"""Where the picture of a ``cover_crop`` result sits in the source frame.
|
||||
|
||||
Returns ``(source_width, source_height, left, top, right, bottom)`` in source pixels, or None when ``resize_data``
|
||||
does not say. The resize scaled the source uniformly by resized/source and the crop kept ``target`` pixels from
|
||||
offset ``crop_left``/``crop_top`` of that, so the kept window maps back by source/resized. A source pixel belongs
|
||||
to it when its centre is inside the window, the rule NH Restore Image Geometry uses for its mask.
|
||||
"""
|
||||
try:
|
||||
source_width, source_height, resized_width, resized_height, target_width, target_height, crop_left, crop_top = (
|
||||
int(resize_data[key])
|
||||
for key in (
|
||||
"source_width", "source_height", "resized_width", "resized_height",
|
||||
"target_width", "target_height", "crop_left", "crop_top",
|
||||
)
|
||||
)
|
||||
except (KeyError, TypeError, ValueError):
|
||||
return None
|
||||
if min(source_width, source_height, resized_width, resized_height, target_width, target_height) <= 0:
|
||||
return None
|
||||
|
||||
def span(crop, target, resized, source):
|
||||
start = _clamp(math.ceil(crop * source / resized - 0.5), 0, source - 1)
|
||||
end = _clamp(math.ceil((crop + target) * source / resized - 0.5), start + 1, source)
|
||||
return start, end
|
||||
|
||||
left, right = span(crop_left, target_width, resized_width, source_width)
|
||||
top, bottom = span(crop_top, target_height, resized_height, source_height)
|
||||
return source_width, source_height, left, top, right, bottom
|
||||
|
||||
|
||||
def _scale_into_source_frame(image, footprint, resize_method):
|
||||
"""Scale [B, H, W, C] into ``footprint`` and fill the rest of the source frame with the nearest edge pixels."""
|
||||
source_width, source_height, left, top, right, bottom = footprint
|
||||
scaled = _resize_image_batch(image, right - left, bottom - top, "stretch", resize_method, "#000000")
|
||||
padding = (left, source_width - right, top, source_height - bottom)
|
||||
return F.pad(scaled.movedim(-1, 1), padding, mode="replicate").movedim(1, -1)
|
||||
|
||||
|
||||
def _cropped_resize_info(resize_data, image_width, image_height, footprint=None):
|
||||
"""Explain what happened to a ``cover_crop`` result; ``footprint`` is None when it was returned unchanged."""
|
||||
source_width = int(resize_data.get("source_width", image_width))
|
||||
source_height = int(resize_data.get("source_height", image_height))
|
||||
target_width = int(resize_data.get("target_width", image_width))
|
||||
target_height = int(resize_data.get("target_height", image_height))
|
||||
crop_text = " ".join(f"{side[0].upper()}{int(resize_data.get(f'crop_{side}', 0))}" for side in ("left", "top", "right", "bottom"))
|
||||
details = (
|
||||
f"source {source_width}×{source_height}; resize target {target_width}×{target_height}; "
|
||||
f"input {image_width}×{image_height}; crop {crop_text}"
|
||||
)
|
||||
if footprint is None:
|
||||
return (
|
||||
"cover_crop resize detected: the image was cropped, so there is no padding to remove and the cropped "
|
||||
"borders cannot be recovered; returned unchanged. output_size=original_size scales it back into the "
|
||||
"source frame (it needs complete resize_data), and NH Restore Image Geometry composites it onto the "
|
||||
f"original image; {details}"
|
||||
)
|
||||
_, _, left, top, right, bottom = footprint
|
||||
return (
|
||||
"cover_crop resize detected: the cropped borders cannot be recovered. The image was scaled back uniformly "
|
||||
f"into the source frame (x {left}-{right}, y {top}-{bottom}); the borders that were cut off repeat its nearest "
|
||||
"edge pixels. NH Restore Image Geometry composites onto the original image instead, which keeps its real "
|
||||
f"borders; {details}; output {source_width}×{source_height}"
|
||||
)
|
||||
|
||||
|
||||
class NH_SmartResolutionPicker:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
@@ -486,13 +588,21 @@ class NH_SmartResolutionPicker:
|
||||
swap_orientation=False,
|
||||
batch_size=1,
|
||||
):
|
||||
# VALIDATE_INPUTS accepts any value, so a saved workflow can carry a model or preset label that no longer
|
||||
# exists. The default is used for it, and the info says so.
|
||||
warnings = []
|
||||
if model_family not in MODEL_ID_BY_LABEL:
|
||||
warnings.append(f"unknown model_family '{model_family}', using {DEFAULT_MODEL_LABEL}")
|
||||
model_family = DEFAULT_MODEL_LABEL
|
||||
|
||||
preset_lookup = PRESET_LOOKUP_BY_MODEL.get(model_family, {})
|
||||
model_presets = PRESET_LABELS_BY_MODEL.get(model_family, [])
|
||||
default_preset = DEFAULT_PRESET if model_family == DEFAULT_MODEL_LABEL else model_presets[0]
|
||||
entry = preset_lookup.get(preset) or preset_lookup[default_preset]
|
||||
entry = preset_lookup.get(preset)
|
||||
if entry is None:
|
||||
entry = preset_lookup[default_preset]
|
||||
if not use_custom: # a custom size does not use the preset's size
|
||||
warnings.append(f"preset '{preset}' not found for {model_family}, using '{default_preset}'")
|
||||
model_spec = MODEL_SPECS[entry["model_id"]]
|
||||
extra_warnings = []
|
||||
|
||||
@@ -530,7 +640,7 @@ class NH_SmartResolutionPicker:
|
||||
"downscale_ratio_spacial": downscale,
|
||||
}
|
||||
aspect_ratio = float(width) / max(float(height), 1.0)
|
||||
info = _build_info(model_spec, tier, ratio, width, height, entry_for_warning, extra_warnings)
|
||||
info = _build_info(model_spec, tier, ratio, width, height, entry_for_warning, extra_warnings, warnings)
|
||||
|
||||
return (int(width), int(height), latent, info, aspect_ratio)
|
||||
|
||||
@@ -754,6 +864,17 @@ class NH_RemoveResizePadding:
|
||||
f"No resize_data metadata found; passed through {image_width}×{image_height}.",
|
||||
)
|
||||
|
||||
if resize_data.get("resize_mode") == "cover_crop" and _has_crop(resize_data):
|
||||
# A cover_crop resize has no padding, and the part it cut off is gone. Stretching what is left to the
|
||||
# source size would distort it (about 2.6% for 1500x1000 -> 1216x832), so for the original size the image
|
||||
# is scaled back by the same uniform factor, put where it came from, and the lost borders repeat its edges.
|
||||
footprint = _cropped_footprint(resize_data) if output_size == "original_size" else None
|
||||
if footprint is None:
|
||||
return (image, image_width, image_height, _cropped_resize_info(resize_data, image_width, image_height))
|
||||
restored = _scale_into_source_frame(image, footprint, resize_method)
|
||||
info = _cropped_resize_info(resize_data, image_width, image_height, footprint)
|
||||
return (restored, footprint[0], footprint[1], info)
|
||||
|
||||
left, top, right, bottom = _scaled_content_bounds(resize_data, image_width, image_height)
|
||||
cropped = image[:, top:bottom, left:right, :]
|
||||
|
||||
|
||||
@@ -13,12 +13,15 @@ def _as_prompt_text(value):
|
||||
|
||||
|
||||
def _remove_placeholder_preserve_lines(text, name):
|
||||
# Chi bo placeholder + khoang trang truoc no (dau dong: ca khoang trang sau no).
|
||||
# Dau phay de _cleanup_prompt_line gop: "a, {x}, b" -> "a,, b" -> "a, b" (khong dinh thanh "ab").
|
||||
placeholder = r"\{" + re.escape(name) + r"\}"
|
||||
horizontal_ws = r"[^\S\r\n]*"
|
||||
return re.sub(
|
||||
horizontal_ws + r",?" + horizontal_ws + placeholder + horizontal_ws + r",?" + horizontal_ws,
|
||||
r"^" + horizontal_ws + placeholder + horizontal_ws + r"|" + horizontal_ws + placeholder,
|
||||
"",
|
||||
text,
|
||||
flags=re.MULTILINE,
|
||||
)
|
||||
|
||||
|
||||
@@ -147,7 +150,8 @@ class NH_PromptScheduler:
|
||||
idx = cycle_len - pos
|
||||
|
||||
elif mode == "random":
|
||||
rng = random.Random(seed + current_step if seed > 0 else None)
|
||||
# Seed dung dung gia tri (0 cung tat dinh): ComfyUI cache output theo input
|
||||
rng = random.Random(seed + current_step)
|
||||
idx = rng.randint(0, n - 1)
|
||||
|
||||
else:
|
||||
|
||||
@@ -191,7 +191,8 @@ class NH_RegexExtract:
|
||||
if mode == "match":
|
||||
m = re.search(pattern, text)
|
||||
if m:
|
||||
groups = "\n".join(m.groups()) if m.groups() else m.group(0)
|
||||
# Group tuy chon khong khop la None -> chuoi rong
|
||||
groups = "\n".join(g or "" for g in m.groups()) if m.groups() else m.group(0)
|
||||
return (m.group(0), True, groups, 1)
|
||||
return ("", False, "", 0)
|
||||
|
||||
|
||||
@@ -70,8 +70,8 @@ class NH_TextSplitLines:
|
||||
if count == 0:
|
||||
return ([""], "", 0, "(empty)", True, True)
|
||||
|
||||
# Reorder by mode
|
||||
rng = random.Random(seed if seed > 0 else None)
|
||||
# Reorder by mode. The seed is used as given (0 too): ComfyUI caches the output per input.
|
||||
rng = random.Random(seed)
|
||||
|
||||
if mode == "shuffle":
|
||||
lines = lines[:]
|
||||
@@ -132,7 +132,8 @@ class NH_TextRandomLine:
|
||||
if total == 0:
|
||||
return ("", -1, 0)
|
||||
|
||||
rng = random.Random(seed if seed > 0 else None)
|
||||
# The seed is used as given (0 too): ComfyUI caches the output per input.
|
||||
rng = random.Random(seed)
|
||||
idx = rng.randint(0, total - 1)
|
||||
|
||||
return (lines[idx], idx, total)
|
||||
|
||||
@@ -1,6 +1,9 @@
|
||||
# Kiểu dữ liệu tùy chỉnh cho đường ống đa năng
|
||||
NH_UNIVERSAL_PIPE = ("NH_UNIVERSAL_PIPE",)
|
||||
|
||||
# Số ô input của Pack (input_0 .. input_9) cũng là độ dài của package
|
||||
_PIPE_SLOTS = 10
|
||||
|
||||
class NH_PackUniversal:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
@@ -19,9 +22,9 @@ class NH_PackUniversal:
|
||||
CATEGORY = "NH-Nodes/Utils/Pipe"
|
||||
|
||||
def pack(self, **kwargs):
|
||||
# Đóng gói tất cả các giá trị đầu vào thành một list
|
||||
# kwargs sẽ là một dictionary chứa tất cả các input đã được kết nối
|
||||
pipe = list(kwargs.values())
|
||||
# Đóng gói theo số ô: package[i] là input_i, ô chưa nối là None.
|
||||
# (list(kwargs.values()) chỉ giữ ô đã nối nên index bị lệch khi bỏ trống input_0.)
|
||||
pipe = [kwargs.get(f"input_{i}") for i in range(_PIPE_SLOTS)]
|
||||
return (pipe,)
|
||||
|
||||
class NH_UnpackUniversal:
|
||||
@@ -29,14 +32,14 @@ class NH_UnpackUniversal:
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"index": ("INT", {"default": 0, "min": 0, "max": 9}),
|
||||
"index": ("INT", {"default": 0, "min": 0, "max": _PIPE_SLOTS - 1}),
|
||||
},
|
||||
"optional": {
|
||||
"package": NH_UNIVERSAL_PIPE,
|
||||
"pipe": NH_UNIVERSAL_PIPE,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
# Đầu ra là kiểu dữ liệu bất kỳ
|
||||
RETURN_TYPES = ("*",)
|
||||
RETURN_NAMES = ("item",)
|
||||
@@ -48,8 +51,8 @@ class NH_UnpackUniversal:
|
||||
if pipe_data is None:
|
||||
return (None,)
|
||||
|
||||
# Kiểm tra xem index có hợp lệ không
|
||||
if index < len(pipe_data):
|
||||
# index = số ô input của Pack; ngoài phạm vi (kể cả âm) thì trả None
|
||||
if 0 <= index < len(pipe_data):
|
||||
return (pipe_data[index],)
|
||||
else:
|
||||
# Trả về None nếu index nằm ngoài phạm vi
|
||||
|
||||
@@ -22,6 +22,9 @@ class NH_UniversalSliderBuilder:
|
||||
FUNCTION = "build_sliders"
|
||||
CATEGORY = "NH-Nodes/Utils"
|
||||
OUTPUT_NODE = True # Đánh dấu đây là node có đầu ra động
|
||||
# Node này không có input/output thật nên không dùng được. DEPRECATED ẩn nó khỏi search,
|
||||
# còn workflow cũ đã lưu vẫn load được. Dùng Multi-Slider (FLOAT/INT) thay thế.
|
||||
DEPRECATED = True
|
||||
|
||||
def build_sliders(self, config, **kwargs):
|
||||
cls = self.__class__
|
||||
|
||||
@@ -7,6 +7,7 @@ use and explains how to install the right wheel.
|
||||
|
||||
import importlib.util
|
||||
import os
|
||||
import weakref
|
||||
|
||||
INSTALL_HINT = (
|
||||
"VTON Ultimate Processor (NH) needs ONNX Runtime, which is not installed in ComfyUI's Python.\n"
|
||||
@@ -16,6 +17,10 @@ INSTALL_HINT = (
|
||||
)
|
||||
|
||||
_warned_cpu_fallback = False
|
||||
# Where the sessions really run, as ONNX Runtime reports it (session.get_providers()[0]). The device widget only
|
||||
# says what was asked for: a CUDA/cuDNN mismatch, or the CPU-only wheel, silently serves a "cuda" request from the CPU.
|
||||
_cpu_fallback_seen = False # a cuda request was served from the CPU at least once in this process (never reset)
|
||||
_cuda_sessions = weakref.WeakSet() # sessions served by the GPU; weak, so a session still dies with its last reference
|
||||
|
||||
|
||||
def ensure_onnxruntime():
|
||||
@@ -28,6 +33,28 @@ def ensure_onnxruntime():
|
||||
raise RuntimeError(f"{INSTALL_HINT}\nThe installed onnxruntime failed to load: {exc}") from exc
|
||||
|
||||
|
||||
def cuda_provider_available():
|
||||
"""True when the installed onnxruntime can use the GPU (the onnxruntime-gpu wheel)."""
|
||||
import onnxruntime as ort
|
||||
|
||||
return "CUDAExecutionProvider" in ort.get_available_providers()
|
||||
|
||||
|
||||
def cpu_fallback_seen():
|
||||
"""True once a session requested on cuda was served from the CPU in this process.
|
||||
|
||||
cuda_provider_available() only says the GPU wheel is installed; a CUDA/cuDNN mismatch still makes
|
||||
ONNX Runtime fall back to the CPU. From then on, asking ComfyUI to free VRAM cannot help these sessions.
|
||||
"""
|
||||
return _cpu_fallback_seen
|
||||
|
||||
|
||||
def cuda_sessions_alive():
|
||||
"""True while a session made by create_session is served by the GPU, i.e. holds VRAM that ComfyUI's
|
||||
model management cannot see or free. Sessions that run on the CPU never count."""
|
||||
return len(_cuda_sessions) > 0
|
||||
|
||||
|
||||
def providers_for(device):
|
||||
# CPU stays in the list so a CUDA request degrades instead of failing outright.
|
||||
if device == "cuda":
|
||||
@@ -36,19 +63,23 @@ def providers_for(device):
|
||||
|
||||
|
||||
def create_session(model_path, device, sess_options=None):
|
||||
global _warned_cpu_fallback
|
||||
global _warned_cpu_fallback, _cpu_fallback_seen
|
||||
import onnxruntime as ort
|
||||
|
||||
session = ort.InferenceSession(model_path, sess_options=sess_options, providers=providers_for(device))
|
||||
active = session.get_providers()
|
||||
if device == "cuda" and (not active or active[0] != "CUDAExecutionProvider") and not _warned_cpu_fallback:
|
||||
_warned_cpu_fallback = True
|
||||
if "CUDAExecutionProvider" in ort.get_available_providers():
|
||||
advice = "The CUDA provider is installed but could not start: check that a GPU is visible and that its CUDA/cuDNN versions match onnxruntime-gpu."
|
||||
else:
|
||||
advice = "Install onnxruntime-gpu built for your CUDA version to use the GPU."
|
||||
print(
|
||||
f"[NH-Nodes] VTON: device=cuda was requested but ONNX Runtime runs "
|
||||
f"{os.path.basename(model_path)} on {active[0] if active else 'an unknown provider'}. {advice}"
|
||||
)
|
||||
if active and active[0] == "CUDAExecutionProvider":
|
||||
_cuda_sessions.add(session)
|
||||
elif device == "cuda":
|
||||
_cpu_fallback_seen = True
|
||||
if not _warned_cpu_fallback:
|
||||
_warned_cpu_fallback = True
|
||||
if "CUDAExecutionProvider" in ort.get_available_providers():
|
||||
advice = "The CUDA provider is installed but could not start: check that a GPU is visible and that its CUDA/cuDNN versions match onnxruntime-gpu."
|
||||
else:
|
||||
advice = "Install onnxruntime-gpu built for your CUDA version to use the GPU."
|
||||
print(
|
||||
f"[NH-Nodes] VTON: device=cuda was requested but ONNX Runtime runs "
|
||||
f"{os.path.basename(model_path)} on {active[0] if active else 'an unknown provider'}. {advice}"
|
||||
)
|
||||
return session
|
||||
|
||||
@@ -1,10 +1,13 @@
|
||||
# VTON Ultimate Processor (NH). Built on ComfyUI-Vton-Mask's mask generator and FitDiT
|
||||
# preprocessing: CC BY-NC-SA 4.0 (non-commercial). See THIRD_PARTY_NOTICES.md.
|
||||
import gc
|
||||
import traceback
|
||||
import torch
|
||||
import numpy as np
|
||||
from PIL import Image, ImageFilter
|
||||
import os
|
||||
import folder_paths
|
||||
from ...core.model_files import download_pinned, make_room_on_gpu, register_model_folder
|
||||
from ...core.utils import resize_image, tensor_to_pil, pil_to_tensor, pil_to_mask
|
||||
|
||||
# onnxruntime (via preprocess/) and scikit-image (via utils_mask) are imported
|
||||
@@ -12,9 +15,86 @@ from ...core.utils import resize_image, tensor_to_pil, pil_to_tensor, pil_to_mas
|
||||
# this node needs, and loading them at ComfyUI start-up makes the pack import
|
||||
# heavier and breaks process memory snapshots (Modal restores segfault).
|
||||
|
||||
# Default location: where the models are downloaded and where existing installs already keep them.
|
||||
MODELS_DIR = os.path.join(folder_paths.models_dir, "ComfyUI-Vton-Mask")
|
||||
LOADED_VTON_MODELS = None
|
||||
|
||||
# extra_model_paths.yaml key. Each location holds dwpose/ and humanparsing/, like MODELS_DIR.
|
||||
VTON_FOLDER_NAME = "vton_mask"
|
||||
VTON_FOLDER_EXTENSIONS = (".onnx",) # what ComfyUI lists for VTON_FOLDER_NAME: the models, not the HF cache metadata
|
||||
VTON_REPO_ID = "kg-09/kg-vton-mask"
|
||||
# The repo is third-party: pin it so a moving "main" cannot change what gets downloaded.
|
||||
VTON_REVISION = "4648778286cac2dca9aab8da167ea5257462487d"
|
||||
VTON_MODEL_FILES = (
|
||||
"dwpose/yolox_l.onnx", "dwpose/dw-ll_ucoco_384.onnx",
|
||||
"humanparsing/parsing_atr.onnx", "humanparsing/parsing_lip.onnx",
|
||||
)
|
||||
_VTON_DOWNLOAD_PATTERNS = ("dwpose/*.onnx", "humanparsing/*.onnx") # skips the repo's other files
|
||||
# Measured with the four CUDA sessions (RTX 5090): ~2.9 GB held after one run, ~0.9 GB of it weights
|
||||
# and the rest ONNX Runtime's arena. Used to ask ComfyUI for that much room (plus the reserve it keeps
|
||||
# for itself, see make_room_on_gpu) before they are created.
|
||||
_VTON_VRAM_ESTIMATE = 3 * 1024 ** 3
|
||||
|
||||
register_model_folder(VTON_FOLDER_NAME, MODELS_DIR, VTON_FOLDER_EXTENSIONS) # make the folder name known to ComfyUI at load time
|
||||
|
||||
|
||||
def _missing_vton_files(root):
|
||||
return [name for name in VTON_MODEL_FILES if not os.path.isfile(os.path.join(root, *name.split("/")))]
|
||||
|
||||
|
||||
def resolve_vton_models_dir(download=True):
|
||||
"""Folder holding all four VTON .onnx files, downloading them into MODELS_DIR if none does.
|
||||
|
||||
Searched: MODELS_DIR (models/ComfyUI-Vton-Mask) and every `vton_mask` location from
|
||||
extra_model_paths.yaml. A folder missing any of the four files, as an interrupted first
|
||||
download leaves it, does not count: the next call finishes the download.
|
||||
"""
|
||||
for root in register_model_folder(VTON_FOLDER_NAME, MODELS_DIR, VTON_FOLDER_EXTENSIONS):
|
||||
if not _missing_vton_files(root):
|
||||
return root
|
||||
if not download:
|
||||
return None
|
||||
download_pinned(VTON_REPO_ID, VTON_REVISION, MODELS_DIR, _VTON_DOWNLOAD_PATTERNS)
|
||||
missing = _missing_vton_files(MODELS_DIR)
|
||||
if missing:
|
||||
raise RuntimeError(f"[NH-Nodes] VTON model files missing after download in {MODELS_DIR}: {', '.join(missing)}")
|
||||
return MODELS_DIR
|
||||
|
||||
|
||||
def release_vton_models():
|
||||
"""Drop the cached ONNX Runtime sessions, which frees the GPU memory they hold.
|
||||
|
||||
ONNX Runtime allocates outside PyTorch, so ComfyUI's model management cannot see or unload
|
||||
these sessions: their VRAM only comes back when the last reference to them goes away.
|
||||
Sessions that ran on the CPU are dropped the same way, but hold no VRAM: torch's cache is left alone.
|
||||
"""
|
||||
global LOADED_VTON_MODELS
|
||||
from .preprocess.ort_session import cuda_sessions_alive
|
||||
|
||||
on_gpu = cuda_sessions_alive() # asked before the sessions go: afterwards nothing is left to ask about
|
||||
LOADED_VTON_MODELS = None
|
||||
gc.collect()
|
||||
if on_gpu and torch.cuda.is_available():
|
||||
try:
|
||||
torch.cuda.empty_cache()
|
||||
except Exception as exc: # runs in a finally: after a CUDA error this must not replace the run's own error
|
||||
print(f"[NH-Nodes] VTON: released the models but torch.cuda.empty_cache() failed: {exc}")
|
||||
|
||||
|
||||
def _should_release(device, keep_model_loaded):
|
||||
"""True when a run on `device` must hand its cached sessions back once it is over.
|
||||
|
||||
That is the case for sessions ONNX Runtime serves from the GPU while the user has not asked to
|
||||
keep them. What counts is where they really run, not the device widget: a cuda request that
|
||||
fell back to the CPU (CUDA/cuDNN mismatch, CPU-only wheel) holds plain RAM, no VRAM, so it keeps
|
||||
its cache like device=cpu does: dropping it would only reload four models for every item of a list.
|
||||
Called after the run, once the sessions exist.
|
||||
"""
|
||||
from .preprocess.ort_session import cuda_sessions_alive
|
||||
|
||||
return device == "cuda" and not keep_model_loaded and cuda_sessions_alive()
|
||||
|
||||
|
||||
LIP_PALETTE = [
|
||||
0,0,0, 128,0,0, 255,0,0, 0,85,0, 170,0,51, 255,85,0, 0,0,85, 0,119,221, 85,85,0, 0,85,85,
|
||||
85,51,0, 52,86,128, 0,128,0, 0,0,255, 51,170,221, 0,255,255, 85,255,170, 170,255,85, 255,255,0, 255,170,85
|
||||
@@ -38,6 +118,17 @@ class NH_VTonUltimateProcessor:
|
||||
"offset_bottom": ("INT", {"default": 0, "min": -200, "max": 200, "step": 1}),
|
||||
"offset_left": ("INT", {"default": 0, "min": -200, "max": 200, "step": 1}),
|
||||
"offset_right": ("INT", {"default": 0, "min": -200, "max": 200, "step": 1}),
|
||||
"keep_model_loaded": ("BOOLEAN", {"default": False, "tooltip": (
|
||||
"Off (default): after each run on cuda the four ONNX models (about 3 GB of VRAM) are released, "
|
||||
"so whatever runs next gets the whole GPU; the next run reloads them. "
|
||||
"On: keep them in VRAM for faster repeated runs. With a list of images ComfyUI calls this node "
|
||||
"once per item (a batch tensor is a single call), so keeping the models loaded is much faster "
|
||||
"for lists and repeated runs. "
|
||||
"ComfyUI cannot see or free that memory, so Free memory does nothing for it "
|
||||
"and a large model run afterwards may run out of VRAM. "
|
||||
"Ignored when the models run on the CPU (device cpu, or cuda without a working ONNX Runtime GPU "
|
||||
"build): they always stay cached."
|
||||
)}),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -50,30 +141,79 @@ class NH_VTonUltimateProcessor:
|
||||
global LOADED_VTON_MODELS
|
||||
if LOADED_VTON_MODELS is not None and LOADED_VTON_MODELS.get("device") == device:
|
||||
return LOADED_VTON_MODELS
|
||||
from .preprocess.ort_session import ensure_onnxruntime
|
||||
from .preprocess.ort_session import cpu_fallback_seen, cuda_provider_available, ensure_onnxruntime
|
||||
ensure_onnxruntime()
|
||||
from .preprocess.dwpose import DWposeDetector
|
||||
from .preprocess.humanparsing.run_parsing import Parsing
|
||||
if not os.path.exists(MODELS_DIR):
|
||||
from huggingface_hub import snapshot_download
|
||||
os.makedirs(MODELS_DIR, exist_ok=True)
|
||||
snapshot_download(repo_id="kg-09/kg-vton-mask", local_dir=MODELS_DIR, local_dir_use_symlinks=False)
|
||||
model_root = resolve_vton_models_dir()
|
||||
if LOADED_VTON_MODELS is not None:
|
||||
release_vton_models() # sessions cached for the other device go before the new ones are allocated
|
||||
# Room is only worth making when the sessions can really use the GPU. Not with the CPU-only wheel, and not
|
||||
# once a cuda request was seen served from the CPU in this process (the GPU wheel is installed but CUDA/cuDNN
|
||||
# do not start; see create_session): evicting ComfyUI's models would not help those sessions. The very first
|
||||
# run cannot know yet and asks once.
|
||||
if device == "cuda" and cuda_provider_available() and not cpu_fallback_seen():
|
||||
make_room_on_gpu(_VTON_VRAM_ESTIMATE)
|
||||
LOADED_VTON_MODELS = {
|
||||
'dwprocessor': DWposeDetector(model_root=MODELS_DIR, device=device),
|
||||
'parsing_model': Parsing(model_root=MODELS_DIR, device=device),
|
||||
'dwprocessor': DWposeDetector(model_root=model_root, device=device),
|
||||
'parsing_model': Parsing(model_root=model_root, device=device),
|
||||
'device': device
|
||||
}
|
||||
return LOADED_VTON_MODELS
|
||||
|
||||
def process(self, human_image, category, mask_feathering, cover_shoes, refine_hands, refine_hair, device,
|
||||
offset_top=0, offset_bottom=0, offset_left=0, offset_right=0):
|
||||
def process(self, human_image, category, mask_feathering, cover_shoes, refine_hands, refine_hair, device,
|
||||
offset_top=0, offset_bottom=0, offset_left=0, offset_right=0, keep_model_loaded=False):
|
||||
# Import trước load_model: thiếu scikit-image thì báo lỗi ngay, không phải tải model xong mới lỗi.
|
||||
from ...core.utils_mask import get_mask_location
|
||||
|
||||
if human_image.shape[0] == 0:
|
||||
raise ValueError("VTON Ultimate Processor: human_image is an empty batch")
|
||||
# Only sessions served by the GPU pin VRAM that ComfyUI cannot manage. Asked after the run, because only
|
||||
# then it is known where load_model's sessions ended up: on the CPU (device=cpu, or a cuda request that fell
|
||||
# back to it) the cache stays whatever the flag says.
|
||||
try:
|
||||
return self._process_batch(
|
||||
human_image, get_mask_location, category, mask_feathering, cover_shoes, refine_hands, refine_hair,
|
||||
device, offset_top, offset_bottom, offset_left, offset_right,
|
||||
)
|
||||
except BaseException as exc:
|
||||
if _should_release(device, keep_model_loaded):
|
||||
# The traceback keeps the failed frames alive and they reference the sessions: clear them,
|
||||
# or the release below frees nothing until the exception has been handled.
|
||||
traceback.clear_frames(exc.__traceback__)
|
||||
raise
|
||||
finally:
|
||||
if _should_release(device, keep_model_loaded):
|
||||
release_vton_models()
|
||||
|
||||
def _process_batch(self, human_image, get_mask_location, category, mask_feathering, cover_shoes, refine_hands,
|
||||
refine_hair, device, offset_top, offset_bottom, offset_left, offset_right):
|
||||
# The models are only referenced from this frame and _process_image's, so release_vton_models() can free them.
|
||||
vton_model = self.load_model(device)
|
||||
human_img_pil = tensor_to_pil(human_image)
|
||||
# Mỗi ảnh trong batch chạy độc lập với cùng setting, rồi nối các output theo chiều batch.
|
||||
results = [
|
||||
self._process_image(
|
||||
vton_model, get_mask_location, human_image[index], index, category, mask_feathering, cover_shoes,
|
||||
refine_hands, refine_hair, offset_top, offset_bottom, offset_left, offset_right,
|
||||
)
|
||||
for index in range(human_image.shape[0])
|
||||
]
|
||||
return tuple(torch.cat(outputs, dim=0) for outputs in zip(*results))
|
||||
|
||||
def _process_image(self, vton_model, get_mask_location, human_image, index, category, mask_feathering, cover_shoes,
|
||||
refine_hands, refine_hair, offset_top, offset_bottom, offset_left, offset_right):
|
||||
"""Xử lý đúng một ảnh [H,W,C]; `index` là vị trí trong batch (chỉ dùng cho thông báo lỗi)."""
|
||||
# DWPose và parsing cần RGB 3 kênh: ảnh xám làm hỏng bước đảo BGR, RGBA đưa 4 kênh vào DWPose (alpha bị bỏ).
|
||||
human_img_pil = tensor_to_pil(human_image).convert("RGB")
|
||||
human_img_resized = resize_image(human_img_pil)
|
||||
|
||||
pose_img_np, _, _, candidate = vton_model['dwprocessor'](np.array(human_img_resized)[:,:,::-1])
|
||||
# DWPose không trả mảng rỗng khi không thấy người: nó vẫn chạy pose trên cả ảnh và đánh dấu mọi keypoint là -1.
|
||||
if len(candidate) == 0 or not np.any(candidate[0] != -1):
|
||||
raise ValueError(
|
||||
f"VTON Ultimate Processor: DWPose found no person in the image at batch index {index}. "
|
||||
"Use a photo where the person is clearly visible."
|
||||
)
|
||||
pose_image_pil = Image.fromarray(pose_img_np[:,:,::-1])
|
||||
parsing_map_pil, parsing_map_colored_pil = vton_model['parsing_model'](human_img_resized)
|
||||
|
||||
@@ -158,7 +298,8 @@ class NH_VTonUltimateProcessor:
|
||||
masked_image = Image.composite(final_mask_for_composite.convert("L"), human_img_pil, final_mask_for_composite)
|
||||
|
||||
pose_image_final = pose_image_pil.resize(human_img_pil.size, Image.LANCZOS)
|
||||
parsing_map_colored_pil = Image.fromarray(parsing_map_full_np.astype(np.uint8), 'P')
|
||||
# No fromarray(mode="P") (deprecated in Pillow 11.3): the uint8 map is read as "L", and putpalette turns it into "P".
|
||||
parsing_map_colored_pil = Image.fromarray(parsing_map_full_np.astype(np.uint8))
|
||||
parsing_map_colored_pil.putpalette(LIP_PALETTE)
|
||||
|
||||
parsing_map_raw_tensor = torch.from_numpy(parsing_map_full_np.astype(np.float32) / 19.0).unsqueeze(0)
|
||||
|
||||
+1
-1
@@ -1,7 +1,7 @@
|
||||
[project]
|
||||
name = "nh-nodes"
|
||||
description = "ComfyUI custom nodes for image workflows, tiling, masks, smart resize, logic routing, prompt and text tools, batch automation, garment/portrait segmentation, and VTON preprocessing."
|
||||
version = "1.4.3"
|
||||
version = "1.6.1"
|
||||
license = { file = "LICENSE" }
|
||||
dependencies = [
|
||||
"numpy",
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,198 @@
|
||||
"""Shared helpers: safe arithmetic evaluation, tensor/PIL conversion, hex colours."""
|
||||
|
||||
import importlib
|
||||
import math
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
safe_eval = importlib.import_module("NH-Nodes.core.safe_eval")
|
||||
utils = importlib.import_module("NH-Nodes.core.utils")
|
||||
|
||||
evaluate = safe_eval.evaluate
|
||||
ExpressionError = safe_eval.ExpressionError
|
||||
|
||||
|
||||
# --- safe_eval ---
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"expression,expected",
|
||||
[
|
||||
("a + b * 2", 7.0),
|
||||
("(a + b) / 2", 2.0),
|
||||
("a ** 2", 1.0),
|
||||
("-a + +b", 2.0),
|
||||
("b // 2", 1.0),
|
||||
("b % 2", 1.0),
|
||||
("clamp(10, 0, 5)", 5.0),
|
||||
("round(3.14159, 2)", 3.14),
|
||||
("sqrt(16) + floor(1.7) + ceil(1.2) + abs(-1) + min(a, b) + max(a, b)", 4 + 1 + 2 + 1 + 1 + 3),
|
||||
],
|
||||
)
|
||||
def test_math_expressions(expression, expected):
|
||||
assert evaluate(expression, {"a": 1, "b": 3}, safe_eval.MATH_FUNCTIONS) == pytest.approx(expected)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("expression", ["9**9**9", "10**400", "(10**100)**100**100", "1e308 * 10"])
|
||||
def test_runaway_numbers_fail_fast_instead_of_freezing(expression):
|
||||
with pytest.raises(ExpressionError):
|
||||
evaluate(expression)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"expression",
|
||||
[
|
||||
"'a' * 10**10", # string constants are not numbers
|
||||
"__import__('os')", # calls outside the whitelist
|
||||
"a.real", # attribute access
|
||||
"[1, 2][0]", # containers / subscripts
|
||||
"lambda: 1",
|
||||
"(-8) ** 0.5", # would be a complex number
|
||||
"sqrt(-1)",
|
||||
"1e308 * 10 - 1e308 * 10", # inf - inf = nan
|
||||
"unknown + 1",
|
||||
"",
|
||||
"1 +",
|
||||
],
|
||||
)
|
||||
def test_rejected_expressions(expression):
|
||||
with pytest.raises(ExpressionError):
|
||||
evaluate(expression, {"a": 1}, safe_eval.MATH_FUNCTIONS)
|
||||
|
||||
|
||||
def test_division_by_zero_modes():
|
||||
warnings = []
|
||||
assert evaluate("a / 0 + 1", {"a": 5}, on_warning=warnings.append) == 1.0
|
||||
assert evaluate("a // 0", {"a": 5}) == 0.0
|
||||
assert evaluate("a % 0", {"a": 5}) == 0.0
|
||||
assert warnings == ["division by zero, using 0"]
|
||||
with pytest.raises(ExpressionError, match="Division by zero"):
|
||||
evaluate("a / 0", {"a": 5}, zero_division="error")
|
||||
|
||||
|
||||
def test_case_insensitive_variables():
|
||||
assert evaluate("BG_W / 2 - w1 / 2", {"bg_w": 100, "w1": 20}, case_insensitive=True) == 40.0
|
||||
with pytest.raises(ExpressionError):
|
||||
evaluate("BG_W", {"bg_w": 100})
|
||||
|
||||
|
||||
def test_non_numeric_variables_are_rejected():
|
||||
with pytest.raises(ExpressionError):
|
||||
evaluate("a", {"a": "5"})
|
||||
with pytest.raises(ExpressionError):
|
||||
evaluate("a", {"a": None})
|
||||
with pytest.raises(ExpressionError):
|
||||
evaluate("a", {"a": torch.ones(2)}) # not a scalar
|
||||
|
||||
|
||||
def test_bools_numpy_and_torch_scalars_count_as_numbers():
|
||||
import numpy as np
|
||||
|
||||
assert evaluate("True + 1") == 2.0
|
||||
assert evaluate("a + b", {"a": np.float32(2), "b": np.int64(3)}) == 5.0
|
||||
assert evaluate("a * 2", {"a": torch.tensor(2.5)}) == 5.0
|
||||
assert evaluate("a + 1", {"a": np.bool_(True)}) == 2.0
|
||||
|
||||
|
||||
def test_absurd_nesting_is_an_expression_error_not_a_crash():
|
||||
with pytest.raises(ExpressionError):
|
||||
evaluate("-" * 100000 + "1")
|
||||
with pytest.raises(ExpressionError):
|
||||
evaluate("(" * 5000 + "1" + ")" * 5000)
|
||||
|
||||
|
||||
def test_result_is_always_a_finite_float():
|
||||
result = evaluate("2 ** 10")
|
||||
assert isinstance(result, float) and math.isfinite(result) and result == 1024.0
|
||||
|
||||
|
||||
# --- tensor_to_pil ---
|
||||
|
||||
def test_tensor_to_pil_clips_instead_of_wrapping():
|
||||
image = torch.tensor([[[1.2, -0.1, 0.5]]]) # 1x1 RGB with out-of-range values
|
||||
pixel = utils.tensor_to_pil(image).getpixel((0, 0))
|
||||
assert pixel == (255, 0, 127)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("channels,mode", [(1, "L"), (3, "RGB"), (4, "RGBA")])
|
||||
def test_tensor_to_pil_channel_modes(channels, mode):
|
||||
image = torch.rand((1, 8, 6, channels))
|
||||
pil = utils.tensor_to_pil(image)
|
||||
assert pil.mode == mode and pil.size == (6, 8)
|
||||
|
||||
|
||||
def test_tensor_to_pil_rejects_batches_with_clear_message():
|
||||
with pytest.raises(ValueError, match="batch of 2"):
|
||||
utils.tensor_to_pil(torch.rand((2, 4, 4, 3)))
|
||||
|
||||
|
||||
def _flat(channels, level=0.5):
|
||||
"""1x1 image whose channel k holds (10 * (k + 1) + level) / 255, i.e. the byte value 10 * (k + 1)."""
|
||||
return torch.tensor([[[(10 * (k + 1) + level) / 255.0 for k in range(channels)]]])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("channels,expected", [(1, (10, 10, 10)), (3, (10, 20, 30)), (4, (10, 20, 30))])
|
||||
def test_tensor_to_pil_rgb_always_gives_rgb(channels, expected):
|
||||
pil = utils.tensor_to_pil(_flat(channels), rgb=True)
|
||||
assert pil.mode == "RGB" and pil.getpixel((0, 0)) == expected # grey is expanded, alpha is dropped
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"channels,mode,expected",
|
||||
[(1, "RGB", (10, 10, 10)), (3, "RGB", (10, 20, 30)), (4, "RGBA", (10, 20, 30, 40))],
|
||||
)
|
||||
def test_tensor_to_pil_keep_alpha_keeps_a_fourth_channel_only(channels, mode, expected):
|
||||
pil = utils.tensor_to_pil(_flat(channels), rgb=True, keep_alpha=True)
|
||||
assert pil.mode == mode and pil.getpixel((0, 0)) == expected
|
||||
|
||||
|
||||
def test_tensor_to_pil_two_channel_image_becomes_rgb_whatever_the_flags():
|
||||
for keep_alpha in (False, True):
|
||||
assert utils.tensor_to_pil(_flat(2), rgb=True, keep_alpha=keep_alpha).mode == "RGB"
|
||||
|
||||
|
||||
def test_tensor_to_pil_keep_alpha_alone_changes_nothing():
|
||||
for channels, mode in ((1, "L"), (3, "RGB"), (4, "RGBA")):
|
||||
assert utils.tensor_to_pil(_flat(channels), keep_alpha=True).mode == mode
|
||||
|
||||
|
||||
def test_tensor_to_pil_flags_are_keyword_only():
|
||||
with pytest.raises(TypeError):
|
||||
utils.tensor_to_pil(_flat(3), True)
|
||||
|
||||
|
||||
def test_tensor_to_pil_accepts_a_single_image_batch_with_the_flags():
|
||||
pil = utils.tensor_to_pil(torch.rand((1, 5, 4, 4)), rgb=True, keep_alpha=True)
|
||||
assert pil.mode == "RGBA" and pil.size == (4, 5)
|
||||
|
||||
|
||||
# --- pil_to_tensor ---
|
||||
|
||||
def test_pil_to_tensor_gives_a_batch_of_one_float32_image():
|
||||
from PIL import Image
|
||||
|
||||
pil = Image.new("RGB", (4, 3), (255, 0, 51))
|
||||
tensor = utils.pil_to_tensor(pil)
|
||||
assert tensor.shape == (1, 3, 4, 3) and tensor.dtype == torch.float32
|
||||
assert torch.allclose(tensor[0, 0, 0], torch.tensor([1.0, 0.0, 51 / 255]))
|
||||
|
||||
|
||||
# --- parse_hex_color ---
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"text,expected",
|
||||
[
|
||||
("#FF8000", (255, 128, 0, 255)),
|
||||
("ff8000", (255, 128, 0, 255)),
|
||||
("#F80", (255, 136, 0, 255)),
|
||||
("#11223344", (17, 34, 51, 68)),
|
||||
(" #000000 ", (0, 0, 0, 255)),
|
||||
],
|
||||
)
|
||||
def test_parse_hex_color(text, expected):
|
||||
assert utils.parse_hex_color(text) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("text", ["", "#12", "#GGGGGG", "red", None])
|
||||
def test_parse_hex_color_invalid_returns_default(text):
|
||||
assert utils.parse_hex_color(text, default=(1, 2, 3, 4)) == (1, 2, 3, 4)
|
||||
@@ -0,0 +1,482 @@
|
||||
"""Executor-level tests: ComfyUI's real PromptExecutor runs small graphs of NH nodes in-process.
|
||||
|
||||
The unit tests in the other files emulate execution.py. These run the real thing (prompt validation, lazy input
|
||||
handling, ExecutionBlocker propagation, the output cache) with a stub in place of the web server, on the CPU, with no
|
||||
GPU, model or network. ``Stub...`` nodes stand in for whatever feeds and consumes the NH nodes: they log what ran and
|
||||
what reached the output nodes.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import copy
|
||||
import gc
|
||||
import importlib
|
||||
import os
|
||||
import weakref
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
import comfy.cli_args
|
||||
import comfy.sd
|
||||
import execution
|
||||
import nodes as comfy_nodes
|
||||
|
||||
logic_nodes = importlib.import_module("NH-Nodes.nodes.logic.logic_nodes")
|
||||
math_nodes = importlib.import_module("NH-Nodes.nodes.logic.math_nodes")
|
||||
folder_loaders = importlib.import_module("NH-Nodes.nodes.image.folder_loaders")
|
||||
loader_nodes = importlib.import_module("NH-Nodes.nodes.loaders.loader_index_nodes")
|
||||
|
||||
# The node only reads the names in the source path, never the folder, so the tests use one that does not exist.
|
||||
SOURCES = Path(os.path.abspath(os.sep)) / "nh_nodes_test_source_photos"
|
||||
|
||||
LOG = SimpleNamespace(ran=[], sunk={}) # reset for every test by the registered_nodes fixture
|
||||
|
||||
|
||||
def _describe(value):
|
||||
"""What the sink keeps of a value: never the value itself, so a sink cannot keep a model alive."""
|
||||
if isinstance(value, torch.Tensor):
|
||||
return (tuple(value.shape), round(float(value.mean()), 3))
|
||||
if value is None or isinstance(value, (bool, int, float, str)):
|
||||
return value
|
||||
return type(value).__name__
|
||||
|
||||
|
||||
# --- stand-ins for the nodes around the NH nodes ----------------------------------------------------------------
|
||||
|
||||
class StubImage:
|
||||
"""Makes a small constant image and logs its tag when it runs."""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"tag": ("STRING", {"default": "image"}), "value": ("FLOAT", {"default": 0.5})}}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "run"
|
||||
|
||||
def run(self, tag, value):
|
||||
LOG.ran.append(tag)
|
||||
return (torch.full((1, 4, 4, 3), value),)
|
||||
|
||||
|
||||
class StubNothing:
|
||||
"""An upstream node that has no value to give (it returns None)."""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"tag": ("STRING", {"default": "nothing"})}}
|
||||
|
||||
RETURN_TYPES = ("*",)
|
||||
FUNCTION = "run"
|
||||
|
||||
def run(self, tag):
|
||||
LOG.ran.append(tag)
|
||||
return (None,)
|
||||
|
||||
|
||||
class StubEcho:
|
||||
"""Hands an image on: any node that needs the image of another one."""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"image": ("IMAGE",)}}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "run"
|
||||
|
||||
def run(self, image):
|
||||
LOG.ran.append("echo")
|
||||
return (image,)
|
||||
|
||||
|
||||
class StubSink:
|
||||
"""Output node: records what reached it under its name."""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"name": ("STRING", {"default": "sink"}), "value": ("*",)}}
|
||||
|
||||
RETURN_TYPES = ()
|
||||
OUTPUT_NODE = True
|
||||
FUNCTION = "run"
|
||||
|
||||
def run(self, name, value):
|
||||
LOG.sunk[name] = _describe(value)
|
||||
return {}
|
||||
|
||||
|
||||
STUBS = {
|
||||
"StubImage": StubImage,
|
||||
"StubNothing": StubNothing,
|
||||
"StubEcho": StubEcho,
|
||||
"StubSink": StubSink,
|
||||
"NH_IfElse": logic_nodes.NH_IfElse,
|
||||
"NH_SwitchN": logic_nodes.NH_SwitchN,
|
||||
"NH_AnySwitchBoolean": logic_nodes.NH_AnySwitchBoolean,
|
||||
"NH_RandomChoice": math_nodes.NH_RandomChoice,
|
||||
"NH_MathEval": math_nodes.NH_MathEval,
|
||||
"NH_LoadImagesMatching": folder_loaders.NH_LoadImagesMatching,
|
||||
"NH_DiffusionModelIndex": loader_nodes.NH_DiffusionModelIndex,
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def registered_nodes(monkeypatch):
|
||||
for name, cls in STUBS.items():
|
||||
monkeypatch.setitem(comfy_nodes.NODE_CLASS_MAPPINGS, name, cls)
|
||||
LOG.ran.clear()
|
||||
LOG.sunk.clear()
|
||||
|
||||
|
||||
# --- running a graph ----------------------------------------------------------------------------------------------
|
||||
|
||||
class _Server:
|
||||
"""The part of PromptServer that PromptExecutor uses: it sends events and has a client id."""
|
||||
|
||||
client_id = None
|
||||
last_node_id = None
|
||||
|
||||
def __init__(self):
|
||||
self.messages = []
|
||||
|
||||
def send_sync(self, event, data, sid=None):
|
||||
self.messages.append((event, data))
|
||||
|
||||
|
||||
def new_executor(cache_type=execution.CacheType.CLASSIC, lru=0, ram=0.0, ram_inactive=0.0):
|
||||
"""An executor with one of ComfyUI's caches (ram and ram_inactive are the RAM-pressure headrooms, in GB)."""
|
||||
cache_args = {"lru": lru, "ram": ram, "ram_inactive": ram_inactive}
|
||||
return execution.PromptExecutor(_Server(), cache_type=cache_type, cache_args=cache_args)
|
||||
|
||||
|
||||
def run(prompt, executor=None):
|
||||
"""Validate `prompt` the way the server does, execute it, and report what happened."""
|
||||
prompt = copy.deepcopy(prompt) # a queued prompt is new data every time; the executor writes into it
|
||||
valid, error, outputs, node_errors = asyncio.run(execution.validate_prompt("test", prompt, None))
|
||||
assert valid and not node_errors, (error, node_errors)
|
||||
|
||||
executor = executor or new_executor()
|
||||
executor.server.messages.clear()
|
||||
LOG.ran.clear()
|
||||
LOG.sunk.clear()
|
||||
executor.execute(prompt, "test", {"client_id": "test"}, outputs)
|
||||
|
||||
errors = [data for event, data in executor.server.messages if event == "execution_error"]
|
||||
return SimpleNamespace(
|
||||
ran=list(LOG.ran),
|
||||
sunk=dict(LOG.sunk),
|
||||
success=executor.success,
|
||||
blocked=[data["exception_message"] for data in errors if data["exception_type"] == "ExecutionBlocked"],
|
||||
failed=[data["exception_message"] for data in errors if data["exception_type"] != "ExecutionBlocked"],
|
||||
)
|
||||
|
||||
|
||||
def sink(name, link):
|
||||
return {"class_type": "StubSink", "inputs": {"name": name, "value": link}}
|
||||
|
||||
|
||||
def image(tag, value=0.5):
|
||||
return {"class_type": "StubImage", "inputs": {"tag": tag, "value": value}}
|
||||
|
||||
|
||||
def nothing(tag="nothing"):
|
||||
return {"class_type": "StubNothing", "inputs": {"tag": tag}}
|
||||
|
||||
|
||||
def math_eval(expression, **values):
|
||||
return {"class_type": "NH_MathEval", "inputs": {"expression": expression, **values}}
|
||||
|
||||
|
||||
GOOD = math_eval("a + b", a=1.0, b=2.0) # 3.0
|
||||
BAD = math_eval("1 +") # a syntax error: the prompt fails if this node ever runs
|
||||
|
||||
|
||||
# --- 1. lazy branches: the branch that is not taken never runs -------------------------------------------------
|
||||
|
||||
def test_the_invalid_math_node_does_fail_when_it_runs():
|
||||
result = run({"1": BAD, "2": sink("result", ["1", 0])})
|
||||
assert result.success is False
|
||||
assert len(result.failed) == 1 and "Math Eval" in result.failed[0]
|
||||
assert result.sunk == {}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("condition", [True, False])
|
||||
def test_if_else_does_not_run_the_unchosen_branch(condition):
|
||||
taken, dropped = ("if_true", "if_false") if condition else ("if_false", "if_true")
|
||||
prompt = {
|
||||
"1": GOOD,
|
||||
"2": BAD,
|
||||
"3": {"class_type": "NH_IfElse", "inputs": {"condition": condition, taken: ["1", 0], dropped: ["2", 0]}},
|
||||
"4": sink("result", ["3", 0]),
|
||||
"5": sink("condition", ["3", 1]),
|
||||
}
|
||||
result = run(prompt)
|
||||
assert result.success and result.failed == [] and result.blocked == []
|
||||
assert result.sunk == {"result": 3.0, "condition": condition}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("condition", [True, False])
|
||||
def test_any_switch_does_not_run_the_unchosen_input(condition):
|
||||
taken, dropped = ("input_true", "input_false") if condition else ("input_false", "input_true")
|
||||
prompt = {
|
||||
"1": GOOD,
|
||||
"2": BAD,
|
||||
"3": {"class_type": "NH_AnySwitchBoolean", "inputs": {"condition": condition, taken: ["1", 0], dropped: ["2", 0]}},
|
||||
"4": sink("result", ["3", 0]),
|
||||
}
|
||||
result = run(prompt)
|
||||
assert result.success and result.failed == []
|
||||
assert result.sunk == {"result": 3.0}
|
||||
|
||||
|
||||
def test_switch_n_does_not_run_the_inputs_that_were_not_picked():
|
||||
prompt = {
|
||||
"1": BAD,
|
||||
"2": GOOD,
|
||||
"3": BAD,
|
||||
"4": {"class_type": "NH_SwitchN", "inputs": {"index": 1, "input_0": ["1", 0], "input_1": ["2", 0], "input_2": ["3", 0]}},
|
||||
"5": sink("result", ["4", 0]),
|
||||
"6": sink("count", ["4", 1]),
|
||||
}
|
||||
result = run(prompt)
|
||||
assert result.success and result.failed == []
|
||||
assert result.sunk == {"result": 3.0, "count": 3}
|
||||
|
||||
|
||||
def test_random_choice_does_not_run_the_inputs_that_were_not_picked():
|
||||
prompt = {
|
||||
"1": BAD,
|
||||
"2": GOOD,
|
||||
"3": BAD,
|
||||
"4": {"class_type": "NH_RandomChoice", "inputs": {"seed": 5, "weights": "0,1,0", "input_0": ["1", 0], "input_1": ["2", 0], "input_2": ["3", 0]}},
|
||||
"5": sink("result", ["4", 0]),
|
||||
"6": sink("picked", ["4", 1]),
|
||||
}
|
||||
result = run(prompt)
|
||||
assert result.success and result.failed == []
|
||||
assert result.sunk == {"result": 3.0, "picked": 1}
|
||||
|
||||
|
||||
def test_a_connected_input_that_produces_nothing_still_counts_as_connected():
|
||||
# What the DESCRIPTION of Switch N and Random Choice says: the picked slot is the connected one, value or not.
|
||||
switch = {
|
||||
"1": image("other"),
|
||||
"2": nothing("empty"),
|
||||
"3": {"class_type": "NH_SwitchN", "inputs": {"index": 1, "input_0": ["1", 0], "input_1": ["2", 0]}},
|
||||
"4": sink("result", ["3", 0]),
|
||||
"5": sink("count", ["3", 1]),
|
||||
}
|
||||
result = run(switch)
|
||||
assert result.success and result.ran == ["empty"] # the other slot did not run
|
||||
assert result.sunk == {"result": None, "count": 2}
|
||||
|
||||
choice = {
|
||||
"1": image("other"),
|
||||
"2": nothing("empty"),
|
||||
"3": {"class_type": "NH_RandomChoice", "inputs": {"seed": 3, "weights": "0,1", "input_0": ["1", 0], "input_1": ["2", 0]}},
|
||||
"4": sink("result", ["3", 0]),
|
||||
"5": sink("picked", ["3", 1]),
|
||||
}
|
||||
result = run(choice)
|
||||
assert result.success and result.ran == ["empty"]
|
||||
assert result.sunk == {"result": None, "picked": 1}
|
||||
|
||||
|
||||
# --- 2. Load Images Matching: a miss, with and without fallback_images ------------------------------------------
|
||||
|
||||
@pytest.fixture
|
||||
def matching_tree(tmp_path):
|
||||
"""<root>/Nu/Anh holds one white 8x8 picture; the source folder is called .../Nu/Anh, so it is a hit."""
|
||||
root = tmp_path / "matching"
|
||||
(root / "Nu" / "Anh").mkdir(parents=True)
|
||||
Image.new("RGB", (8, 8), (255, 255, 255)).save(root / "Nu" / "Anh" / "01.png")
|
||||
return SimpleNamespace(root=root, source=SOURCES / "Nu" / "Anh")
|
||||
|
||||
|
||||
def matching(tree, **inputs):
|
||||
settings = {
|
||||
"source_folder_path": str(tree.source), "matching_root_path": str(tree.root), "match_text": "",
|
||||
"match_mode": "contains", "recursive": True, "index": 0, "step": 1, "repeat_count": 1,
|
||||
"context_folder_names": "Nam\nNu", "source_file_names": "",
|
||||
}
|
||||
settings.update(inputs)
|
||||
return {"class_type": "NH_LoadImagesMatching", "inputs": settings}
|
||||
|
||||
|
||||
# Without match_text the source folder's name ("Anh") is looked for, and it is there: a hit. "Zed" is not: a miss.
|
||||
MISS = {"match_text": "Zed"}
|
||||
|
||||
|
||||
def test_a_hit_loads_the_pictures_and_never_runs_the_fallback(matching_tree):
|
||||
prompt = {
|
||||
"1": matching(matching_tree, fallback_images=["2", 0]),
|
||||
"2": image("fallback", 0.25),
|
||||
"3": sink("images", ["1", 0]),
|
||||
"4": sink("found", ["1", 4]),
|
||||
}
|
||||
result = run(prompt)
|
||||
assert result.success and result.blocked == [] and result.failed == []
|
||||
assert result.sunk == {"images": ((1, 8, 8, 3), 1.0), "found": True}
|
||||
assert result.ran == [] # fallback_images is lazy and nothing asked for it
|
||||
|
||||
|
||||
def test_a_miss_with_fallback_images_delivers_the_fallback(matching_tree):
|
||||
prompt = {
|
||||
"1": matching(matching_tree, fallback_images=["2", 0], **MISS),
|
||||
"2": image("fallback", 0.25),
|
||||
"3": sink("images", ["1", 0]),
|
||||
"4": sink("found", ["1", 4]),
|
||||
"5": sink("matched_text", ["1", 3]),
|
||||
"6": sink("filename", ["1", 1]),
|
||||
}
|
||||
result = run(prompt)
|
||||
assert result.success and result.blocked == [] and result.failed == []
|
||||
assert result.sunk == {
|
||||
"images": ((1, 4, 4, 3), 0.25), # the fallback, as it is
|
||||
"found": False,
|
||||
"matched_text": "Zed",
|
||||
"filename": "",
|
||||
}
|
||||
assert result.ran == ["fallback"]
|
||||
|
||||
|
||||
def test_a_miss_with_a_fallback_that_an_earlier_prompt_cached(matching_tree):
|
||||
executor = new_executor()
|
||||
prompt = {
|
||||
"1": matching(matching_tree, fallback_images=["2", 0], **MISS),
|
||||
"2": image("fallback", 0.25),
|
||||
"3": sink("images", ["1", 0]),
|
||||
}
|
||||
first = run(prompt, executor)
|
||||
assert first.ran == ["fallback"] and first.sunk == {"images": ((1, 4, 4, 3), 0.25)}
|
||||
second = run(prompt, executor) # the miss is never remembered, the fallback comes from the cache
|
||||
assert second.success and second.ran == []
|
||||
assert second.sunk == {"images": ((1, 4, 4, 3), 0.25)}
|
||||
|
||||
|
||||
def test_a_miss_without_fallback_images_blocks_only_its_own_consumers(matching_tree):
|
||||
prompt = {
|
||||
"1": matching(matching_tree, **MISS),
|
||||
"2": {"class_type": "StubEcho", "inputs": {"image": ["1", 0]}},
|
||||
"3": sink("echoed", ["2", 0]),
|
||||
"4": image("independent", 0.75),
|
||||
"5": sink("independent", ["4", 0]),
|
||||
"6": sink("found", ["1", 4]),
|
||||
}
|
||||
result = run(prompt)
|
||||
assert result.success and result.failed == []
|
||||
assert result.ran == ["independent"] # the node that needed the images never ran
|
||||
assert result.sunk == {"independent": ((1, 4, 4, 3), 0.75), "found": False} # nothing reached the echo's sink
|
||||
assert len(result.blocked) == 1 # the reason is reported once
|
||||
assert result.blocked[0].startswith("Execution Blocked: No matching folder found for 'Zed'")
|
||||
|
||||
|
||||
def test_routing_the_held_back_images_through_if_else_does_not_reach_the_fallback(matching_tree):
|
||||
# The reason fallback_images exists: a node with a held-back input is blocked before it can choose, whichever
|
||||
# input it would have picked, so the other branch of the If/Else is never reached.
|
||||
prompt = {
|
||||
"1": matching(matching_tree, **MISS),
|
||||
"2": image("fallback", 0.25),
|
||||
"3": {"class_type": "NH_IfElse", "inputs": {"condition": ["1", 4], "if_true": ["1", 0], "if_false": ["2", 0]}},
|
||||
"4": sink("result", ["3", 0]),
|
||||
}
|
||||
result = run(prompt)
|
||||
assert result.sunk == {} and result.ran == []
|
||||
assert len(result.blocked) == 1
|
||||
|
||||
|
||||
def test_fallback_images_does_the_job_the_if_else_cannot(matching_tree):
|
||||
prompt = {
|
||||
"1": matching(matching_tree, fallback_images=["2", 0], **MISS),
|
||||
"2": image("fallback", 0.25),
|
||||
"3": sink("result", ["1", 0]),
|
||||
}
|
||||
result = run(prompt)
|
||||
assert result.blocked == [] and result.sunk == {"result": ((1, 4, 4, 3), 0.25)}
|
||||
|
||||
|
||||
def test_a_setting_that_cannot_work_fails_even_with_fallback_images(matching_tree):
|
||||
prompt = {
|
||||
"1": matching(matching_tree, matching_root_path=str(matching_tree.root / "missing"), fallback_images=["2", 0]),
|
||||
"2": image("fallback", 0.25),
|
||||
"3": sink("images", ["1", 0]),
|
||||
}
|
||||
result = run(prompt)
|
||||
assert result.success is False and result.ran == []
|
||||
assert len(result.failed) == 1 and "Matching root folder does not exist" in result.failed[0]
|
||||
|
||||
|
||||
# --- 3. Indexed loaders under ComfyUI's own output caches -------------------------------------------------------
|
||||
|
||||
MODEL_FILES = ["a.safetensors", "b.safetensors", "c.safetensors", "d.safetensors"]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def diffusion_models(tmp_path, monkeypatch):
|
||||
"""The indexed diffusion loader with only the file reading faked: counts loads and watches the loaded objects."""
|
||||
loads = SimpleNamespace(count=0, watchers=[])
|
||||
|
||||
class Loaded:
|
||||
"""A plain object, like the ModelPatcher the real loader returns."""
|
||||
|
||||
def load_diffusion_model(path, model_options=None):
|
||||
loads.count += 1
|
||||
model = Loaded()
|
||||
loads.watchers.append(weakref.ref(model))
|
||||
return model
|
||||
|
||||
def full_path(folder, name):
|
||||
path = tmp_path / name
|
||||
if not path.exists(): # written once: the node tells a replaced file by its size and time
|
||||
path.write_bytes(b"model")
|
||||
return str(path)
|
||||
|
||||
monkeypatch.setattr(comfy.sd, "load_diffusion_model", load_diffusion_model)
|
||||
monkeypatch.setattr(loader_nodes.folder_paths, "get_full_path_or_raise", full_path)
|
||||
monkeypatch.setattr(loader_nodes.folder_paths, "get_filename_list", lambda folder: list(MODEL_FILES))
|
||||
loads.alive = lambda: sum(watcher() is not None for watcher in loads.watchers)
|
||||
return loads
|
||||
|
||||
|
||||
def edit_unselected_slot_three_times(executor):
|
||||
"""Queue the same loader prompt with a slot that nothing selected changed each time: the node runs again."""
|
||||
for unselected in MODEL_FILES[1:]:
|
||||
inputs = {
|
||||
"select_mode": "index", "index": 1, "model_count": 2, "weight_dtype": "default",
|
||||
"diffusion_model_1": "a.safetensors", "diffusion_model_2": unselected,
|
||||
}
|
||||
prompt = {"1": {"class_type": "NH_DiffusionModelIndex", "inputs": inputs}, "2": sink("model", ["1", 0])}
|
||||
result = run(prompt, executor)
|
||||
assert result.success and result.sunk == {"model": "Loaded"}
|
||||
gc.collect()
|
||||
|
||||
|
||||
def test_default_cache_the_loader_reuses_its_model_when_an_unselected_slot_changes(diffusion_models, monkeypatch):
|
||||
monkeypatch.setattr(comfy.cli_args.args, "cache_classic", False)
|
||||
executor = new_executor(execution.CacheType.RAM_PRESSURE, ram=0.001, ram_inactive=0.001) # plenty of free RAM
|
||||
edit_unselected_slot_three_times(executor)
|
||||
assert diffusion_models.count == 1 and diffusion_models.alive() == 1
|
||||
|
||||
|
||||
def test_default_cache_memory_pressure_can_free_the_loaded_model(diffusion_models, monkeypatch):
|
||||
monkeypatch.setattr(comfy.cli_args.args, "cache_classic", False)
|
||||
# A headroom no machine has: the RAM-pressure cache is always short of memory and evicts what it holds.
|
||||
executor = new_executor(execution.CacheType.RAM_PRESSURE, ram=1e6, ram_inactive=1e6)
|
||||
edit_unselected_slot_three_times(executor)
|
||||
assert diffusion_models.alive() == 0 # nothing of the node's kept the evicted models in memory
|
||||
|
||||
|
||||
def test_cache_classic_the_loader_holds_the_model_the_classic_cache_dropped(diffusion_models, monkeypatch):
|
||||
monkeypatch.setattr(comfy.cli_args.args, "cache_classic", True) # what --cache-classic sets
|
||||
executor = new_executor(execution.CacheType.CLASSIC)
|
||||
edit_unselected_slot_three_times(executor)
|
||||
assert diffusion_models.count == 1 and diffusion_models.alive() == 1
|
||||
|
||||
|
||||
def test_lru_cache_the_loader_reuses_its_model_like_the_default_cache(diffusion_models, monkeypatch):
|
||||
monkeypatch.setattr(comfy.cli_args.args, "cache_classic", False)
|
||||
executor = new_executor(execution.CacheType.LRU, lru=100)
|
||||
edit_unselected_slot_three_times(executor)
|
||||
assert diffusion_models.count == 1 and diffusion_models.alive() == 1
|
||||
@@ -0,0 +1,280 @@
|
||||
"""Regression tests (2026-10-18 review): the Image Grid Composite size guard and the
|
||||
file-name part of Save Image Path."""
|
||||
|
||||
import importlib
|
||||
import os
|
||||
import types
|
||||
from pathlib import Path
|
||||
|
||||
import psutil
|
||||
import pytest
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
grid = importlib.import_module("NH-Nodes.nodes.image.image_grid_nodes")
|
||||
save_image_path = importlib.import_module("NH-Nodes.nodes.image.save_image_path")
|
||||
utils = importlib.import_module("NH-Nodes.core.utils")
|
||||
|
||||
GIB = 2**30
|
||||
|
||||
|
||||
# --- Image Grid Composite: the size guard follows the free memory, not Pillow's decompression-bomb limit ---
|
||||
|
||||
@pytest.fixture
|
||||
def free_memory(monkeypatch):
|
||||
"""Make psutil report the given available RAM. The tests only exercise the decision: nothing big is allocated."""
|
||||
|
||||
def report(nbytes):
|
||||
monkeypatch.setattr(psutil, "virtual_memory", lambda: types.SimpleNamespace(available=int(nbytes)))
|
||||
|
||||
return report
|
||||
|
||||
|
||||
def _check(width, height, columns=1, rows=1):
|
||||
grid._check_canvas_size(width, height, columns, rows, width // columns, height // rows)
|
||||
|
||||
|
||||
def _composite(images, **overrides):
|
||||
params = dict(
|
||||
columns=2, rows=2, cell_size_mode="manual", cell_width=8, cell_height=8,
|
||||
resize_mode="stretch", gap_x=0, gap_y=0, background_color_hex="#000000",
|
||||
)
|
||||
params.update(overrides)
|
||||
return grid.NH_ImageGridComposite().composite(images, **params)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("pillow_limit", [89_478_485, 1000, None]) # Pillow's default, a low value, "disabled"
|
||||
@pytest.mark.parametrize(
|
||||
"columns,rows,cell",
|
||||
[(40, 1, 1024), (4, 4, 4096), (1, 1, 16384)], # a 40960x1024 filmstrip; 16384x16384 as 4x4 cells and as one cell
|
||||
)
|
||||
def test_grid_guard_accepts_big_canvases_whatever_pillow_limit_other_packs_set(
|
||||
monkeypatch, free_memory, pillow_limit, columns, rows, cell
|
||||
):
|
||||
# Image.MAX_IMAGE_PIXELS is process-wide and only guards Image.open; an in-memory IMAGE never goes through it.
|
||||
monkeypatch.setattr(Image, "MAX_IMAGE_PIXELS", pillow_limit)
|
||||
free_memory(64 * GIB)
|
||||
|
||||
grid._check_canvas_size(columns * cell, rows * cell, columns, rows, cell, cell)
|
||||
|
||||
|
||||
def test_grid_estimate_is_the_peak_of_the_canvas_to_tensor_conversion():
|
||||
# PIL keeps an RGB canvas as 4 bytes per pixel; pil_to_tensor then holds two float32 RGB arrays (astype, / 255.0).
|
||||
pil_rgb, float32_rgb = 4, 3 * 4
|
||||
peak_per_pixel = pil_rgb + 2 * float32_rgb
|
||||
|
||||
assert grid._estimate_canvas_bytes(1, 1) == peak_per_pixel
|
||||
assert grid._estimate_canvas_bytes(40960, 1024) == 40960 * 1024 * peak_per_pixel
|
||||
assert grid._estimate_canvas_bytes(16384, 16384) == 7 * GIB
|
||||
|
||||
|
||||
def test_grid_estimate_covers_what_the_node_keeps_alive():
|
||||
canvas = Image.new("RGB", (64, 48))
|
||||
tensor = utils.pil_to_tensor(canvas)
|
||||
|
||||
held = tensor.numel() * tensor.element_size() + 64 * 48 * 4 # the IMAGE tensor and the PIL canvas
|
||||
assert grid._estimate_canvas_bytes(64, 48) >= held
|
||||
|
||||
|
||||
def test_grid_guard_needs_headroom_on_top_of_the_estimate(free_memory):
|
||||
width = height = 8192
|
||||
need = grid._estimate_canvas_bytes(width, height)
|
||||
|
||||
free_memory(4 * need)
|
||||
_check(width, height)
|
||||
|
||||
for free in (need * 1.05, need, need // 2, 0): # barely enough is not enough: some memory stays free
|
||||
free_memory(free)
|
||||
with pytest.raises(ValueError, match=r"(?i)too large"):
|
||||
_check(width, height)
|
||||
|
||||
|
||||
def test_grid_guard_switches_exactly_at_the_usable_share_of_the_free_memory(free_memory):
|
||||
need = grid._estimate_canvas_bytes(4096, 4096)
|
||||
threshold = need / grid._USABLE_MEMORY_FRACTION
|
||||
|
||||
free_memory(threshold + 2)
|
||||
_check(4096, 4096)
|
||||
|
||||
free_memory(threshold - 2)
|
||||
with pytest.raises(ValueError, match=r"(?i)too large"):
|
||||
_check(4096, 4096)
|
||||
|
||||
|
||||
def test_grid_guard_error_names_the_grid_and_the_memory(free_memory):
|
||||
free_memory(2 * GIB)
|
||||
|
||||
with pytest.raises(ValueError) as error:
|
||||
grid._check_canvas_size(16384, 16384, 4, 4, 4096, 4096)
|
||||
|
||||
message = str(error.value)
|
||||
assert "16384x16384" in message and "4x4 cells of 4096x4096" in message
|
||||
assert "7.0 GiB" in message and "2.0 GiB" in message # what it needs / what is free
|
||||
|
||||
|
||||
def test_grid_node_checks_the_memory_before_it_allocates_the_canvas(monkeypatch, free_memory):
|
||||
free_memory(1024) # 1 KiB free: even this 16x16 canvas (2x2 cells of 8x8) cannot be planned
|
||||
created_sizes = []
|
||||
real_new = Image.new
|
||||
|
||||
def spy(mode, size, *args, **kwargs):
|
||||
created_sizes.append(size)
|
||||
return real_new(mode, size, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(Image, "new", spy)
|
||||
|
||||
with pytest.raises(ValueError, match=r"(?i)too large"):
|
||||
_composite([torch.rand((1, 8, 8, 3))])
|
||||
|
||||
assert (16, 16) not in created_sizes # Pillow itself makes 8x8 images while converting the input
|
||||
|
||||
|
||||
def test_grid_node_ignores_the_pillow_pixel_limit(monkeypatch):
|
||||
monkeypatch.setattr(Image, "MAX_IMAGE_PIXELS", 1000) # the first guard refused every canvas over 2 x 1000 px
|
||||
|
||||
image = _composite([torch.rand((1, 8, 8, 3))], cell_width=32, cell_height=32)[0]
|
||||
|
||||
assert image.shape == (1, 64, 64, 3) # 4096 px
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"columns,rows,cell_width,cell_height,shape",
|
||||
[
|
||||
(3, 1, 16384, 1, (1, 1, 49152, 3)), # 49152 px wide
|
||||
(1, 3, 1, 16384, (1, 49152, 1, 3)), # 49152 px tall
|
||||
],
|
||||
)
|
||||
def test_grid_node_builds_strips_longer_than_the_old_32768_px_cap(columns, rows, cell_width, cell_height, shape):
|
||||
image = _composite(
|
||||
[torch.rand((1, 2, 2, 3))], columns=columns, rows=rows, cell_width=cell_width, cell_height=cell_height
|
||||
)[0]
|
||||
|
||||
assert image.shape == shape
|
||||
|
||||
|
||||
def test_grid_guard_does_nothing_without_psutil(monkeypatch):
|
||||
monkeypatch.setattr(grid, "psutil", None)
|
||||
|
||||
_check(16384, 16384) # no information about the memory: no guard, as before the guard existed
|
||||
|
||||
|
||||
def test_grid_guard_does_nothing_when_psutil_cannot_read_the_memory(monkeypatch):
|
||||
def broken():
|
||||
raise OSError("memory status is not available")
|
||||
|
||||
monkeypatch.setattr(psutil, "virtual_memory", broken)
|
||||
|
||||
_check(16384, 16384)
|
||||
|
||||
|
||||
def test_grid_guard_refuses_a_side_pillow_cannot_create(free_memory):
|
||||
free_memory(1024 * GIB)
|
||||
ceiling = grid._PIL_MAX_SIDE # Image.new raises a bare MemoryError for a width above INT_MAX // 4 - 1
|
||||
|
||||
_check(ceiling, 1)
|
||||
for width, height in ((ceiling + 1, 1), (1, ceiling + 1)):
|
||||
with pytest.raises(ValueError, match=r"(?i)too large.*Pillow"):
|
||||
_check(width, height)
|
||||
|
||||
|
||||
# --- Save Image Path: the file-name part of filename_prefix is checked too ---
|
||||
|
||||
@pytest.fixture
|
||||
def written(monkeypatch):
|
||||
"""Record the files Save Image Path would write instead of writing them, so that a prefix that escapes the
|
||||
target folder cannot leave a stray file on the real disk."""
|
||||
paths = []
|
||||
monkeypatch.setattr(save_image_path, "_save_pil_image", lambda image, path, file_format, dpi: paths.append(path))
|
||||
return paths
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def as_windows(monkeypatch):
|
||||
"""Judge prefixes by the Windows rules even where the tests run on another OS (a no-op on Windows itself)."""
|
||||
monkeypatch.setattr(save_image_path, "_WINDOWS", True, raising=False)
|
||||
|
||||
|
||||
def _save_image_path(tmp_path, prefix, **kwargs):
|
||||
settings = dict(folder_path=str(tmp_path / "out"), create_new_folder=False, new_folder_name="")
|
||||
settings.update(kwargs)
|
||||
return save_image_path.NH_SaveImagePath().save_images(torch.rand((1, 4, 4, 3)), filename_prefix=prefix, **settings)
|
||||
|
||||
|
||||
def _png_files_outside_out(tmp_path):
|
||||
return [p for p in tmp_path.rglob("*.png") if p.relative_to(tmp_path).parts[0] != "out"]
|
||||
|
||||
|
||||
# "D:x" and "x:" are paths on another drive, "C:x" is relative to the current folder of drive C: (os.path.join keeps
|
||||
# it below the target but drops the "C:", so the reported name is wrong and every run overwrites x0001.png), and
|
||||
# "name:stream" writes an NTFS alternate data stream. A colon cannot be part of a Windows file name at all.
|
||||
COLON_PREFIXES = [
|
||||
"D:x", "C:x", "c:x", "a:b", "x:", ":", "12:30", "name:stream",
|
||||
"sub/D:x", "sub\\C:x", "C:sub/name", "D:\\x", "D:/x",
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("as_windows")
|
||||
@pytest.mark.parametrize("number_suffix", [True, False])
|
||||
@pytest.mark.parametrize("prefix", COLON_PREFIXES)
|
||||
def test_windows_prefix_with_a_colon_is_refused_before_anything_is_written(tmp_path, written, prefix, number_suffix):
|
||||
with pytest.raises(ValueError, match="filename_prefix"):
|
||||
_save_image_path(tmp_path, prefix, number_suffix=number_suffix)
|
||||
|
||||
assert written == []
|
||||
assert not (tmp_path / "out").exists() # not even the target folder is created
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"prefix",
|
||||
["/x", "\\x", "//server/share/x", "\\\\server\\share\\x", "\\\\?\\C:\\x", "../x", "..\\x", "sub/../../x"],
|
||||
)
|
||||
def test_rooted_unc_and_parent_prefixes_are_refused_on_every_os(tmp_path, written, prefix):
|
||||
with pytest.raises(ValueError, match="filename_prefix"):
|
||||
_save_image_path(tmp_path, prefix)
|
||||
|
||||
assert written == []
|
||||
assert not (tmp_path / "out").exists()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"prefix,expected",
|
||||
[
|
||||
("name", "name0001.png"),
|
||||
("sub/name", "sub/name0001.png"),
|
||||
("sub\\name", "sub\\name0001.png"),
|
||||
("a/b/name", "a/b/name0001.png"),
|
||||
("my..name", "my..name0001.png"),
|
||||
("sub/", "sub/0001.png"),
|
||||
],
|
||||
)
|
||||
def test_plain_and_subfolder_prefixes_keep_working(tmp_path, prefix, expected):
|
||||
filenames, path, count = _save_image_path(tmp_path, prefix)
|
||||
|
||||
assert (filenames, path, count) == (expected, str(tmp_path / "out"), 1)
|
||||
assert (tmp_path / "out" / Path(expected.replace("\\", "/"))).is_file()
|
||||
assert _png_files_outside_out(tmp_path) == []
|
||||
|
||||
|
||||
def test_plain_prefix_is_numbered_across_runs(tmp_path):
|
||||
first = _save_image_path(tmp_path, "sub/name")[0]
|
||||
second = _save_image_path(tmp_path, "sub\\name")[0]
|
||||
|
||||
assert (first, second) == ("sub/name0001.png", "sub\\name0002.png")
|
||||
assert (tmp_path / "out" / "sub" / "name0002.png").is_file()
|
||||
|
||||
|
||||
@pytest.mark.skipif(os.name == "nt", reason="a colon is not a valid character in a Windows file name")
|
||||
def test_colons_stay_allowed_outside_windows(tmp_path):
|
||||
filenames, _, _ = _save_image_path(tmp_path, "render 12:30")
|
||||
|
||||
assert filenames == "render 12:300001.png"
|
||||
assert (tmp_path / "out" / "render 12:300001.png").is_file()
|
||||
|
||||
|
||||
def test_same_size_rgba_resize_keeps_colour_under_transparent_pixels():
|
||||
picker = importlib.import_module("NH-Nodes.nodes.resolution.smart_resolution_picker")
|
||||
image = torch.zeros((1, 4, 8, 8))
|
||||
image[:, :3] = torch.tensor([0.2, 0.6, 0.9]).view(1, 3, 1, 1) # colour kept under alpha 0
|
||||
for method in ("lanczos", "box", "hamming"):
|
||||
resized = picker._resize_bchw(image, 8, 8, method)
|
||||
assert torch.allclose(resized, image), method
|
||||
@@ -0,0 +1,462 @@
|
||||
"""Regression tests: image core nodes (VTON batches, face paste, agnostic image, mask-aware
|
||||
resize, large image save, cache manager, grid composite)."""
|
||||
|
||||
import importlib
|
||||
import importlib.util
|
||||
import os
|
||||
|
||||
import comfy.cli_args
|
||||
import cv2
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
vton = importlib.import_module("NH-Nodes.nodes.vton.vton_preprocessor_nodes")
|
||||
face_paste = importlib.import_module("NH-Nodes.nodes.image.face_paste")
|
||||
agnostic = importlib.import_module("NH-Nodes.nodes.image.agnostic_image")
|
||||
mask_resize = importlib.import_module("NH-Nodes.nodes.image.mask_resize_image")
|
||||
large_save = importlib.import_module("NH-Nodes.nodes.image.large_image_save")
|
||||
cache_manager = importlib.import_module("NH-Nodes.nodes.image.large_preview_cache_manager")
|
||||
grid = importlib.import_module("NH-Nodes.nodes.image.image_grid_nodes")
|
||||
|
||||
RED, GREEN, BLUE = (1.0, 0.0, 0.0), (0.0, 1.0, 0.0), (0.0, 0.0, 1.0)
|
||||
YELLOW, WHITE, CYAN = (1.0, 1.0, 0.0), (1.0, 1.0, 1.0), (0.0, 1.0, 1.0)
|
||||
|
||||
|
||||
def _solid(color, h=32, w=32):
|
||||
"""One image [H,W,C] filled with a colour."""
|
||||
return torch.tensor(color, dtype=torch.float32).expand(h, w, len(color)).clone()
|
||||
|
||||
|
||||
def _boxes(boxes, h=32, w=32):
|
||||
"""MASK batch [B,H,W]; each box is (x1, y1, x2, y2) with an exclusive end."""
|
||||
masks = torch.zeros((len(boxes), h, w))
|
||||
for index, (x1, y1, x2, y2) in enumerate(boxes):
|
||||
masks[index, y1:y2, x1:x2] = 1.0
|
||||
return masks
|
||||
|
||||
|
||||
def _bbox(condition):
|
||||
"""(top, bottom, left, right) of the True pixels of a 2-D bool tensor."""
|
||||
rows, cols = torch.nonzero(condition, as_tuple=True)
|
||||
return int(rows.min()), int(rows.max()), int(cols.min()), int(cols.max())
|
||||
|
||||
|
||||
# --- VTON Ultimate Processor: batches, RGB conversion, no-person error ---
|
||||
|
||||
_MODEL_FILES = (
|
||||
"dwpose/yolox_l.onnx", "dwpose/dw-ll_ucoco_384.onnx",
|
||||
"humanparsing/parsing_atr.onnx", "humanparsing/parsing_lip.onnx",
|
||||
)
|
||||
|
||||
|
||||
def _vton_ready():
|
||||
if importlib.util.find_spec("onnxruntime") is None or importlib.util.find_spec("skimage") is None:
|
||||
return False
|
||||
return all(os.path.isfile(os.path.join(vton.MODELS_DIR, name)) for name in _MODEL_FILES)
|
||||
|
||||
|
||||
class _FakePose:
|
||||
"""DWposeDetector stand-in. Like the real one it never returns an empty array for a person-less
|
||||
image: it marks every keypoint -1. A completely black image counts as "no person"."""
|
||||
|
||||
def __init__(self):
|
||||
self.seen = []
|
||||
self.empty_when_no_person = False
|
||||
|
||||
def __call__(self, bgr):
|
||||
self.seen.append(bgr)
|
||||
height, width = bgr.shape[:2]
|
||||
pose_image = np.zeros((height, width, 3), np.uint8)
|
||||
if bgr.max() == 0:
|
||||
candidate = np.zeros((0, 134, 2), np.float32) if self.empty_when_no_person else np.full((1, 134, 2), -1.0, np.float32)
|
||||
else:
|
||||
candidate = np.full((1, 134, 2), 0.5, np.float32)
|
||||
return pose_image, None, None, candidate
|
||||
|
||||
|
||||
class _FakeParsing:
|
||||
"""Parsing stand-in: the left part of the image becomes upper clothes (label 5), wider when brighter."""
|
||||
|
||||
def __call__(self, image):
|
||||
pixels = np.array(image)
|
||||
labels = np.zeros(pixels.shape[:2], np.uint8)
|
||||
labels[:, : max(1, int(pixels.shape[1] * pixels.mean() / 255.0))] = 5
|
||||
return Image.fromarray(labels), Image.fromarray(labels)
|
||||
|
||||
|
||||
def _fake_get_mask_location(category, parsing_map, keypoints, width, height, *offsets):
|
||||
return Image.fromarray(((np.array(parsing_map) == 5) * 255).astype(np.uint8)), None
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fake_vton(monkeypatch):
|
||||
pytest.importorskip("skimage") # core.utils_mask imports it at module level
|
||||
utils_mask = importlib.import_module("NH-Nodes.core.utils_mask")
|
||||
pose = _FakePose()
|
||||
monkeypatch.setattr(utils_mask, "get_mask_location", _fake_get_mask_location)
|
||||
monkeypatch.setattr(vton, "LOADED_VTON_MODELS", {"dwprocessor": pose, "parsing_model": _FakeParsing(), "device": "cpu"})
|
||||
return pose
|
||||
|
||||
|
||||
def _process(images, category="Upper-body"):
|
||||
return vton.NH_VTonUltimateProcessor().process(images, category, 4, False, True, True, "cpu")
|
||||
|
||||
|
||||
def _flat_batch(levels, channels=3, h=48, w=64):
|
||||
return torch.stack([torch.full((h, w, channels), level) for level in levels])
|
||||
|
||||
|
||||
def test_vton_processes_every_batch_image_independently(fake_vton):
|
||||
batch = _flat_batch([0.2, 0.5, 0.8])
|
||||
outputs = _process(batch)
|
||||
|
||||
assert len(outputs) == 8
|
||||
assert all(output.shape[0] == 3 for output in outputs)
|
||||
for index in range(3):
|
||||
single = _process(batch[index:index + 1])
|
||||
for got, want in zip(outputs, single):
|
||||
assert torch.equal(got[index:index + 1], want)
|
||||
# The items really differ, so the comparison above is not trivially true.
|
||||
assert not torch.equal(outputs[0][0], outputs[0][2])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("channels", [1, 4])
|
||||
def test_vton_converts_grayscale_and_rgba_input_to_rgb(fake_vton, channels):
|
||||
outputs = _process(_flat_batch([0.5], channels=channels))
|
||||
|
||||
assert fake_vton.seen[0].ndim == 3 and fake_vton.seen[0].shape[2] == 3 # what DWPose receives
|
||||
assert outputs[4].shape == (1, 48, 64, 3) # masked_image
|
||||
|
||||
|
||||
@pytest.mark.parametrize("empty_candidate", [False, True])
|
||||
def test_vton_reports_batch_index_when_no_person_is_found(fake_vton, empty_candidate):
|
||||
fake_vton.empty_when_no_person = empty_candidate
|
||||
batch = _flat_batch([0.5, 0.0, 0.5]) # the middle image is black: nothing to detect
|
||||
|
||||
with pytest.raises(ValueError, match="batch index 1"):
|
||||
_process(batch)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not _vton_ready(), reason="VTON models or onnxruntime not available")
|
||||
def test_vton_batch_of_two_matches_single_results_on_cpu(monkeypatch):
|
||||
from skimage import data
|
||||
|
||||
monkeypatch.setattr(vton, "LOADED_VTON_MODELS", None)
|
||||
astronaut = torch.from_numpy(data.astronaut().astype(np.float32) / 255.0)
|
||||
batch = torch.stack([astronaut, astronaut.flip(1)]) # the second image is mirrored left-right
|
||||
node = vton.NH_VTonUltimateProcessor()
|
||||
|
||||
outputs = node.process(batch, "Upper-body", 4, False, True, True, "cpu")
|
||||
|
||||
assert len(outputs) == 8
|
||||
assert all(outputs[i].shape == (2, 512, 512) for i in (0, 1, 5, 6, 7))
|
||||
assert all(outputs[i].shape == (2, 512, 512, 3) for i in (2, 3, 4))
|
||||
for index in range(2):
|
||||
single = node.process(batch[index:index + 1], "Upper-body", 4, False, True, True, "cpu")
|
||||
for got, want in zip(outputs, single):
|
||||
torch.testing.assert_close(got[index:index + 1], want)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not _vton_ready(), reason="VTON models or onnxruntime not available")
|
||||
def test_vton_real_models_report_batch_index_for_person_less_image(monkeypatch):
|
||||
from skimage import data
|
||||
|
||||
monkeypatch.setattr(vton, "LOADED_VTON_MODELS", None)
|
||||
astronaut = torch.from_numpy(data.astronaut().astype(np.float32) / 255.0)
|
||||
blank = torch.full_like(astronaut, 0.5)
|
||||
|
||||
with pytest.raises(ValueError, match="batch index 1"):
|
||||
vton.NH_VTonUltimateProcessor().process(torch.stack([astronaut, blank]), "Upper-body", 4, False, True, True, "cpu")
|
||||
|
||||
|
||||
# --- NH Simple Face Paste ---
|
||||
|
||||
def _paste(dest, dest_mask, source, source_mask, feathering=0):
|
||||
return face_paste.NH_SimpleFacePaste().paste_face(dest, dest_mask, source, source_mask, feathering)[0]
|
||||
|
||||
|
||||
def test_face_paste_processes_every_item_and_broadcasts_single_inputs():
|
||||
dest = torch.stack([_solid(RED), _solid(GREEN)])
|
||||
|
||||
out = _paste(dest, _boxes([(8, 8, 24, 24)]), torch.stack([_solid(BLUE)]), _boxes([(0, 0, 16, 16)]))
|
||||
|
||||
assert out.shape == (2, 32, 32, 3)
|
||||
for index, color in enumerate((RED, GREEN)):
|
||||
assert torch.equal(out[index, 16, 16], torch.tensor(BLUE)) # pasted
|
||||
assert torch.equal(out[index, 2, 2], torch.tensor(color)) # destination untouched
|
||||
|
||||
|
||||
def test_face_paste_pairs_batch_items_by_index():
|
||||
dest = torch.stack([_solid(RED), _solid(GREEN)])
|
||||
source = torch.stack([_solid(BLUE), _solid(YELLOW)])
|
||||
|
||||
out = _paste(dest, _boxes([(0, 0, 8, 8), (24, 24, 32, 32)]), source, _boxes([(0, 0, 16, 16)]))
|
||||
|
||||
assert torch.equal(out[0, 4, 4], torch.tensor(BLUE)) and torch.equal(out[0, 28, 28], torch.tensor(RED))
|
||||
assert torch.equal(out[1, 28, 28], torch.tensor(YELLOW)) and torch.equal(out[1, 4, 4], torch.tensor(GREEN))
|
||||
|
||||
|
||||
def test_face_paste_clips_overshoot_instead_of_wrapping():
|
||||
# Resampled images can leave [0, 1]; uint8 conversion must saturate, not wrap around to black.
|
||||
dest = torch.stack([_solid((1.2, 1.2, 1.2))])
|
||||
source = torch.stack([_solid((-0.2, 1.2, 1.0))])
|
||||
|
||||
out = _paste(dest, _boxes([(8, 8, 24, 24)]), source, _boxes([(0, 0, 16, 16)]))
|
||||
|
||||
assert torch.equal(out[0, 16, 16], torch.tensor([0.0, 1.0, 1.0]))
|
||||
assert torch.equal(out[0, 2, 2], torch.tensor(WHITE))
|
||||
|
||||
|
||||
def test_face_paste_leaves_items_without_a_mask_untouched():
|
||||
dest = torch.rand((2, 32, 32, 3))
|
||||
dest_mask = _boxes([(8, 8, 24, 24), (0, 0, 0, 0)]) # the second mask is empty
|
||||
|
||||
out = _paste(dest, dest_mask, torch.stack([_solid(BLUE)]), _boxes([(0, 0, 16, 16)]))
|
||||
|
||||
assert out.shape == dest.shape
|
||||
assert torch.equal(out[0, 16, 16], torch.tensor(BLUE))
|
||||
assert torch.equal(out[1], dest[1]) # exact, not re-quantised through uint8
|
||||
|
||||
|
||||
def test_face_paste_accepts_2d_masks():
|
||||
out = _paste(
|
||||
torch.stack([_solid(RED)]), _boxes([(8, 8, 24, 24)])[0],
|
||||
torch.stack([_solid(BLUE)]), _boxes([(0, 0, 16, 16)])[0],
|
||||
)
|
||||
assert torch.equal(out[0, 16, 16], torch.tensor(BLUE))
|
||||
|
||||
|
||||
# --- NH Agnostic Image Generator ---
|
||||
|
||||
def _generate(image, mask, fill_mode="gray", blur_radius=65, noise_strength=1.0, gray_value=0.5, feathering=0):
|
||||
return agnostic.NH_AgnosticImageGenerator().generate(
|
||||
image, mask, fill_mode, blur_radius, noise_strength, gray_value, feathering,
|
||||
)
|
||||
|
||||
|
||||
def test_agnostic_resizes_a_mask_of_a_different_size():
|
||||
image = torch.rand((1, 16, 16, 3))
|
||||
mask = torch.zeros((1, 8, 8))
|
||||
mask[:, 2:6, 2:6] = 1.0
|
||||
|
||||
agnostic_img, masked_img, composite = _generate(image, mask)
|
||||
|
||||
assert agnostic_img.shape == masked_img.shape == composite.shape == (1, 16, 16, 3)
|
||||
assert torch.allclose(agnostic_img[0, 8, 8], torch.full((3,), 0.5)) # inside the (rescaled) mask
|
||||
assert torch.equal(agnostic_img[0, 0, 0], image[0, 0, 0]) # outside
|
||||
|
||||
|
||||
def test_agnostic_treats_a_2d_mask_as_a_batch_of_one():
|
||||
image = torch.rand((1, 16, 16, 3))
|
||||
mask = torch.zeros((16, 16))
|
||||
mask[4:12, 4:12] = 1.0
|
||||
|
||||
agnostic_img = _generate(image, mask)[0]
|
||||
|
||||
assert agnostic_img.shape == (1, 16, 16, 3)
|
||||
assert torch.allclose(agnostic_img[0, 8, 8], torch.full((3,), 0.5))
|
||||
assert torch.equal(agnostic_img[0, 0, 0], image[0, 0, 0])
|
||||
|
||||
|
||||
def test_agnostic_single_mask_broadcasts_over_the_image_batch():
|
||||
image = torch.rand((3, 16, 16, 3))
|
||||
mask = torch.zeros((1, 16, 16))
|
||||
mask[:, 4:12, 4:12] = 1.0
|
||||
|
||||
agnostic_img = _generate(image, mask)[0]
|
||||
|
||||
assert agnostic_img.shape == (3, 16, 16, 3)
|
||||
assert torch.allclose(agnostic_img[:, 8, 8], torch.full((3, 3), 0.5))
|
||||
|
||||
|
||||
def test_agnostic_binary_mask_result_is_unchanged():
|
||||
image = torch.rand((1, 16, 16, 3))
|
||||
mask = torch.zeros((1, 16, 16))
|
||||
mask[:, 4:12, 4:12] = 1.0
|
||||
|
||||
agnostic_img, masked_img, _ = _generate(image, mask)
|
||||
|
||||
expected = image.clone()
|
||||
expected[:, 4:12, 4:12] = 0.5
|
||||
assert torch.equal(agnostic_img, expected)
|
||||
expected[:, 4:12, 4:12] = 1.0
|
||||
assert torch.equal(masked_img, expected)
|
||||
|
||||
|
||||
def test_agnostic_keeps_soft_mask_values():
|
||||
image = torch.rand((1, 16, 16, 3))
|
||||
mask = torch.full((1, 16, 16), 0.25) # a soft mask must blend, not be thresholded away
|
||||
|
||||
agnostic_img = _generate(image, mask)[0]
|
||||
|
||||
assert torch.allclose(agnostic_img, image * 0.75 + 0.5 * 0.25, atol=1e-6)
|
||||
|
||||
|
||||
def test_agnostic_feathering_starts_from_the_soft_mask():
|
||||
image = torch.rand((1, 16, 16, 3))
|
||||
mask = torch.full((1, 16, 16), 0.5)
|
||||
|
||||
agnostic_img = _generate(image, mask, feathering=3)[0]
|
||||
|
||||
assert torch.allclose(agnostic_img, image * 0.5 + 0.5 * 0.5, atol=1e-5)
|
||||
|
||||
|
||||
def test_agnostic_preview_outline_follows_the_half_level_of_a_soft_mask():
|
||||
image = torch.zeros((1, 64, 64, 3))
|
||||
box = np.zeros((64, 64), np.float32)
|
||||
box[20:44, 20:44] = 1.0
|
||||
soft = torch.from_numpy(cv2.GaussianBlur(box, (31, 31), 5)).unsqueeze(0)
|
||||
|
||||
composite = _generate(image, soft)[2]
|
||||
|
||||
assert float(composite[0, 32, 8].sum()) == 0.0 # faint fringe (alpha ~0.01) is not outlined
|
||||
assert float(composite[0, 32, 20, 1]) > 0.5 # the box edge (alpha 0.5) is
|
||||
|
||||
|
||||
def test_agnostic_noise_is_deterministic_per_batch_index():
|
||||
image = torch.rand((2, 8, 8, 3))
|
||||
mask = torch.ones((2, 8, 8))
|
||||
|
||||
first = _generate(image, mask, "noise")[0]
|
||||
second = _generate(image, mask, "noise")[0]
|
||||
|
||||
assert torch.equal(first, second)
|
||||
assert not torch.equal(first[0], first[1])
|
||||
for index in range(2):
|
||||
expected = np.random.default_rng(index).random((8, 8, 3), dtype=np.float32)
|
||||
assert torch.allclose(first[index], torch.from_numpy(expected))
|
||||
|
||||
|
||||
# --- NH Mask-Aware Resize ---
|
||||
|
||||
def _marker_image(h=40, w=80):
|
||||
"""Blue background, red/green end stripes (12.5% wide each) and a white 8x8 marker in the middle."""
|
||||
image = torch.zeros((h, w, 3))
|
||||
image[..., 2] = 1.0
|
||||
image[:, :10] = torch.tensor(RED)
|
||||
image[:, 70:] = torch.tensor(GREEN)
|
||||
image[16:24, 36:44] = 1.0
|
||||
return image.unsqueeze(0)
|
||||
|
||||
|
||||
def _is_white(image):
|
||||
return (image[..., 0] > 0.5) & (image[..., 1] > 0.5) & (image[..., 2] > 0.5)
|
||||
|
||||
|
||||
def test_mask_resize_empty_mask_pad_fits_the_whole_image():
|
||||
image, mask = _marker_image(), torch.zeros((1, 40, 80))
|
||||
|
||||
out_img, out_mask = mask_resize.NH_MaskResizeImage().resize_by_mask(image, mask, 64, 64, "pad", "#00FFFF")
|
||||
|
||||
assert out_img.shape == (1, 64, 64, 3) and out_mask.shape == (1, 64, 64)
|
||||
top, bottom, left, right = _bbox((out_img[0] != torch.tensor(CYAN)).any(-1))
|
||||
assert (bottom - top + 1, right - left + 1) == (32, 64) # 80x40 contained in 64x64, padded top/bottom
|
||||
m_top, m_bottom, m_left, m_right = _bbox(_is_white(out_img[0]))
|
||||
assert abs((m_bottom - m_top) - (m_right - m_left)) <= 1 # the square marker stays square
|
||||
|
||||
|
||||
def test_mask_resize_empty_mask_crop_covers_the_target_with_a_centred_crop():
|
||||
image, mask = _marker_image(), torch.zeros((1, 40, 80))
|
||||
|
||||
out_img, out_mask = mask_resize.NH_MaskResizeImage().resize_by_mask(image, mask, 64, 64, "crop", "#00FFFF")
|
||||
|
||||
assert out_img.shape == (1, 64, 64, 3) and out_mask.shape == (1, 64, 64)
|
||||
red = (out_img[0, ..., 0] > 0.5) & (out_img[0, ..., 1] < 0.5) & (out_img[0, ..., 2] < 0.5)
|
||||
green = (out_img[0, ..., 1] > 0.5) & (out_img[0, ..., 0] < 0.5) & (out_img[0, ..., 2] < 0.5)
|
||||
assert not red.any() and not green.any() # the end stripes are cropped away, not squeezed in
|
||||
assert not (out_img[0] == torch.tensor(CYAN)).all(-1).any() # no padding: the image covers the target
|
||||
m_top, m_bottom, m_left, m_right = _bbox(_is_white(out_img[0]))
|
||||
assert abs((m_bottom - m_top) - (m_right - m_left)) <= 1
|
||||
|
||||
|
||||
def test_mask_resize_output_is_clipped_to_the_unit_range():
|
||||
image = torch.zeros((1, 64, 64, 3))
|
||||
image[:, :, 32:] = 1.0 # a hard edge makes Lanczos ring below 0 and above 1
|
||||
mask = torch.ones((1, 64, 64))
|
||||
|
||||
out_img, _ = mask_resize.NH_MaskResizeImage().resize_by_mask(image, mask, 100, 100, "pad", "#000000")
|
||||
|
||||
assert float(out_img.min()) >= 0.0 and float(out_img.max()) <= 1.0
|
||||
|
||||
|
||||
# --- NH Save Large Image ---
|
||||
|
||||
def _save(tmp_path, image, file_format, **kwargs):
|
||||
node = large_save.NH_SaveLargeImage()
|
||||
node.output_dir = str(tmp_path)
|
||||
return node.save(image, "nh_test", file_format, **kwargs)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"file_format,width",
|
||||
[("webp", 16384), ("jpg", 65501)],
|
||||
)
|
||||
def test_save_large_rejects_oversized_images_before_writing(tmp_path, file_format, width):
|
||||
batch = torch.zeros((2, 1, width, 3)) # 1 px tall keeps the test light
|
||||
|
||||
with pytest.raises(ValueError, match=r"(?is)too large.*png.*tiff"):
|
||||
_save(tmp_path, batch, file_format)
|
||||
|
||||
assert list(tmp_path.rglob("*")) == [] # not even the output folder was touched
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"file_format,width",
|
||||
[("webp", 16383), ("jpg", 65500), ("png", 70000), ("tiff", 70000)],
|
||||
)
|
||||
def test_save_large_writes_images_at_the_format_limit(tmp_path, file_format, width):
|
||||
result = _save(tmp_path, torch.zeros((1, 1, width, 3)), file_format)
|
||||
|
||||
saved_count = result["result"][2]
|
||||
assert saved_count == 1
|
||||
assert os.path.getsize(result["result"][1]) > 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("disabled", [False, True])
|
||||
def test_save_large_png_metadata_follows_disable_metadata(tmp_path, monkeypatch, disabled):
|
||||
monkeypatch.setattr(comfy.cli_args.args, "disable_metadata", disabled)
|
||||
|
||||
result = _save(
|
||||
tmp_path, torch.zeros((1, 4, 4, 3)), "png",
|
||||
prompt={"1": {"class_type": "Example"}}, extra_pnginfo={"workflow": {"nodes": []}},
|
||||
)
|
||||
|
||||
with Image.open(result["result"][1]) as saved:
|
||||
text = dict(saved.text)
|
||||
if disabled:
|
||||
assert "prompt" not in text and "workflow" not in text
|
||||
else:
|
||||
assert "prompt" in text and "workflow" in text
|
||||
|
||||
|
||||
# --- NH Large Preview Cache Manager ---
|
||||
|
||||
def test_cache_manager_runs_on_every_queue():
|
||||
changed = cache_manager.NH_LargePreviewCacheManager.IS_CHANGED(
|
||||
action="info", max_age_hours=24.0, max_total_mb=2048.0, cache_id="",
|
||||
)
|
||||
assert changed != changed # NaN never equals itself -> ComfyUI never serves the cached output
|
||||
|
||||
|
||||
# --- Image Grid Composite ---
|
||||
|
||||
def _grid(images, **overrides):
|
||||
params = dict(
|
||||
columns=2, rows=2, cell_size_mode="manual", cell_width=8, cell_height=8,
|
||||
resize_mode="stretch", gap_x=0, gap_y=0, background_color_hex="#000000",
|
||||
)
|
||||
params.update(overrides)
|
||||
return grid.NH_ImageGridComposite().composite(images, **params)
|
||||
|
||||
|
||||
def test_grid_builds_a_normal_canvas():
|
||||
image, slot_count, used_count, _info = _grid([torch.rand((1, 8, 8, 3))])
|
||||
|
||||
assert image.shape == (1, 16, 16, 3)
|
||||
assert (slot_count, used_count) == (4, 1)
|
||||
|
||||
|
||||
# The canvas size guard (free memory, not Pillow's decompression-bomb limit) is tested in
|
||||
# test_grid_save_path_resize_guards.py.
|
||||
@@ -0,0 +1,103 @@
|
||||
"""Channel handling of the nodes that turn IMAGE tensors into PIL pictures through core.utils.tensor_to_pil.
|
||||
|
||||
NH Save Large Image and the large preview cache keep an alpha channel where the format can hold it and give RGB
|
||||
otherwise (grey included); Image Grid Composite always builds RGB cells and drops alpha without compositing.
|
||||
"""
|
||||
|
||||
import importlib
|
||||
import os
|
||||
|
||||
import folder_paths
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
large_save = importlib.import_module("NH-Nodes.nodes.image.large_image_save")
|
||||
preview_cache = importlib.import_module("NH-Nodes.core.large_preview_cache")
|
||||
grid = importlib.import_module("NH-Nodes.nodes.image.image_grid_nodes")
|
||||
|
||||
|
||||
def _flat(height, width, channels):
|
||||
"""[1, H, W, C] image whose channel k holds the byte value 10 * (k + 1); the +0.5 keeps float rounding out of it."""
|
||||
values = torch.tensor([(10 * (k + 1) + 0.5) / 255.0 for k in range(channels)])
|
||||
return values.expand(1, height, width, channels).clone()
|
||||
|
||||
|
||||
def _save(tmp_path, image, file_format):
|
||||
node = large_save.NH_SaveLargeImage()
|
||||
node.output_dir = str(tmp_path)
|
||||
return node.save(image, "nh_channels", file_format)["result"][1].split("\n")[0]
|
||||
|
||||
|
||||
# --- NH Save Large Image ---
|
||||
|
||||
def test_png_keeps_the_alpha_channel(tmp_path):
|
||||
with Image.open(_save(tmp_path, _flat(4, 6, 4), "png")) as saved:
|
||||
assert saved.mode == "RGBA" and saved.getpixel((2, 1)) == (10, 20, 30, 40)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("file_format", ["jpg", "webp"])
|
||||
def test_formats_that_need_rgb_drop_the_alpha_channel(tmp_path, file_format):
|
||||
with Image.open(_save(tmp_path, _flat(4, 6, 4), file_format)) as saved:
|
||||
assert saved.mode == "RGB"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("channels,pixel", [(3, (10, 20, 30)), (1, (10, 10, 10))])
|
||||
def test_rgb_and_grey_images_are_saved_as_rgb(tmp_path, channels, pixel):
|
||||
with Image.open(_save(tmp_path, _flat(4, 6, channels), "png")) as saved:
|
||||
assert saved.mode == "RGB" and saved.getpixel((2, 1)) == pixel
|
||||
|
||||
|
||||
# --- large preview cache ---
|
||||
|
||||
@pytest.fixture
|
||||
def temp_dir(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(folder_paths, "temp_directory", str(tmp_path))
|
||||
return tmp_path
|
||||
|
||||
|
||||
def _thumbnail(manifest):
|
||||
thumb = manifest["thumb"]
|
||||
return Image.open(os.path.join(folder_paths.get_temp_directory(), thumb["subfolder"], thumb["filename"]))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("preview_format,mode", [("png", "RGBA"), ("jpg", "RGB"), ("webp", "RGB")])
|
||||
def test_preview_thumbnail_keeps_alpha_only_in_png(temp_dir, preview_format, mode):
|
||||
manifest = preview_cache.create_preview_cache(_flat(40, 30, 4)[0], preview_format=preview_format, tile_size=64)
|
||||
|
||||
with _thumbnail(manifest) as thumb:
|
||||
assert thumb.mode == mode and thumb.size == (30, 40)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", [(40, 30, 1), (40, 30, 3), (40, 30, 2), (40, 30)])
|
||||
def test_preview_accepts_grey_rgb_and_two_dimensional_images(temp_dir, shape):
|
||||
image = torch.rand(shape)
|
||||
|
||||
manifest = preview_cache.create_preview_cache(image, preview_format="png", tile_size=64)
|
||||
|
||||
assert (manifest["width"], manifest["height"]) == (30, 40)
|
||||
with _thumbnail(manifest) as thumb:
|
||||
assert thumb.mode == "RGB"
|
||||
|
||||
|
||||
def test_preview_rejects_anything_but_a_tensor(temp_dir):
|
||||
with pytest.raises(TypeError, match="Expected IMAGE tensor, got ndarray"):
|
||||
preview_cache.create_preview_cache(np.zeros((8, 8, 3), np.float32))
|
||||
|
||||
|
||||
# --- Image Grid Composite ---
|
||||
|
||||
def test_grid_cells_are_rgb_whatever_the_input_channels():
|
||||
transparent = _flat(4, 4, 4)
|
||||
transparent[..., 3] = 0.0 # fully transparent: the colour channels must survive, not be composited away
|
||||
grey = _flat(4, 4, 1)
|
||||
|
||||
image, _slots, _used, _info = grid.NH_ImageGridComposite().composite(
|
||||
[transparent, grey], columns=2, rows=1, cell_size_mode="manual", cell_width=4, cell_height=4,
|
||||
resize_mode="stretch", gap_x=0, gap_y=0, background_color_hex="#000000",
|
||||
)
|
||||
|
||||
assert image.shape == (1, 4, 8, 3)
|
||||
assert torch.allclose(image[0, 0, 0], torch.tensor([10.0, 20.0, 30.0]) / 255.0)
|
||||
assert torch.allclose(image[0, 0, 4], torch.tensor([10.0, 10.0, 10.0]) / 255.0)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,633 @@
|
||||
"""Wave 3 tests: lazy branch selection (If/Else, Switch N, Random Choice) and the torch Mask Morphology.
|
||||
|
||||
The lazy tests run every node through a small emulation of ComfyUI's executor (execution.py, execute()):
|
||||
connected-but-unevaluated lazy inputs arrive as None, unconnected optional inputs are absent from the kwargs, a
|
||||
name returned by check_lazy_status is only honoured while that input is still unevaluated, and asking for an input
|
||||
that is not connected is an error (NodeInputError). The ``old_*`` functions are verbatim copies of the node code
|
||||
before the lazy change, so the tests prove the choice itself did not move.
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import importlib
|
||||
import io
|
||||
import itertools
|
||||
import random
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from scipy.ndimage import (
|
||||
binary_dilation,
|
||||
binary_erosion,
|
||||
binary_fill_holes,
|
||||
gaussian_filter,
|
||||
maximum_filter1d,
|
||||
minimum_filter1d,
|
||||
)
|
||||
|
||||
logic_nodes = importlib.import_module("NH-Nodes.nodes.logic.logic_nodes")
|
||||
math_nodes = importlib.import_module("NH-Nodes.nodes.logic.math_nodes")
|
||||
morphology = importlib.import_module("NH-Nodes.nodes.mask.mask_morphology")
|
||||
|
||||
|
||||
# --- executor emulation -------------------------------------------------------------------------------------
|
||||
|
||||
def run_lazy(node, function, widgets, connected):
|
||||
"""Run `node` like ComfyUI does.
|
||||
|
||||
widgets: plain (non-lazy) inputs, always available. connected: {input name: upstream value} of every connected
|
||||
lazy input; only the ones check_lazy_status asks for are delivered, the rest arrive as None.
|
||||
Returns (outputs, requested) with `requested` = the lazy inputs that had to be evaluated upstream, in order.
|
||||
"""
|
||||
evaluated = {}
|
||||
requested = []
|
||||
while True:
|
||||
kwargs = {**widgets, **{name: evaluated.get(name) for name in connected}}
|
||||
wanted = [name for name in node.check_lazy_status(**kwargs) if name not in evaluated]
|
||||
if not wanted:
|
||||
return getattr(node, function)(**kwargs), requested
|
||||
for name in wanted:
|
||||
assert name in connected, f"asked for {name}, which is not connected (ComfyUI raises NodeInputError)"
|
||||
evaluated[name] = connected[name]
|
||||
requested.append(name)
|
||||
|
||||
|
||||
# --- the node code before the lazy change (reference) -------------------------------------------------------
|
||||
|
||||
def old_switch_n(index, **kwargs):
|
||||
connected = {}
|
||||
for i in range(10):
|
||||
key = f"input_{i}"
|
||||
if key in kwargs and kwargs[key] is not None:
|
||||
connected[i] = kwargs[key]
|
||||
|
||||
count = len(connected)
|
||||
if count == 0:
|
||||
return (None, 0)
|
||||
if index in connected:
|
||||
return (connected[index], count)
|
||||
first_key = sorted(connected.keys())[0]
|
||||
return (connected[first_key], count)
|
||||
|
||||
|
||||
def old_random_choice(seed, weights, **kwargs):
|
||||
slots = 6
|
||||
connected = []
|
||||
for i in range(slots):
|
||||
key = f"input_{i}"
|
||||
if key in kwargs and kwargs[key] is not None:
|
||||
connected.append((i, kwargs[key]))
|
||||
|
||||
if not connected:
|
||||
return (None, -1, "")
|
||||
|
||||
slot_weights = []
|
||||
for s in weights.split(",")[:slots]:
|
||||
try:
|
||||
slot_weights.append(max(0.0, float(s.strip())))
|
||||
except (ValueError, TypeError):
|
||||
slot_weights.append(1.0)
|
||||
slot_weights += [1.0] * (slots - len(slot_weights))
|
||||
w = [slot_weights[i] for i, _ in connected]
|
||||
|
||||
total = sum(w) or 1.0
|
||||
probs = [x / total for x in w]
|
||||
|
||||
rng = random.Random(seed)
|
||||
r = rng.random()
|
||||
cumsum = 0.0
|
||||
for idx, (orig_i, item) in enumerate(connected):
|
||||
cumsum += probs[idx]
|
||||
if r <= cumsum:
|
||||
prob_str = ", ".join(f"{p:.1%}" for p in probs)
|
||||
return (item, orig_i, prob_str)
|
||||
|
||||
orig_i, item = connected[-1]
|
||||
prob_str = ", ".join(f"{p:.1%}" for p in probs)
|
||||
return (item, orig_i, prob_str)
|
||||
|
||||
|
||||
# --- If/Else ------------------------------------------------------------------------------------------------
|
||||
|
||||
def test_if_else_branches_are_lazy_and_the_interface_is_unchanged():
|
||||
spec = logic_nodes.NH_IfElse.INPUT_TYPES()
|
||||
assert list(spec["required"]) == ["condition", "if_true", "if_false"]
|
||||
assert "optional" not in spec
|
||||
assert spec["required"]["if_true"] == ("*", {"lazy": True})
|
||||
assert spec["required"]["if_false"] == ("*", {"lazy": True})
|
||||
assert spec["required"]["condition"] == ("BOOLEAN", {"default": True}) # the condition must run first
|
||||
assert logic_nodes.NH_IfElse.RETURN_TYPES == ("*", "BOOLEAN")
|
||||
assert logic_nodes.NH_IfElse.RETURN_NAMES == ("result", "condition_out")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"condition,true_value,false_value,expected",
|
||||
[
|
||||
(True, None, None, ["if_true"]),
|
||||
(False, None, None, ["if_false"]),
|
||||
(True, None, "B", ["if_true"]), # the branch that is not chosen being ready does not help
|
||||
(False, "A", None, ["if_false"]),
|
||||
(True, "A", None, []),
|
||||
(False, None, "B", []),
|
||||
(True, "A", "B", []),
|
||||
(False, "A", "B", []),
|
||||
],
|
||||
)
|
||||
def test_if_else_check_lazy_status_asks_only_for_the_chosen_branch(condition, true_value, false_value, expected):
|
||||
node = logic_nodes.NH_IfElse()
|
||||
assert node.check_lazy_status(condition, if_true=true_value, if_false=false_value) == expected
|
||||
assert node.check_lazy_status(condition) == (["if_true"] if condition else ["if_false"]) # defaults: nothing ready
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", [0, 0.0, "", False, [], torch.zeros(2)])
|
||||
def test_if_else_an_evaluated_falsy_value_is_not_asked_for_again(value):
|
||||
node = logic_nodes.NH_IfElse()
|
||||
assert node.check_lazy_status(True, if_true=value) == []
|
||||
assert node.check_lazy_status(False, if_false=value) == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize("condition", [True, False])
|
||||
def test_if_else_runs_only_the_chosen_branch_and_returns_it(condition):
|
||||
node = logic_nodes.NH_IfElse()
|
||||
connected = {"if_true": "TRUE-VALUE", "if_false": "FALSE-VALUE"}
|
||||
outputs, requested = run_lazy(node, "switch", {"condition": condition}, connected)
|
||||
assert outputs == ("TRUE-VALUE" if condition else "FALSE-VALUE", condition)
|
||||
assert requested == (["if_true"] if condition else ["if_false"])
|
||||
|
||||
|
||||
def test_if_else_keeps_the_exact_object_and_handles_an_empty_chosen_value():
|
||||
node = logic_nodes.NH_IfElse()
|
||||
tensor = torch.rand(1, 4, 4, 3)
|
||||
outputs, requested = run_lazy(node, "switch", {"condition": False}, {"if_true": "x", "if_false": tensor})
|
||||
assert outputs[0] is tensor and requested == ["if_false"]
|
||||
# an upstream node that returned None: asked once (the executor drops the repeated ask), result is None
|
||||
outputs, requested = run_lazy(node, "switch", {"condition": True}, {"if_true": None, "if_false": "x"})
|
||||
assert outputs == (None, True) and requested == ["if_true"]
|
||||
|
||||
|
||||
# --- Switch N -----------------------------------------------------------------------------------------------
|
||||
|
||||
def test_switch_n_inputs_are_lazy_and_the_interface_is_unchanged():
|
||||
spec = logic_nodes.NH_SwitchN.INPUT_TYPES()
|
||||
assert list(spec["required"]) == ["index"]
|
||||
assert spec["required"]["index"] == ("INT", {"default": 0, "min": 0, "max": 9})
|
||||
assert list(spec["optional"]) == [f"input_{i}" for i in range(10)]
|
||||
assert all(spec["optional"][f"input_{i}"] == ("*", {"lazy": True}) for i in range(10))
|
||||
assert logic_nodes.NH_SwitchN.RETURN_TYPES == ("*", "INT")
|
||||
assert logic_nodes.NH_SwitchN.RETURN_NAMES == ("result", "count")
|
||||
|
||||
|
||||
def test_switch_n_description_says_that_connected_inputs_count_even_without_a_value():
|
||||
# With lazy inputs the node cannot see a value before it asks for it, so a connected slot whose upstream node
|
||||
# returns None is picked (the old code skipped None values). The user is told.
|
||||
description = logic_nodes.NH_SwitchN.DESCRIPTION
|
||||
assert "Connected inputs count even if they produce no value" in description
|
||||
assert "Only the chosen input is executed" in description
|
||||
|
||||
|
||||
def test_switch_n_check_lazy_status_values():
|
||||
node = logic_nodes.NH_SwitchN()
|
||||
unevaluated = {"input_0": None, "input_2": None, "input_5": None} # connected, nothing run yet
|
||||
assert node.check_lazy_status(2, **unevaluated) == ["input_2"]
|
||||
assert node.check_lazy_status(5, **unevaluated) == ["input_5"]
|
||||
assert node.check_lazy_status(1, **unevaluated) == ["input_0"] # slot 1 is not connected: lowest connected slot
|
||||
assert node.check_lazy_status(9, **unevaluated) == ["input_0"]
|
||||
assert node.check_lazy_status(2, input_0=None, input_2="two", input_5=None) == [] # chosen one is ready
|
||||
assert node.check_lazy_status(2, input_0="zero", input_2=None, input_5="five") == ["input_2"]
|
||||
assert node.check_lazy_status(3) == [] # nothing connected: nothing to ask for
|
||||
assert node.check_lazy_status(0, input_4=0) == [] # a falsy value is an evaluated value
|
||||
|
||||
|
||||
def test_switch_n_runs_only_the_selected_input():
|
||||
node = logic_nodes.NH_SwitchN()
|
||||
connected = {f"input_{i}": f"value-{i}" for i in range(4)}
|
||||
outputs, requested = run_lazy(node, "switch", {"index": 2}, connected)
|
||||
assert outputs == ("value-2", 4)
|
||||
assert requested == ["input_2"]
|
||||
|
||||
|
||||
def test_switch_n_unconnected_selected_slot_falls_back_to_the_lowest_connected_one(capsys):
|
||||
node = logic_nodes.NH_SwitchN()
|
||||
connected = {"input_2": "two", "input_5": "five", "input_7": "seven"}
|
||||
outputs, requested = run_lazy(node, "switch", {"index": 3}, connected)
|
||||
assert outputs == ("two", 3)
|
||||
assert requested == ["input_2"] # the fallback slot, not every connected one
|
||||
assert capsys.readouterr().out.count("[NH-Nodes] SwitchN: index 3 not connected, falling back to input_2") == 1
|
||||
outputs, requested = run_lazy(node, "switch", {"index": 5}, connected)
|
||||
assert (outputs, requested) == (("five", 3), ["input_5"])
|
||||
assert capsys.readouterr().out == "" # no fallback message when the selected slot exists
|
||||
|
||||
|
||||
def test_switch_n_with_nothing_connected_warns_and_returns_none(capsys):
|
||||
node = logic_nodes.NH_SwitchN()
|
||||
outputs, requested = run_lazy(node, "switch", {"index": 4}, {})
|
||||
assert outputs == (None, 0) and requested == []
|
||||
assert capsys.readouterr().out.count("[NH-Nodes] Warning: SwitchN has no connected inputs") == 1
|
||||
|
||||
|
||||
def test_switch_n_count_is_the_number_of_connected_slots_not_of_evaluated_ones():
|
||||
node = logic_nodes.NH_SwitchN()
|
||||
outputs, requested = run_lazy(node, "switch", {"index": 1}, {"input_0": "a", "input_1": "b", "input_9": "c"})
|
||||
assert outputs == ("b", 3) and requested == ["input_1"]
|
||||
# A connected slot whose upstream produced None still counts as connected (it used to be skipped).
|
||||
outputs, requested = run_lazy(node, "switch", {"index": 1}, {"input_0": "a", "input_1": None})
|
||||
assert outputs == (None, 2) and requested == ["input_1"]
|
||||
|
||||
|
||||
def test_switch_n_makes_the_same_choice_as_before_for_every_combination():
|
||||
node = logic_nodes.NH_SwitchN()
|
||||
checked = 0
|
||||
with contextlib.redirect_stdout(io.StringIO()):
|
||||
for subset in range(1 << 10):
|
||||
slots = [i for i in range(10) if subset >> i & 1]
|
||||
values = {f"input_{i}": f"value-{i}" for i in slots}
|
||||
for index in range(10):
|
||||
outputs, requested = run_lazy(node, "switch", {"index": index}, values)
|
||||
assert outputs == old_switch_n(index, **values), (slots, index)
|
||||
if slots:
|
||||
chosen = index if index in slots else slots[0]
|
||||
assert requested == [f"input_{chosen}"], (slots, index)
|
||||
else:
|
||||
assert requested == []
|
||||
checked += 1
|
||||
assert checked == 10 * 1024
|
||||
|
||||
|
||||
# --- Random Choice ------------------------------------------------------------------------------------------
|
||||
|
||||
def test_random_choice_inputs_are_lazy_and_the_interface_is_unchanged():
|
||||
spec = math_nodes.NH_RandomChoice.INPUT_TYPES()
|
||||
assert list(spec["required"]) == ["seed", "weights"]
|
||||
assert spec["required"]["seed"] == ("INT", {"default": 0, "min": 0, "max": 2**32 - 1})
|
||||
assert spec["required"]["weights"] == ("STRING", {"default": "1,1,1,1,1,1"})
|
||||
assert list(spec["optional"]) == [f"input_{i}" for i in range(6)]
|
||||
assert all(spec["optional"][f"input_{i}"] == ("*", {"lazy": True}) for i in range(6))
|
||||
assert math_nodes.NH_RandomChoice.RETURN_TYPES == ("*", "INT", "STRING")
|
||||
assert math_nodes.NH_RandomChoice.RETURN_NAMES == ("result", "picked_index", "probabilities")
|
||||
|
||||
|
||||
def test_random_choice_description_says_that_connected_inputs_count_even_without_a_value():
|
||||
description = math_nodes.NH_RandomChoice.DESCRIPTION
|
||||
assert "Connected inputs count even if they produce no value" in description
|
||||
assert "Only the picked input is executed" in description
|
||||
|
||||
|
||||
def test_random_choice_check_lazy_status_values():
|
||||
node = math_nodes.NH_RandomChoice()
|
||||
unevaluated = {f"input_{i}": None for i in range(6)}
|
||||
assert node.check_lazy_status(1, "0,0,5,0,0,0", **unevaluated) == ["input_2"] # only slot 2 has any weight
|
||||
assert node.check_lazy_status(1, "0,0,5,0,0,0", **{**unevaluated, "input_2": "ready"}) == []
|
||||
assert node.check_lazy_status(1, "1,1,1,1,1,1") == [] # nothing connected
|
||||
only_two = {"input_1": None, "input_4": None}
|
||||
for seed in range(20):
|
||||
picked = old_random_choice(seed, "1,1,1,1,1,1", input_1="a", input_4="b")[1]
|
||||
assert node.check_lazy_status(seed, "1,1,1,1,1,1", **only_two) == [f"input_{picked}"]
|
||||
|
||||
|
||||
def test_random_choice_runs_only_the_picked_input():
|
||||
node = math_nodes.NH_RandomChoice()
|
||||
for slot in range(6):
|
||||
weights = ",".join("5" if i == slot else "0" for i in range(6))
|
||||
connected = {f"input_{i}": f"value-{i}" for i in range(6)}
|
||||
outputs, requested = run_lazy(node, "choose", {"seed": 11, "weights": weights}, connected)
|
||||
assert outputs[:2] == (f"value-{slot}", slot)
|
||||
assert requested == [f"input_{slot}"]
|
||||
|
||||
|
||||
def test_random_choice_with_nothing_connected_returns_none():
|
||||
outputs, requested = run_lazy(math_nodes.NH_RandomChoice(), "choose", {"seed": 3, "weights": "1,1"}, {})
|
||||
assert outputs == (None, -1, "") and requested == []
|
||||
|
||||
|
||||
def test_random_choice_makes_the_same_pick_as_before_for_every_seed_weights_and_connection():
|
||||
node = math_nodes.NH_RandomChoice()
|
||||
weight_sets = ["1,1,1,1,1,1", "3,0,1,2,5,0", "0,0,0,0,0,0", "1,1,10", "x,2,,4", "", "2", "0.5, 0.25 ,0.25,9,9,9,9,9"]
|
||||
seeds = [0, 1, 2, 3, 7, 42, 99, 12345, 2**32 - 1] + list(range(10, 40))
|
||||
checked = 0
|
||||
for subset in range(1, 1 << 6):
|
||||
slots = [i for i in range(6) if subset >> i & 1]
|
||||
values = {f"input_{i}": f"value-{i}" for i in slots}
|
||||
for weights, seed in itertools.product(weight_sets, seeds):
|
||||
outputs, requested = run_lazy(node, "choose", {"seed": seed, "weights": weights}, values)
|
||||
expected = old_random_choice(seed, weights, **values)
|
||||
assert outputs == expected, (slots, weights, seed)
|
||||
assert requested == [f"input_{expected[1]}"], (slots, weights, seed)
|
||||
checked += 1
|
||||
assert checked == 63 * len(weight_sets) * len(seeds)
|
||||
|
||||
|
||||
def test_random_choice_empty_picked_value_is_returned_as_is():
|
||||
node = math_nodes.NH_RandomChoice()
|
||||
outputs, requested = run_lazy(node, "choose", {"seed": 1, "weights": "0,1"}, {"input_0": "a", "input_1": None})
|
||||
assert outputs == (None, 1, "0.0%, 100.0%") and requested == ["input_1"]
|
||||
|
||||
|
||||
# --- Mask Morphology ----------------------------------------------------------------------------------------
|
||||
|
||||
def process(mask, horizontal=0, vertical=0, fill_holes=False, blur=0.0):
|
||||
return morphology.NH_MaskMorphology().process(mask, horizontal, vertical, fill_holes, blur)[0]
|
||||
|
||||
|
||||
def blobs(shape, seed, sigma=3.0):
|
||||
"""Binary mask made of blobs a few pixels wide, so dilation and erosion both have something to do."""
|
||||
rng = np.random.default_rng(seed)
|
||||
noise = gaussian_filter(rng.random(shape), sigma=sigma)
|
||||
return (noise > np.median(noise)).astype(np.float32)
|
||||
|
||||
|
||||
def old_expand(mask_np, horizontal_expand, vertical_expand):
|
||||
"""The scipy code of Mask Morphology before the torch rewrite (expand steps only)."""
|
||||
m = mask_np
|
||||
if horizontal_expand != 0:
|
||||
op = binary_dilation if horizontal_expand > 0 else binary_erosion
|
||||
m = op(m, structure=np.ones((1, abs(horizontal_expand)))).astype(m.dtype)
|
||||
if vertical_expand != 0:
|
||||
op = binary_dilation if vertical_expand > 0 else binary_erosion
|
||||
m = op(m, structure=np.ones((abs(vertical_expand), 1))).astype(m.dtype)
|
||||
return m
|
||||
|
||||
|
||||
def test_mask_morphology_interface_is_unchanged():
|
||||
spec = morphology.NH_MaskMorphology.INPUT_TYPES()
|
||||
assert list(spec["required"]) == ["mask", "horizontal_expand", "vertical_expand", "fill_holes", "blur_radius"]
|
||||
assert spec["required"]["horizontal_expand"] == ("INT", {"default": 0, "min": -512, "max": 512, "step": 1})
|
||||
assert spec["required"]["vertical_expand"] == ("INT", {"default": 0, "min": -512, "max": 512, "step": 1})
|
||||
assert spec["required"]["blur_radius"] == ("FLOAT", {"default": 0.0, "min": 0.0, "max": 100.0, "step": 0.1})
|
||||
assert morphology.NH_MaskMorphology.RETURN_TYPES == ("MASK",)
|
||||
assert morphology.NODE_CLASS_MAPPINGS == {"NH_MaskAdvancedMorphology": morphology.NH_MaskMorphology}
|
||||
assert morphology.NODE_DISPLAY_NAME_MAPPINGS == {"NH_MaskAdvancedMorphology": "Mask Morphology (NH)"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("horizontal,vertical", list(itertools.product([-9, -5, -3, 0, 3, 5, 9], repeat=2)))
|
||||
def test_binary_mask_with_odd_kernels_equals_the_old_result_away_from_the_borders(horizontal, vertical):
|
||||
for seed in range(3):
|
||||
mask_np = blobs((72, 80), seed) # not square: catches a swapped axis
|
||||
expected = old_expand(mask_np, horizontal, vertical)
|
||||
got = process(torch.from_numpy(mask_np)[None], horizontal, vertical)[0].numpy()
|
||||
assert got.dtype == np.float32 and got.shape == expected.shape
|
||||
rh, rv = abs(horizontal) // 2, abs(vertical) // 2
|
||||
interior = (slice(rv, 72 - rv), slice(rh, 80 - rh))
|
||||
assert np.array_equal(got[interior], expected[interior]), (horizontal, vertical, seed)
|
||||
# The old code treated the outside as empty, so it could only eat more. Growing is identical everywhere.
|
||||
assert np.all(got >= expected)
|
||||
if horizontal >= 0 and vertical >= 0:
|
||||
assert np.array_equal(got, expected)
|
||||
|
||||
|
||||
def test_soft_masks_keep_their_gradient_instead_of_being_binarised():
|
||||
soft = torch.zeros(1, 9, 9)
|
||||
soft[0, 4, 4] = 0.3
|
||||
grown = process(soft, 5, 0)
|
||||
assert torch.equal(grown[0, 4], torch.tensor([0, 0, 0.3, 0.3, 0.3, 0.3, 0.3, 0, 0])) # old code gave 1.0 here
|
||||
assert float(grown.sum()) == pytest.approx(5 * 0.3)
|
||||
|
||||
ramp = torch.linspace(0, 1, 11).repeat(5, 1)[None] # values rise from 0 to 1 along the width
|
||||
assert torch.allclose(process(ramp, 3, 0)[0, 0, :-1], ramp[0, 0, 1:]) # max over 3 = the next pixel
|
||||
assert torch.allclose(process(ramp, -3, 0)[0, 0, 1:], ramp[0, 0, :-1]) # min over 3 = the previous pixel
|
||||
|
||||
plateau = torch.full((1, 7, 7), 0.8)
|
||||
plateau[0, 3, 3] = 0.2
|
||||
eroded = process(plateau, -3, -3)
|
||||
assert float(eroded[0, 3, 3]) == pytest.approx(0.2) and float(eroded[0, 0, 0]) == pytest.approx(0.8)
|
||||
assert bool(((eroded > 0.19) & (eroded < 0.81)).all()) # nothing was pushed to 0 or 1
|
||||
|
||||
|
||||
def test_erosion_does_not_eat_a_mask_from_the_image_border():
|
||||
full = torch.ones(1, 12, 12)
|
||||
assert torch.equal(process(full, -9, -9), full) # the old code emptied the border band of 4 px
|
||||
|
||||
left = torch.zeros(1, 6, 14)
|
||||
left[0, :, :8] = 1.0 # the mask touches the left edge, the empty part starts at column 8
|
||||
eroded = process(left, -5, 0)[0]
|
||||
assert torch.equal(eroded[:, :6], torch.ones(6, 6)) # columns 0..5 stay (the old code lost columns 0 and 1)
|
||||
assert float(eroded[:, 6:].sum()) == 0.0 # the side facing the empty area still shrinks by 2
|
||||
|
||||
top = torch.zeros(1, 14, 6)
|
||||
top[0, :8, :] = 1.0
|
||||
assert torch.equal(process(top, 0, -5)[0, :6], torch.ones(6, 6))
|
||||
assert float(process(top, 0, -5)[0, 6:].sum()) == 0.0
|
||||
|
||||
|
||||
def test_dilation_stops_at_the_image_border_and_does_not_wrap_around():
|
||||
mask = torch.zeros(1, 5, 10)
|
||||
mask[0, 2, 0] = 1.0
|
||||
grown = process(mask, 5, 5)[0]
|
||||
assert torch.equal(grown[:, :3], torch.ones(5, 3)) # radius 2 around (2, 0): rows 0..4, columns 0..2
|
||||
assert float(grown[:, 3:].sum()) == 0.0 # nothing wrapped to the far side
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"expand,columns",
|
||||
[
|
||||
(1, (5, 5)), # kernel length 1: nothing changes
|
||||
(2, (4, 6)), # even sizes are rounded up to the next odd length: 2 -> 3
|
||||
(3, (4, 6)),
|
||||
(4, (3, 7)), # 4 -> 5
|
||||
(5, (3, 7)),
|
||||
(10, (0, 10)), # 10 -> 11 (radius 5), clipped by the 11 px wide image
|
||||
],
|
||||
)
|
||||
def test_even_kernel_lengths_are_rounded_up_so_growing_stays_centred(expand, columns):
|
||||
mask = torch.zeros(1, 3, 11)
|
||||
mask[0, :, 5] = 1.0
|
||||
grown = process(mask, expand, 0)[0]
|
||||
first, last = columns
|
||||
assert torch.equal(grown[1], torch.tensor([1.0 if first <= x <= last else 0.0 for x in range(11)]))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("expand,columns", [(-2, (2, 8)), (-3, (2, 8)), (-4, (3, 7)), (-5, (3, 7)), (-6, (4, 6))])
|
||||
def test_even_kernel_lengths_are_rounded_up_so_shrinking_stays_centred(expand, columns):
|
||||
mask = torch.zeros(1, 3, 11)
|
||||
mask[0, :, 1:10] = 1.0 # run of 9 ones, columns 1..9
|
||||
shrunk = process(mask, expand, 0)[0]
|
||||
first, last = columns
|
||||
assert torch.equal(shrunk[1], torch.tensor([1.0 if first <= x <= last else 0.0 for x in range(11)]))
|
||||
|
||||
|
||||
def test_horizontal_expand_acts_on_columns_and_vertical_expand_on_rows_in_that_order():
|
||||
mask = torch.zeros(1, 9, 9)
|
||||
mask[0, 4, 4] = 1.0
|
||||
wide = process(mask, 5, 0)[0]
|
||||
assert torch.equal(wide.nonzero()[:, 0].unique(), torch.tensor([4])) and int(wide.sum()) == 5
|
||||
tall = process(mask, 0, 5)[0]
|
||||
assert torch.equal(tall.nonzero()[:, 1].unique(), torch.tensor([4])) and int(tall.sum()) == 5
|
||||
assert int(process(mask, 5, 5).sum()) == 25
|
||||
|
||||
|
||||
def test_horizontal_step_runs_before_the_vertical_step_like_before():
|
||||
diagonal = torch.zeros(1, 9, 9)
|
||||
for step in (3, 4, 5):
|
||||
diagonal[0, step, step] = 1.0
|
||||
# Grow sideways first: the three pixels overlap in column 4 only, so one pixel survives the vertical erosion.
|
||||
# (Eroding first would leave nothing.)
|
||||
result = process(diagonal, 3, -3)[0]
|
||||
assert int(result.sum()) == 1 and result[4, 4] == 1.0
|
||||
expected = old_expand(diagonal[0].numpy(), 3, -3)
|
||||
assert np.array_equal(result.numpy(), expected)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("size", [3, 8, 31, 64, 255, 513])
|
||||
def test_any_kernel_size_matches_an_independent_window_filter(size):
|
||||
rng = np.random.default_rng(size)
|
||||
soft = rng.random((2, 37, 41), dtype=np.float32)
|
||||
radius_size = 2 * (size // 2) + 1 # odd length actually used (even lengths round up)
|
||||
for sign, window_filter in ((1, maximum_filter1d), (-1, minimum_filter1d)):
|
||||
horizontal_only = process(torch.from_numpy(soft), sign * size, 0).numpy()
|
||||
expected = np.stack([window_filter(m, size=radius_size, axis=1, mode="nearest") for m in soft])
|
||||
assert np.array_equal(horizontal_only, expected), (size, sign)
|
||||
|
||||
vertical_only = process(torch.from_numpy(soft), 0, sign * size).numpy()
|
||||
expected = np.stack([window_filter(m, size=radius_size, axis=0, mode="nearest") for m in soft])
|
||||
assert np.array_equal(vertical_only, expected), (size, sign)
|
||||
|
||||
both = process(torch.from_numpy(soft), sign * size, sign * size).numpy()
|
||||
expected = np.stack([window_filter(window_filter(m, size=radius_size, axis=1, mode="nearest"),
|
||||
size=radius_size, axis=0, mode="nearest") for m in soft])
|
||||
assert np.array_equal(both, expected), (size, sign)
|
||||
|
||||
|
||||
def test_a_kernel_wider_than_the_image_gives_the_global_maximum_or_minimum():
|
||||
soft = torch.rand(1, 6, 9)
|
||||
assert torch.equal(process(soft, 512, 0)[0], soft[0].amax(dim=1, keepdim=True).expand(6, 9))
|
||||
assert torch.equal(process(soft, 0, -512)[0], soft[0].amin(dim=0, keepdim=True).expand(6, 9))
|
||||
tiny = torch.rand(1, 1, 1)
|
||||
assert torch.equal(process(tiny, 512, -512), tiny)
|
||||
|
||||
|
||||
def test_batch_items_are_independent_and_the_input_is_not_modified():
|
||||
batch = torch.from_numpy(np.stack([blobs((40, 50), seed) for seed in range(4)]))
|
||||
original = batch.clone()
|
||||
together = process(batch, 5, -3)
|
||||
assert together.shape == batch.shape and together.dtype == torch.float32
|
||||
assert torch.equal(batch, original)
|
||||
for index in range(4):
|
||||
assert torch.equal(together[index], process(batch[index : index + 1], 5, -3)[0])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("budget", [1, 42, 100, 126, 10**9])
|
||||
def test_long_batches_are_processed_in_chunks_with_the_same_result(monkeypatch, budget):
|
||||
batch = torch.from_numpy(np.stack([blobs((6, 7), seed) for seed in range(5)])) # 42 pixels per mask
|
||||
expected = process(batch, 3, -3)
|
||||
monkeypatch.setattr(morphology, "_CHUNK_PIXELS", budget) # 1 and 42: one mask per chunk; 100: 2+2+1; 126: 3+2
|
||||
assert torch.equal(process(batch, 3, -3), expected)
|
||||
assert torch.equal(process(batch, 0, 0), batch) and process(batch).data_ptr() != batch.data_ptr()
|
||||
filled = process(batch, 3, 3, True, 0.8)
|
||||
monkeypatch.setattr(morphology, "_CHUNK_PIXELS", 1 << 24)
|
||||
assert torch.equal(filled, process(batch, 3, 3, True, 0.8))
|
||||
|
||||
|
||||
def test_unchanged_mask_comes_back_as_a_copy():
|
||||
mask = torch.rand(2, 5, 6)
|
||||
out = process(mask)
|
||||
assert torch.equal(out, mask) and out.data_ptr() != mask.data_ptr()
|
||||
out += 1.0
|
||||
assert not torch.equal(out, mask) # a node downstream may modify its input in place
|
||||
# Kernel length 1 and an image only one pixel wide have nothing to do either, and still return a copy.
|
||||
for horizontal, vertical in ((1, -1), (-1, 1)):
|
||||
same = process(mask, horizontal, vertical)
|
||||
assert torch.equal(same, mask) and same.data_ptr() != mask.data_ptr()
|
||||
thin = torch.rand(1, 5, 1)
|
||||
assert process(thin, 9, 0).data_ptr() != thin.data_ptr() and torch.equal(process(thin, -9, 0), thin)
|
||||
|
||||
|
||||
def test_two_dimensional_masks_are_accepted_and_other_shapes_are_rejected_clearly():
|
||||
flat = torch.zeros(7, 7)
|
||||
flat[3, 3] = 1.0
|
||||
out = process(flat, 3, 3)
|
||||
assert out.shape == (1, 7, 7) and int(out.sum()) == 9
|
||||
with pytest.raises(ValueError, match="MASK"):
|
||||
process(torch.zeros(1, 1, 4, 4), 3, 3)
|
||||
with pytest.raises(ValueError, match="MASK"):
|
||||
process(torch.zeros(5), 3, 3)
|
||||
|
||||
|
||||
def test_empty_batch_is_returned_empty():
|
||||
out = process(torch.zeros(0, 8, 8), 5, -5, True, 1.0)
|
||||
assert out.shape == (0, 8, 8)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", [(1, 0, 8), (1, 8, 0), (3, 0, 0), (0, 8)])
|
||||
@pytest.mark.parametrize(
|
||||
"settings",
|
||||
[
|
||||
(0, 0),
|
||||
(5, 0),
|
||||
(0, 5),
|
||||
(-5, 0),
|
||||
(0, -5),
|
||||
(5, -5, False, 0.0),
|
||||
(9, 9, True, 0.0),
|
||||
(0, 0, False, 1.5),
|
||||
(3, 3, True, 2.0),
|
||||
],
|
||||
)
|
||||
def test_a_mask_without_pixels_comes_back_unchanged_like_before(shape, settings):
|
||||
# Zero height or width: there is nothing to grow, fill or blur. The scipy version passed such a mask through
|
||||
# whatever the settings; replicate padding cannot pad it (it raises), so the node must not try.
|
||||
mask = torch.zeros(shape)
|
||||
out = process(mask, *settings)
|
||||
assert out is not mask # a copy, like every other unchanged result
|
||||
assert out.shape == ((1, *shape) if len(shape) == 2 else shape) # a 2-D mask is a batch of one
|
||||
assert out.dtype == torch.float32 and out.numel() == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"dtype,expected",
|
||||
[(torch.float32, torch.float32), (torch.float64, torch.float64), (torch.float16, torch.float16), (torch.bool, torch.float32)],
|
||||
)
|
||||
def test_a_mask_without_pixels_follows_the_same_dtype_rule_as_any_other(dtype, expected):
|
||||
out = process(torch.zeros(2, 0, 8, dtype=dtype), 5, 5, True, 1.0)
|
||||
assert out.shape == (2, 0, 8) and out.dtype == expected
|
||||
|
||||
|
||||
def test_non_float_and_float64_masks_are_handled():
|
||||
as_bool = torch.zeros(1, 5, 5, dtype=torch.bool)
|
||||
as_bool[0, 2, 2] = True
|
||||
out = process(as_bool, 3, 3)
|
||||
assert out.dtype == torch.float32 and int(out.sum()) == 9
|
||||
as_double = torch.zeros(1, 5, 5, dtype=torch.float64)
|
||||
as_double[0, 2, 2] = 1.0
|
||||
out = process(as_double, 3, 3)
|
||||
assert out.dtype == torch.float64 and int(out.sum()) == 9
|
||||
|
||||
|
||||
def test_fill_holes_and_blur_still_use_the_same_scipy_calls():
|
||||
ring = np.zeros((30, 30), np.float32)
|
||||
yy, xx = np.mgrid[:30, :30]
|
||||
ring[(np.hypot(yy - 15, xx - 15) < 10) & (np.hypot(yy - 15, xx - 15) > 4)] = 1.0
|
||||
ring_t = torch.from_numpy(ring)[None]
|
||||
|
||||
filled = process(ring_t, fill_holes=True)[0].numpy()
|
||||
assert np.array_equal(filled, binary_fill_holes(ring > 0.5).astype(np.float32))
|
||||
assert filled[15, 15] == 1.0 and ring[15, 15] == 0.0
|
||||
|
||||
blurred = process(ring_t, blur=2.5)[0].numpy()
|
||||
assert np.array_equal(blurred, gaussian_filter(ring, sigma=2.5))
|
||||
|
||||
# expand -> fill holes -> blur, in that order
|
||||
both = process(ring_t, 3, 3, True, 1.5)[0].numpy()
|
||||
grown = process(ring_t, 3, 3)[0].numpy()
|
||||
assert np.array_equal(both, gaussian_filter(binary_fill_holes(grown > 0.5).astype(np.float32), sigma=1.5))
|
||||
|
||||
|
||||
def test_fill_holes_binarises_a_soft_mask_at_one_half_like_before():
|
||||
soft = torch.zeros(1, 9, 9)
|
||||
soft[0, 2:7, 2:7] = 0.9
|
||||
soft[0, 4, 4] = 0.0 # hole
|
||||
soft[0, 2, 2] = 0.4 # below the 0.5 threshold: not part of the mask
|
||||
out = process(soft, fill_holes=True)[0]
|
||||
assert out[4, 4] == 1.0 and out[2, 2] == 0.0 and out[3, 3] == 1.0
|
||||
assert set(out.unique().tolist()) == {0.0, 1.0}
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device")
|
||||
def test_morphology_runs_on_the_tensor_device_and_matches_the_cpu_result():
|
||||
mask = torch.from_numpy(np.stack([blobs((64, 80), seed) for seed in range(2)]))
|
||||
mask[0, 10:20, 10:20] *= 0.5 # keep some soft values
|
||||
cpu = process(mask, 15, -9)
|
||||
gpu = process(mask.cuda(), 15, -9)
|
||||
assert gpu.device.type == "cuda" and gpu.dtype == torch.float32
|
||||
assert torch.equal(gpu.cpu(), cpu)
|
||||
gpu_filled = process(mask.cuda(), 5, 5, True, 1.0)
|
||||
assert gpu_filled.device.type == "cuda"
|
||||
assert torch.allclose(gpu_filled.cpu(), process(mask, 5, 5, True, 1.0))
|
||||
@@ -0,0 +1,435 @@
|
||||
"""Regression tests: logic, text, batch, utils and mask nodes.
|
||||
|
||||
Each fix has tests that fail on the code before release 1.5.0; a few cases only pin behaviour that must
|
||||
keep working (variables, rounding, plain membership, filled template variables). Sections follow the
|
||||
1.5.0 entries of docs/CHANGELOG.md.
|
||||
"""
|
||||
|
||||
import importlib
|
||||
import random
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
math_nodes = importlib.import_module("NH-Nodes.nodes.logic.math_nodes")
|
||||
logic_nodes = importlib.import_module("NH-Nodes.nodes.logic.logic_nodes")
|
||||
prompt_nodes = importlib.import_module("NH-Nodes.nodes.text.prompt_nodes")
|
||||
text_nodes = importlib.import_module("NH-Nodes.nodes.text.text_nodes")
|
||||
split_lines = importlib.import_module("NH-Nodes.nodes.text.text_split_lines")
|
||||
batch_nodes = importlib.import_module("NH-Nodes.nodes.batch.batch_nodes")
|
||||
pipe_nodes = importlib.import_module("NH-Nodes.nodes.utils.universal_pipe_nodes")
|
||||
utils_nodes = importlib.import_module("NH-Nodes.nodes.utils.utils_nodes")
|
||||
aspect = importlib.import_module("NH-Nodes.nodes.mask.mask_aspect_match")
|
||||
utils_mask = importlib.import_module("NH-Nodes.core.utils_mask")
|
||||
|
||||
|
||||
# --- 1. Math Eval: bad expressions raise, x / 0 only zeroes its own term ---
|
||||
|
||||
def _evaluate(expression, **kwargs):
|
||||
return math_nodes.NH_MathEval().evaluate(expression, **kwargs)
|
||||
|
||||
|
||||
def test_math_eval_keeps_variables_functions_and_rounding():
|
||||
assert _evaluate("a + b * 2", a=1.0, b=3.0) == (7.0, 7, "7.0")
|
||||
assert _evaluate("clamp(a * 10, b, c) + max(d, 1)", a=0.5, b=1.0, c=3.0, d=0.0) == (4.0, 4, "4.0")
|
||||
assert _evaluate("a / 3", a=10.0, round_to=2) == (3.33, 3, "3.33")
|
||||
assert _evaluate("a", a=2.7, round_to=0) == (3.0, 3, "3.0") # result_int is int() of the rounded value
|
||||
assert _evaluate("a", a=2.7) == (2.7, 2, "2.7") # round_to = -1: no rounding
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"expression",
|
||||
["", "1 +", "unknown + 1", "__import__('os')", "'a' * 3", "(-8) ** 0.5", "1e308 * 10", "9 ** 9 ** 9", "sqrt(-1)"],
|
||||
)
|
||||
def test_math_eval_invalid_expression_raises_value_error(expression):
|
||||
with pytest.raises(ValueError):
|
||||
_evaluate(expression)
|
||||
|
||||
|
||||
def test_math_eval_division_and_modulo_by_zero_only_zero_their_own_term(capsys):
|
||||
node = math_nodes.NH_MathEval()
|
||||
assert node.evaluate("a + b / 0", a=1.0, b=5.0)[0] == 1.0
|
||||
assert node.evaluate("a + b // 0", a=1.0, b=5.0)[0] == 1.0 # used to zero the whole result
|
||||
assert node.evaluate("a + b % 0", a=1.0, b=5.0)[0] == 1.0
|
||||
assert capsys.readouterr().out.count("[NH-Nodes] MathEval warning: division by zero") == 3
|
||||
|
||||
|
||||
# --- 2. Random Choice: weight N belongs to input_N, the seed is used as given ---
|
||||
|
||||
def test_random_choice_weight_belongs_to_its_input_slot():
|
||||
node = math_nodes.NH_RandomChoice()
|
||||
# input_0 is not connected: input_1 keeps the weight of slot 1 and input_2 the weight of slot 2.
|
||||
assert node.choose(1, "1,1,10", input_1="B", input_2="C")[2] == "9.1%, 90.9%"
|
||||
picks = [node.choose(seed, "1,1,10", input_1="B", input_2="C")[1] for seed in range(1, 401)]
|
||||
assert picks.count(2) > 300 # about 91 % of 400; matching weights to connected inputs gave 50 %
|
||||
# A slot without a weight counts as 1: input_3 is the 4th weight, not the 2nd.
|
||||
assert node.choose(1, "1,2", input_0="A", input_3="D")[2] == "50.0%, 50.0%"
|
||||
|
||||
|
||||
def test_random_choice_zero_weight_input_is_never_picked():
|
||||
node = math_nodes.NH_RandomChoice()
|
||||
picks = {node.choose(seed, "0,1,1", input_0="A", input_1="B", input_2="C")[1] for seed in range(200)}
|
||||
assert picks == {1, 2}
|
||||
|
||||
|
||||
def test_random_choice_seed_zero_is_deterministic():
|
||||
node = math_nodes.NH_RandomChoice()
|
||||
inputs = {f"input_{index}": chr(ord("A") + index) for index in range(6)}
|
||||
weights = "1,1,1,1,1,1"
|
||||
assert len({node.choose(0, weights, **inputs)[1] for _ in range(30)}) == 1
|
||||
assert node.choose(5, weights, **inputs) == node.choose(5, weights, **inputs)
|
||||
assert {node.choose(seed, weights, **inputs)[1] for seed in range(60)} == set(range(6))
|
||||
|
||||
|
||||
# --- 3. Prompt Template: an empty variable never glues its neighbours together ---
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"template,names,values,expected,missing",
|
||||
[
|
||||
("red dress, {style}, 8k photo", "style", {}, "red dress, 8k photo", "style"),
|
||||
("{style}, red dress, 8k photo", "style", {}, "red dress, 8k photo", "style"), # at the start
|
||||
("red dress, 8k photo, {style}", "style", {}, "red dress, 8k photo", "style"), # at the end
|
||||
("red dress, {a}, {b}, 8k photo", "a,b", {}, "red dress, 8k photo", "a, b"), # consecutive
|
||||
("{a}, {b}, red dress", "a,b", {}, "red dress", "a, b"),
|
||||
("red dress, {a}, {b}", "a,b", {}, "red dress", "a, b"),
|
||||
("{a}, {b}", "a,b", {}, "", "a, b"),
|
||||
("{a}", "a", {}, "", "a"),
|
||||
("a {color} dress", "color", {}, "a dress", "color"), # inside a sentence
|
||||
("a {color} {garment}, {style}, 8k photo", "color,garment,style", {"var2": "dress"}, "a dress, 8k photo", "color, style"),
|
||||
("a red {garment}, {style}, 8k photo", "garment,style", {}, "a red, 8k photo", "garment, style"),
|
||||
("line one, {x}\nline two, {y}, tail\n{x}, last", "x,y", {}, "line one\nline two, tail\nlast", "x, y"),
|
||||
("a, {x}\r\nb", "x", {}, "a\r\nb", "x"),
|
||||
("a, {x}, b", "x", {"var1": None}, "a, b", "x"), # an unconnected input counts as empty
|
||||
],
|
||||
)
|
||||
def test_prompt_template_empty_variable_keeps_neighbours_apart(template, names, values, expected, missing):
|
||||
assert prompt_nodes.NH_PromptTemplate().render(template, names, **values) == (expected, missing)
|
||||
|
||||
|
||||
def test_prompt_template_filled_variables_are_substituted():
|
||||
node = prompt_nodes.NH_PromptTemplate()
|
||||
result = node.render("a {color} {garment}, {style}, 8k photo", "color,garment,style", var1="red", var2="dress", var3="studio")
|
||||
assert result == ("a red dress, studio, 8k photo", "")
|
||||
|
||||
|
||||
# --- 4. Seeds are used as given, so seed 0 is reproducible too ---
|
||||
|
||||
@pytest.mark.parametrize("seed", [0, 3])
|
||||
def test_prompt_scheduler_random_mode_uses_the_seed_literally(seed):
|
||||
node = prompt_nodes.NH_PromptScheduler()
|
||||
prompts = "\n".join("abcdef")
|
||||
for step in range(8):
|
||||
index = node.schedule(prompts, step, 10, "random", seed=seed)[2]
|
||||
assert index == random.Random(seed + step).randint(0, 5)
|
||||
assert index == node.schedule(prompts, step, 10, "random", seed=seed)[2]
|
||||
|
||||
|
||||
def test_text_split_lines_seed_zero_is_deterministic():
|
||||
node = split_lines.NH_TextSplitLines()
|
||||
text = "\n".join("abcdef")
|
||||
expected = list("abcdef")
|
||||
random.Random(0).shuffle(expected)
|
||||
assert node.split_lines(text, 0, "shuffle", 0)[0] == expected
|
||||
assert len({tuple(node.split_lines(text, 0, "shuffle", 0)[0]) for _ in range(20)}) == 1
|
||||
# index -1 picks a random line with the same seeded generator
|
||||
picks = {node.split_lines(text, -1, "sequential", 0)[1] for _ in range(30)}
|
||||
assert picks == {text.split("\n")[random.Random(0).randint(0, 5)]}
|
||||
|
||||
|
||||
def test_text_random_line_seed_zero_is_deterministic():
|
||||
node = split_lines.NH_TextRandomLine()
|
||||
text = "\n".join("abcdef")
|
||||
expected = random.Random(0).randint(0, 5)
|
||||
assert [node.pick(text, 0)[1] for _ in range(30)] == [expected] * 30
|
||||
|
||||
|
||||
# --- 5. Regex Extract: a group that did not take part is an empty string ---
|
||||
|
||||
def test_regex_extract_unmatched_optional_group_is_empty():
|
||||
node = text_nodes.NH_RegexExtract()
|
||||
assert node.extract("b", r"(a)|(b)", "match") == ("b", True, "\nb", 1)
|
||||
assert node.extract("x-y", r"(x)(-)?(z)?", "match") == ("x-", True, "x\n-\n", 1)
|
||||
assert node.extract("ab", r"(a)(c)?", "match") == ("a", True, "a\n", 1)
|
||||
assert node.extract("ab", r"a", "match") == ("a", True, "a", 1) # no group: the whole match
|
||||
|
||||
|
||||
# --- 6. Pack Universal: the package is indexed by slot ---
|
||||
|
||||
def test_pack_universal_keeps_slot_positions():
|
||||
pack = pipe_nodes.NH_PackUniversal().pack(input_1="ONE", input_3="THREE")[0]
|
||||
unpack = pipe_nodes.NH_UnpackUniversal()
|
||||
assert unpack.unpack(0, package=pack) == (None,) # input_0 was left empty
|
||||
assert unpack.unpack(1, package=pack) == ("ONE",)
|
||||
assert unpack.unpack(2, package=pack) == (None,)
|
||||
assert unpack.unpack(3, package=pack) == ("THREE",)
|
||||
assert unpack.unpack(9, package=pack) == (None,)
|
||||
assert unpack.unpack(1, pipe=pack) == ("ONE",) # the legacy `pipe` input reads the same package
|
||||
|
||||
|
||||
def test_pack_universal_ignores_argument_order_and_keeps_falsy_values():
|
||||
pack = pipe_nodes.NH_PackUniversal().pack(input_2="C", input_0=0, input_1="")[0]
|
||||
unpack = pipe_nodes.NH_UnpackUniversal()
|
||||
assert [unpack.unpack(index, package=pack)[0] for index in range(4)] == [0, "", "C", None]
|
||||
assert pipe_nodes.NH_UnpackUniversal().unpack(0, package=pipe_nodes.NH_PackUniversal().pack()[0]) == (None,)
|
||||
|
||||
|
||||
# --- 7. Compare "in": plain membership / substring ---
|
||||
|
||||
@pytest.mark.parametrize("needle,haystack", [(13, "1024x768 3"), ("13", "1024x768 3"), (15, "1,05"), (25, "2 and 5")])
|
||||
def test_compare_in_is_not_a_subsequence_match(needle, haystack):
|
||||
assert logic_nodes.NH_Compare().compare(needle, haystack, "in")[0] is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"needle,haystack",
|
||||
[
|
||||
("cat", "concatenate"),
|
||||
(13, "size 1024x13"),
|
||||
(7, [1, 7]),
|
||||
("7", [1, 7]),
|
||||
(7, ["1", "7"]),
|
||||
("k", {"k": 1}),
|
||||
(1, {"1": "x"}),
|
||||
("", "anything"),
|
||||
],
|
||||
)
|
||||
def test_compare_in_finds_substrings_and_members(needle, haystack):
|
||||
assert logic_nodes.NH_Compare().compare(needle, haystack, "in")[0] is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize("needle,haystack", [("dog", "concatenate"), ("CAT", "concatenate"), (4, [1, 7]), ("z", {"k": 1})])
|
||||
def test_compare_in_rejects_non_members(needle, haystack):
|
||||
assert logic_nodes.NH_Compare().compare(needle, haystack, "in")[0] is False
|
||||
|
||||
|
||||
# --- 8. Batch Merge: channel counts follow batch_a ---
|
||||
|
||||
def test_batch_merge_aligns_channels_to_batch_a():
|
||||
node = batch_nodes.NH_BatchMerge()
|
||||
rgb = torch.full((2, 4, 4, 3), 0.25)
|
||||
rgba = torch.full((1, 4, 4, 4), 0.75)
|
||||
rgba[..., 3] = 0.5
|
||||
|
||||
merged, count = node.merge(rgb, rgba) # RGB first: the alpha channel is dropped
|
||||
assert merged.shape == (3, 4, 4, 3) and count == 3
|
||||
assert torch.equal(merged[:2], rgb) and torch.equal(merged[2], rgba[0, ..., :3])
|
||||
|
||||
merged, count = node.merge(rgba, rgb) # RGBA first: RGB images get a fully opaque alpha
|
||||
assert merged.shape == (3, 4, 4, 4) and count == 3
|
||||
assert torch.equal(merged[0], rgba[0])
|
||||
assert torch.equal(merged[1:, ..., :3], rgb)
|
||||
assert torch.all(merged[1:, ..., 3] == 1.0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("resize_mode", ["stretch", "crop", "pad"])
|
||||
def test_batch_merge_handles_size_and_channel_mismatch_together(resize_mode):
|
||||
node = batch_nodes.NH_BatchMerge()
|
||||
rgba = torch.rand((1, 4, 4, 4))
|
||||
rgb_big = torch.rand((2, 8, 6, 3))
|
||||
rgb_same = torch.rand((1, 4, 4, 3))
|
||||
merged, count = node.merge(rgba, rgb_big, rgb_same, resize_mode=resize_mode)
|
||||
assert merged.shape == (4, 4, 4, 4) and count == 4
|
||||
if resize_mode == "pad":
|
||||
# 8x6 fits as 4x3: the picture is opaque, the 1 px letterbox column is transparent (like RGBA input)
|
||||
assert torch.all(merged[1:3, :, :3, 3] == 1.0) and torch.all(merged[1:3, :, 3, 3] == 0.0)
|
||||
assert torch.all(merged[3, ..., 3] == 1.0)
|
||||
else:
|
||||
assert torch.all(merged[1:, ..., 3] == 1.0)
|
||||
|
||||
|
||||
def test_batch_merge_rejects_unsupported_channel_counts_clearly():
|
||||
node = batch_nodes.NH_BatchMerge()
|
||||
with pytest.raises(ValueError, match="channel"):
|
||||
node.merge(torch.rand((1, 4, 4, 3)), torch.rand((1, 4, 4, 2)))
|
||||
|
||||
|
||||
# --- 9. NH Mask Aspect Ratio Match ---
|
||||
|
||||
def _canvas(*boxes, size=100):
|
||||
"""MASK batch [B, size, size]; each box is (x1, y1, x2, y2) with an exclusive end."""
|
||||
masks = torch.zeros((len(boxes), size, size))
|
||||
for index, (x1, y1, x2, y2) in enumerate(boxes):
|
||||
masks[index, y1:y2, x1:x2] = 1.0
|
||||
return masks
|
||||
|
||||
|
||||
def _bounds(mask):
|
||||
"""(x1, y1, x2, y2), inclusive, of the pixels above 0.5."""
|
||||
rows, cols = torch.nonzero(mask > 0.5, as_tuple=True)
|
||||
return int(cols.min()), int(rows.min()), int(cols.max()), int(rows.max())
|
||||
|
||||
|
||||
TALL = _canvas((20, 10, 50, 70)) # 30 x 60 -> ratio 0.5
|
||||
WIDE = _canvas((10, 20, 70, 50)) # 60 x 30 -> ratio 2.0
|
||||
|
||||
|
||||
def _match(target, box, mode):
|
||||
return aspect.NH_MaskAspectRatioMatch().match_aspect_ratio(target, box, mode)[0]
|
||||
|
||||
|
||||
def test_aspect_match_pad_grows_the_box_to_the_requested_size():
|
||||
box = _canvas((30, 30, 50, 50)) # 20 x 20 in the middle
|
||||
out = _match(TALL, box, "pad")
|
||||
assert out.shape == box.shape
|
||||
assert _bounds(out[0]) == (30, 20, 49, 59) # 20 x 40, same centre, nothing shifted
|
||||
assert float(out.sum()) == 20 * 40
|
||||
|
||||
|
||||
def test_aspect_match_pad_keeps_the_original_pixels_where_they_were():
|
||||
box = _canvas((30, 30, 50, 50))
|
||||
box[0, 38:42, 38:42] = 0.0 # a hole must stay exactly there, the padding is added around it
|
||||
out = _match(TALL, box, "pad")[0]
|
||||
assert float(out[38:42, 38:42].sum()) == 0.0
|
||||
assert torch.equal(out[30:50, 30:50], box[0, 30:50, 30:50])
|
||||
assert float(out[20:30, 30:50].min()) == 1.0 and float(out[50:60, 30:50].min()) == 1.0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode,expected", [("stretch", (0, 40, 39, 59)), ("pad", (0, 40, 39, 59)), ("crop", (0, 45, 19, 54))])
|
||||
def test_aspect_match_box_touching_the_left_edge(mode, expected):
|
||||
box = _canvas((0, 40, 20, 60)) # x1 = 0 used to crash pad and crop (negative paste offset)
|
||||
out = _match(WIDE, box, mode)
|
||||
assert out.shape == box.shape
|
||||
# The widened box slides inwards instead of being cut by the edge, so it keeps the requested ratio.
|
||||
assert _bounds(out[0]) == expected
|
||||
|
||||
|
||||
def test_aspect_match_stretch_is_centred_and_scaled_not_shifted():
|
||||
box = _canvas((30, 30, 50, 50))
|
||||
box[0, 38:42, 38:42] = 0.0
|
||||
out = _match(TALL, box, "stretch")[0]
|
||||
assert _bounds(out) == (30, 20, 49, 59) # the unchanged width must not move by a pixel
|
||||
assert float(out.sum()) == 20 * 40 - 4 * 8 # the 4 x 4 hole is stretched to 4 x 8
|
||||
|
||||
|
||||
def test_aspect_match_stretch_larger_than_the_image_shows_the_middle_of_the_content():
|
||||
box = _canvas((40, 40, 60, 60))
|
||||
box[0, 45:55, 40:44] = 0.0 # a hole on the far left of the box
|
||||
extreme = _canvas((0, 0, 1000, 10), size=1100) # ratio 100: the stretched box would be 2000 px wide
|
||||
out = _match(extreme, box, "stretch")[0]
|
||||
assert out.shape == box.shape[1:]
|
||||
# Only the middle of the stretched content fits into the 100 px wide image, so the hole is not in view.
|
||||
assert _bounds(out) == (0, 40, 99, 59)
|
||||
assert float(out.sum()) == 100 * 20
|
||||
|
||||
|
||||
def test_aspect_match_crop_trims_around_the_centre_without_shifting():
|
||||
out = _match(TALL, _canvas((30, 30, 50, 50)), "crop")
|
||||
assert _bounds(out[0]) == (35, 30, 44, 49) # 10 x 20 taken from the middle of 20 x 20
|
||||
assert float(out.sum()) == 10 * 20
|
||||
|
||||
|
||||
def test_aspect_match_processes_the_whole_batch():
|
||||
boxes = _canvas((30, 30, 50, 50), (10, 10, 40, 40), (0, 0, 0, 0)) # the last one is empty
|
||||
out = _match(TALL, boxes, "pad")
|
||||
assert out.shape == boxes.shape
|
||||
assert _bounds(out[0]) == (30, 20, 49, 59)
|
||||
assert _bounds(out[1]) == (10, 0, 39, 59) # 30 x 60; grown upwards by sliding back inside the canvas
|
||||
assert float(out[2].sum()) == 0.0 # an empty mask is returned as it is
|
||||
|
||||
|
||||
def test_aspect_match_pairs_target_and_box_by_batch_index():
|
||||
targets = torch.cat([TALL, WIDE])
|
||||
boxes = _canvas((30, 30, 50, 50), (30, 30, 50, 50))
|
||||
out = _match(targets, boxes, "pad")
|
||||
assert _bounds(out[0]) == (30, 20, 49, 59) # tall target -> 20 x 40
|
||||
assert _bounds(out[1]) == (20, 30, 59, 49) # wide target -> 40 x 20
|
||||
assert _match(TALL, boxes, "pad").shape == boxes.shape # a single target serves the whole batch
|
||||
|
||||
|
||||
def test_aspect_match_accepts_masks_without_batch_dimension():
|
||||
out = _match(TALL[0], _canvas((30, 30, 50, 50))[0], "pad")
|
||||
assert out.shape == (1, 100, 100)
|
||||
assert _bounds(out[0]) == (30, 20, 49, 59)
|
||||
|
||||
|
||||
# --- 10. core/utils_mask.py (FitDiT garment masks) ---
|
||||
|
||||
def test_remove_small_clears_only_the_pixels_of_small_regions():
|
||||
mask = np.zeros((60, 60), np.uint8)
|
||||
mask[11:50, 11:50] = 255 # large region, 39 x 39
|
||||
mask[8, 8:20] = 255 # thin L whose bounding box (rows/cols 8..19) overlaps the corner of the square
|
||||
mask[8:20, 8] = 255
|
||||
cleaned = utils_mask.remove_small(mask, min_area=100)
|
||||
assert np.array_equal(cleaned[11:50, 11:50], mask[11:50, 11:50]) # nothing cut out of the large region
|
||||
assert int((cleaned > 0).sum()) == 39 * 39
|
||||
assert not cleaned[8, 8:20].any() and not cleaned[8:20, 8].any()
|
||||
|
||||
|
||||
POSE_W, POSE_H = 200, 300
|
||||
|
||||
|
||||
def _pose(points):
|
||||
"""pose_data [134, 2] like the VTON node builds it: DWPose marks an undetected keypoint -1 and the node
|
||||
scales x by the width and y by the height, so a missing keypoint is (-W, -H).
|
||||
Rows: body 0-17, foot1 18-20 (right), foot2 21-23 (left), face 24-91, hands1 92-112, hands2 113-133."""
|
||||
pose = np.empty((134, 2), np.float32)
|
||||
pose[:] = (-POSE_W, -POSE_H)
|
||||
for row, xy in points.items():
|
||||
pose[row] = xy
|
||||
return pose
|
||||
|
||||
|
||||
def _garment_mask(category, pose):
|
||||
parse = Image.new("L", (POSE_W, POSE_H), 0) # no clothes, no head: only the keypoints decide
|
||||
mask, _ = utils_mask.get_mask_location(category, parse, pose, POSE_W, POSE_H, 0, 0, 0, 0)
|
||||
return mask
|
||||
|
||||
|
||||
def _rect(mask):
|
||||
"""(x1, y1, x2, y2), inclusive, of the white pixels."""
|
||||
rows, cols = np.nonzero(np.array(mask))
|
||||
return int(cols.min()), int(rows.min()), int(cols.max()), int(rows.max())
|
||||
|
||||
|
||||
def test_upper_mask_ignores_keypoints_that_were_not_detected():
|
||||
# Left shoulder, left hip and every hand/face point are missing (negative).
|
||||
pose = _pose({2: (80, 100), 3: (70, 150), 4: (65, 190), 6: (130, 150), 7: (135, 190), 8: (85, 200)})
|
||||
# x: 65 - 25 .. 135 + 25, y: 100 - 20 .. 200 + 10. A missing point used to pull the box to the image top/left.
|
||||
assert _rect(_garment_mask("Upper-body", pose)) == (40, 80, 160, 210)
|
||||
|
||||
|
||||
def test_upper_mask_uses_only_detected_face_points():
|
||||
pose = _pose({2: (80, 100), 5: (120, 100), 8: (85, 200), 11: (115, 200), 29: (100, 60)}) # face point 5 only
|
||||
assert _rect(_garment_mask("Upper-body", pose)) == (55, 60, 145, 210)
|
||||
|
||||
|
||||
def test_lower_mask_starts_at_the_detected_hip():
|
||||
# Left hip, left knee and the left foot are missing; the right hip is at y = 150.
|
||||
pose = _pose({8: (85, 150), 9: (84, 220), 10: (88, 240), 13: (115, 245), 20: (92, 262)})
|
||||
# top = 150 - 50 (no garment found), bottom = foot 262 + 10. min(right hip, missing left hip) put it at 0.
|
||||
assert _rect(_garment_mask("Lower-body", pose)) == (54, 100, 145, 272)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("left_foot", [(0, 0), (-POSE_W, -POSE_H)])
|
||||
def test_lower_mask_bottom_uses_the_right_foot_when_the_left_foot_is_missing(left_foot):
|
||||
pose = _pose({8: (85, 150), 11: (115, 150), 20: (92, 262), 23: left_foot}) # ankles not detected
|
||||
assert _rect(_garment_mask("Lower-body", pose)) == (55, 100, 145, 272)
|
||||
|
||||
|
||||
def test_dresses_mask_ignores_keypoints_that_were_not_detected():
|
||||
pose = _pose({
|
||||
2: (80, 100), 3: (70, 150), 6: (130, 150), # right shoulder side; left shoulder and both wrists missing
|
||||
8: (85, 170), 9: (84, 220), 10: (88, 250), 11: (115, 170), 12: (116, 220), 13: (118, 252),
|
||||
18: (95, 262), 21: (120, 264),
|
||||
})
|
||||
# x: 70 - 50 .. 130 + 50, y: 100 - 20 .. 264 + 10; the missing left shoulder used to put the top at 0.
|
||||
assert _rect(_garment_mask("Dresses", pose)) == (20, 80, 180, 274)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("category", ["Upper-body", "Lower-body", "Dresses"])
|
||||
def test_mask_without_any_detected_keypoint_falls_back_to_the_whole_image(category):
|
||||
mask = _garment_mask(category, _pose({}))
|
||||
assert np.array(mask).all()
|
||||
|
||||
|
||||
# --- 11. Universal Slider Builder cannot work, so it is hidden from the search ---
|
||||
|
||||
def test_universal_slider_builder_is_deprecated_but_still_loads_old_workflows():
|
||||
node_class = utils_nodes.NH_UniversalSliderBuilder
|
||||
assert node_class.DEPRECATED is True
|
||||
assert utils_nodes.NODE_CLASS_MAPPINGS["NH_UniversalSliderBuilder"] is node_class
|
||||
assert utils_nodes.NODE_DISPLAY_NAME_MAPPINGS["NH_UniversalSliderBuilder"] == "Universal Slider Builder (NH)"
|
||||
assert list(node_class.INPUT_TYPES()["required"]) == ["config"]
|
||||
assert not getattr(utils_nodes.NH_BooleanSwitch, "DEPRECATED", False)
|
||||
@@ -0,0 +1,141 @@
|
||||
"""core.mask_ops.mask_bbox and the mask nodes built on it (Mask Properties, Create Box Mask)."""
|
||||
|
||||
import importlib
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
mask_ops = importlib.import_module("NH-Nodes.core.mask_ops")
|
||||
properties = importlib.import_module("NH-Nodes.nodes.mask.mask_properties")
|
||||
box_mask = importlib.import_module("NH-Nodes.nodes.mask.mask_to_bbox")
|
||||
|
||||
mask_bbox = mask_ops.mask_bbox
|
||||
|
||||
|
||||
def _mask(h, w, *boxes):
|
||||
"""float32 [H, W] mask; each box is (x1, y1, x2, y2) with an exclusive end."""
|
||||
mask = np.zeros((h, w), np.float32)
|
||||
for x1, y1, x2, y2 in boxes:
|
||||
mask[y1:y2, x1:x2] = 1.0
|
||||
return mask
|
||||
|
||||
|
||||
# --- mask_bbox ---
|
||||
|
||||
def test_bbox_is_inclusive_and_ordered_x1_y1_x2_y2():
|
||||
assert mask_bbox(_mask(10, 12, (3, 2, 9, 5))) == (3, 2, 8, 4) # width 6, height 3
|
||||
|
||||
|
||||
def test_bbox_covers_several_blobs():
|
||||
assert mask_bbox(_mask(20, 30, (2, 3, 5, 6), (20, 15, 27, 18))) == (2, 3, 26, 17)
|
||||
|
||||
|
||||
def test_bbox_of_an_empty_mask_is_none():
|
||||
assert mask_bbox(np.zeros((6, 7), np.float32)) is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", [(0, 5), (5, 0), (0, 0)])
|
||||
def test_bbox_of_a_zero_sized_mask_is_none(shape):
|
||||
assert mask_bbox(np.zeros(shape, np.float32)) is None
|
||||
|
||||
|
||||
def test_bbox_of_a_single_pixel_and_of_a_full_mask():
|
||||
single = np.zeros((9, 11), np.float32)
|
||||
single[7, 4] = 1.0
|
||||
assert mask_bbox(single) == (4, 7, 4, 7)
|
||||
assert mask_bbox(np.ones((9, 11), np.float32)) == (0, 0, 10, 8)
|
||||
|
||||
|
||||
def test_threshold_is_strict_and_defaults_to_half():
|
||||
half = np.full((4, 4), 0.5, np.float32)
|
||||
assert mask_bbox(half) is None # exactly 0.5 is outside, as in every NH mask node
|
||||
assert mask_bbox(np.full((4, 4), np.nextafter(np.float32(0.5), np.float32(1)), np.float32)) == (0, 0, 3, 3)
|
||||
|
||||
soft = _mask(8, 8, (2, 2, 6, 6)) * 0.4
|
||||
assert mask_bbox(soft) is None
|
||||
assert mask_bbox(soft, threshold=0.25) == (2, 2, 5, 5)
|
||||
assert mask_bbox(soft, threshold=0.0) == (2, 2, 5, 5) # "any non-zero pixel"
|
||||
|
||||
|
||||
def test_nan_pixels_are_outside():
|
||||
mask = _mask(5, 5, (1, 1, 3, 3))
|
||||
mask[0, 0] = np.nan
|
||||
assert mask_bbox(mask) == (1, 1, 2, 2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", [(), (7,), (4, 5, 1), (2, 4, 5)])
|
||||
def test_bbox_rejects_masks_that_are_not_2d(shape):
|
||||
with pytest.raises(ValueError, match="2-D"):
|
||||
mask_bbox(np.ones(shape, np.float32))
|
||||
|
||||
|
||||
def test_bbox_is_four_python_ints():
|
||||
bbox = mask_bbox(_mask(6, 6, (1, 2, 4, 5)))
|
||||
assert all(type(value) is int for value in bbox)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [np.float16, np.float64, np.uint8, np.bool_])
|
||||
def test_bbox_accepts_other_array_types(dtype):
|
||||
mask = _mask(8, 9, (2, 1, 7, 4)).astype(dtype)
|
||||
if dtype is np.uint8:
|
||||
mask *= 255
|
||||
assert mask_bbox(mask) == (2, 1, 6, 3)
|
||||
|
||||
|
||||
def test_bbox_matches_a_brute_force_search():
|
||||
rng = np.random.default_rng(5)
|
||||
for _ in range(200):
|
||||
h, w = rng.integers(1, 30, 2)
|
||||
mask = (rng.random((h, w)) * (rng.random((h, w)) < rng.choice([0.0, 0.03, 0.4]))).astype(np.float32)
|
||||
points = [(x, y) for y in range(h) for x in range(w) if mask[y, x] > 0.5]
|
||||
expected = None if not points else (
|
||||
min(x for x, _ in points), min(y for _, y in points), max(x for x, _ in points), max(y for _, y in points)
|
||||
)
|
||||
assert mask_bbox(mask) == expected
|
||||
|
||||
|
||||
# --- Mask Properties ---
|
||||
|
||||
def _properties(mask):
|
||||
return properties.NH_MaskProperties().get_properties(torch.from_numpy(mask))
|
||||
|
||||
|
||||
def test_properties_report_size_corners_and_bbox():
|
||||
mask = torch.from_numpy(_mask(10, 12, (3, 2, 9, 5))).unsqueeze(0)
|
||||
|
||||
width, height, x1, y1, x2, y2, bbox = properties.NH_MaskProperties().get_properties(mask)
|
||||
|
||||
assert (width, height, x1, y1, x2, y2) == (6, 3, 3, 2, 8, 4)
|
||||
assert bbox == (3, 2, 8, 4)
|
||||
|
||||
|
||||
def test_properties_of_an_empty_mask_are_all_zero():
|
||||
assert _properties(np.zeros((1, 6, 7), np.float32)) == (0, 0, 0, 0, 0, 0, (0, 0, 0, 0))
|
||||
|
||||
|
||||
def test_properties_only_look_at_the_first_batch_item():
|
||||
batch = torch.from_numpy(np.stack([_mask(8, 8, (1, 1, 3, 3)), _mask(8, 8, (4, 4, 8, 8))]))
|
||||
assert _properties(batch.numpy())[:6] == (2, 2, 1, 1, 2, 2)
|
||||
|
||||
|
||||
# --- Create Box Mask ---
|
||||
|
||||
def test_box_mask_fills_the_bounding_box_of_every_batch_item():
|
||||
first = _mask(10, 10, (2, 2, 4, 3), (6, 5, 8, 8))
|
||||
second = np.zeros((10, 10), np.float32) # empty: stays empty
|
||||
third = _mask(10, 10, (0, 0, 1, 1))
|
||||
masks = torch.from_numpy(np.stack([first, second, third]))
|
||||
|
||||
(result,) = box_mask.NH_MaskToBBox().create_bbox_mask(masks)
|
||||
|
||||
assert result.shape == masks.shape and result.dtype == masks.dtype
|
||||
assert torch.equal(result[0], torch.from_numpy(_mask(10, 10, (2, 2, 8, 8))))
|
||||
assert float(result[1].sum()) == 0.0
|
||||
assert torch.equal(result[2], torch.from_numpy(_mask(10, 10, (0, 0, 1, 1))))
|
||||
|
||||
|
||||
def test_box_mask_ignores_values_at_or_below_half():
|
||||
soft = torch.from_numpy(_mask(6, 6, (1, 1, 5, 5)) * 0.5).unsqueeze(0)
|
||||
(result,) = box_mask.NH_MaskToBBox().create_bbox_mask(soft)
|
||||
assert float(result.sum()) == 0.0
|
||||
@@ -0,0 +1,91 @@
|
||||
"""Regression tests for Math Eval scalars, Batch Merge padding and VTON garment boxes."""
|
||||
|
||||
import importlib
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
math_nodes = importlib.import_module("NH-Nodes.nodes.logic.math_nodes")
|
||||
batch_nodes = importlib.import_module("NH-Nodes.nodes.batch.batch_nodes")
|
||||
utils_mask = importlib.import_module("NH-Nodes.core.utils_mask")
|
||||
|
||||
|
||||
# --- Math Eval accepts the scalar types upstream nodes actually send ---
|
||||
|
||||
@pytest.mark.parametrize("value", [np.float32(2), np.int64(2), torch.tensor(2.0), True])
|
||||
def test_math_eval_accepts_numpy_torch_and_bool_scalars(value):
|
||||
result_float, result_int, result_string = math_nodes.NH_MathEval().evaluate("a + b", a=value, b=1.0)
|
||||
assert result_float == float(value) + 1.0 and result_int == int(result_float) and result_string == str(result_float)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("expression", ["foo * 2", "__import__('os')", "a < b", "''"])
|
||||
def test_math_eval_errors_always_name_the_expression(expression):
|
||||
with pytest.raises(ValueError) as error:
|
||||
math_nodes.NH_MathEval().evaluate(expression, a=1.0, b=2.0)
|
||||
assert repr(expression) in str(error.value)
|
||||
|
||||
|
||||
# --- Batch Merge ---
|
||||
|
||||
def test_batch_merge_pad_never_rounds_a_side_to_zero():
|
||||
batch_a = torch.zeros((1, 64, 64, 3))
|
||||
thin = torch.ones((1, 1, 1000, 3)) # scale 0.064 -> 0.064 px tall before the fix
|
||||
result, count = batch_nodes.NH_BatchMerge().merge(batch_a, thin, resize_mode="pad")
|
||||
assert result.shape == (2, 64, 64, 3) and count == 2
|
||||
|
||||
|
||||
def test_batch_merge_pad_bars_get_the_same_alpha_for_rgb_and_rgba_inputs():
|
||||
batch_a = torch.ones((1, 8, 8, 4))
|
||||
rgb = torch.ones((1, 4, 8, 3)) # letterboxed: 2 px bars on top and bottom
|
||||
rgba = torch.ones((1, 4, 8, 4))
|
||||
merged, _ = batch_nodes.NH_BatchMerge().merge(batch_a, rgb, rgba, resize_mode="pad")
|
||||
bars_rgb, bars_rgba = merged[1, :2, :, 3], merged[2, :2, :, 3]
|
||||
assert torch.equal(bars_rgb, bars_rgba) and float(bars_rgb.max()) == 0.0
|
||||
assert float(merged[1, 2:6, :, 3].min()) == 1.0 # the picture itself is opaque
|
||||
|
||||
|
||||
# --- VTON garment boxes when a whole keypoint group is missing ---
|
||||
|
||||
W, H = 200, 300
|
||||
|
||||
|
||||
def _pose(points):
|
||||
pose = np.empty((134, 2), np.float32)
|
||||
pose[:] = (-W, -H) # DWPose marks undetected keypoints -1, the node scales them by W/H
|
||||
for row, xy in points.items():
|
||||
pose[row] = xy
|
||||
return pose
|
||||
|
||||
|
||||
def _rect(mask):
|
||||
rows, cols = np.nonzero(np.array(mask))
|
||||
return int(cols.min()), int(rows.min()), int(cols.max()), int(rows.max())
|
||||
|
||||
|
||||
def test_lower_mask_uses_the_trousers_when_both_hips_are_missing():
|
||||
parse = np.zeros((H, W), np.uint8)
|
||||
parse[150:280, 70:130] = 6 # trousers
|
||||
pose = _pose({9: (84, 220), 10: (88, 250), 12: (116, 220), 13: (118, 252)}) # knees + ankles only
|
||||
mask, _ = utils_mask.get_mask_location("Lower-body", Image.fromarray(parse), pose, W, H, 0, 0, 0, 0)
|
||||
x1, y1, x2, y2 = _rect(mask)
|
||||
assert 140 <= y1 <= 150 # trousers top minus the 5 px pad; the missing hips used to put it at 0
|
||||
assert x1 < 70 and x2 >= 129 and y2 >= 279
|
||||
|
||||
|
||||
def test_upper_mask_with_only_face_points_stays_around_the_shirt():
|
||||
parse = np.zeros((H, W), np.uint8)
|
||||
parse[90:200, 60:140] = 4 # upper clothes
|
||||
pose = _pose({29: (100, 60)}) # one face point, no body or hand keypoints
|
||||
mask, _ = utils_mask.get_mask_location("Upper-body", Image.fromarray(parse), pose, W, H, 0, 0, 0, 0)
|
||||
x1, y1, x2, y2 = _rect(mask)
|
||||
assert (x1, x2) == (35, 165) # shirt contour 60..140 (exclusive end) plus 25 px padding, not the whole width
|
||||
assert y1 == 60 and y2 < H - 1 # from the face line down to the shirt hem, not the whole height
|
||||
|
||||
|
||||
@pytest.mark.parametrize("category", ["Upper-body", "Lower-body", "Dresses"])
|
||||
def test_no_keypoints_and_no_garment_still_cover_the_whole_image(category):
|
||||
parse = Image.fromarray(np.zeros((H, W), np.uint8))
|
||||
mask, _ = utils_mask.get_mask_location(category, parse, _pose({}), W, H, 0, 0, 0, 0)
|
||||
assert np.array(mask).all()
|
||||
@@ -0,0 +1,74 @@
|
||||
"""Pillow deprecations on the garment-mask and VTON paths.
|
||||
|
||||
Pillow 11.3 deprecated Image.fromarray(array, mode) and warned at every call (12.0 took most of that back and 13.0
|
||||
only rejects a mode that changes the data type). The code reads uint8 label maps with a plain Image.fromarray(array)
|
||||
instead, which behaves the same on every Pillow version and stays quiet.
|
||||
|
||||
Both tests turn any "will be removed in Pillow N" DeprecationWarning into an error, so a `mode=` that comes back
|
||||
fails here.
|
||||
"""
|
||||
|
||||
import importlib
|
||||
import warnings
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
pytest.importorskip("skimage") # core.utils_mask imports it at module level
|
||||
|
||||
utils_mask = importlib.import_module("NH-Nodes.core.utils_mask")
|
||||
vton = importlib.import_module("NH-Nodes.nodes.vton.vton_preprocessor_nodes")
|
||||
|
||||
H, W = 48, 64
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def pillow_deprecations_are_errors():
|
||||
with warnings.catch_warnings():
|
||||
warnings.filterwarnings("error", message=r".*removed in Pillow", category=DeprecationWarning)
|
||||
yield
|
||||
|
||||
|
||||
def _parsing_labels():
|
||||
labels = np.zeros((H, W), np.uint8)
|
||||
labels[5:12, 20:40] = 11 # head
|
||||
labels[12:30, 15:45] = 4 # upper clothes
|
||||
labels[30:44, 18:42] = 5 # trousers
|
||||
labels[44:48, 18:42] = 9 # shoes
|
||||
return labels
|
||||
|
||||
|
||||
@pytest.mark.parametrize("category", ["Upper-body", "Lower-body", "Dresses"])
|
||||
def test_garment_masks_do_not_trip_a_pillow_deprecation(category):
|
||||
pose = np.full((134, 2), -1.0, np.float32) # nothing detected: the masks fall back to the garment contours
|
||||
|
||||
mask, mask_gray = utils_mask.get_mask_location(category, Image.fromarray(_parsing_labels()), pose, W, H, 0, 0, 0, 0)
|
||||
|
||||
assert mask.mode == mask_gray.mode == "L" and mask.size == (W, H)
|
||||
|
||||
|
||||
class _Pose:
|
||||
def __call__(self, bgr):
|
||||
return np.zeros(bgr.shape, np.uint8), None, None, np.full((1, 134, 2), 0.5, np.float32)
|
||||
|
||||
|
||||
class _Parsing:
|
||||
def __call__(self, image):
|
||||
labels = Image.fromarray(_parsing_labels())
|
||||
return labels, labels
|
||||
|
||||
|
||||
@pytest.mark.parametrize("category", ["Upper-body", "Lower-body", "Dresses"])
|
||||
def test_vton_node_builds_its_coloured_parsing_image_without_a_pillow_deprecation(monkeypatch, category):
|
||||
monkeypatch.setattr(vton, "LOADED_VTON_MODELS", {"dwprocessor": _Pose(), "parsing_model": _Parsing(), "device": "cpu"})
|
||||
|
||||
outputs = vton.NH_VTonUltimateProcessor().process(torch.full((1, H, W, 3), 0.5), category, 2, False, True, True, "cpu")
|
||||
|
||||
parsing_image = outputs[3]
|
||||
assert parsing_image.shape == (1, H, W, 3)
|
||||
# The palette is applied label by label: row 20, column 30 is label 4 (upper clothes), row 46, column 25 label 9 (shoes).
|
||||
for (row, column), label in (((20, 30), 4), ((46, 25), 9), ((0, 0), 0)):
|
||||
colour = torch.tensor(vton.LIP_PALETTE[3 * label:3 * label + 3], dtype=torch.float32) / 255.0
|
||||
assert torch.allclose(parsing_image[0, row, column], colour), label
|
||||
@@ -0,0 +1,843 @@
|
||||
"""HTTP routes (core/folder_routes.py, core/large_preview_routes.py) and the frontend extensions (web/*.js).
|
||||
|
||||
The routes are served from a throw-away aiohttp app through each module's ``register_routes``; no ComfyUI server runs.
|
||||
The extensions run in Node against a fake ComfyUI app and fake nodes (two kinds of widget: the one of the current
|
||||
frontend, which has ``hidden``, and a plain one from an older frontend, which has not).
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import importlib
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
import threading
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from aiohttp import web
|
||||
from aiohttp.test_utils import TestClient, TestServer
|
||||
|
||||
import folder_paths
|
||||
|
||||
pack = importlib.import_module("NH-Nodes")
|
||||
folder_routes = importlib.import_module("NH-Nodes.core.folder_routes")
|
||||
preview_routes = importlib.import_module("NH-Nodes.core.large_preview_routes")
|
||||
preview_cache = importlib.import_module("NH-Nodes.core.large_preview_cache")
|
||||
resolution_data = importlib.import_module("NH-Nodes.core.resolution_data")
|
||||
picker = importlib.import_module("NH-Nodes.nodes.resolution.smart_resolution_picker")
|
||||
|
||||
PACK_DIR = Path(__file__).resolve().parent.parent
|
||||
WEB_FILES = sorted((PACK_DIR / "web").glob("*.js"))
|
||||
SHARED_MODULE = "nh_widget_utils.js"
|
||||
|
||||
COUNT = folder_routes.COUNT_PATH
|
||||
PRESETS = folder_routes.PRESETS_PATH
|
||||
CLEANUP = "/nh-nodes/large-preview/cleanup"
|
||||
NO_CACHE_ID = "0" * 32
|
||||
|
||||
|
||||
def call(method, path, **kwargs):
|
||||
"""One request to the pack's routes. Returns (status, parsed JSON) or (status, raw bytes) for other content."""
|
||||
|
||||
async def run():
|
||||
routes = web.RouteTableDef()
|
||||
folder_routes.register_routes(routes)
|
||||
preview_routes.register_routes(routes)
|
||||
app = web.Application()
|
||||
app.add_routes(routes)
|
||||
async with TestClient(TestServer(app)) as client:
|
||||
response = await client.request(method, path, **kwargs)
|
||||
return response.status, response.content_type, await response.read()
|
||||
|
||||
status, content_type, body = asyncio.run(run())
|
||||
return status, json.loads(body) if content_type == "application/json" else body
|
||||
|
||||
|
||||
# --- registration -----------------------------------------------------------------------------------------------
|
||||
|
||||
def test_every_route_is_registered_and_the_package_imports_both_route_modules():
|
||||
routes = web.RouteTableDef()
|
||||
folder_routes.register_routes(routes)
|
||||
preview_routes.register_routes(routes)
|
||||
assert {(route.method, route.path) for route in routes} == {
|
||||
("GET", "/nh-nodes/load-images-folder/count"),
|
||||
("GET", "/nh-nodes/smart-resolution-picker/presets"),
|
||||
("GET", "/nh-nodes/large-preview/manifest/{cache_id}"),
|
||||
("GET", "/nh-nodes/large-preview/thumb/{cache_id}"),
|
||||
("GET", "/nh-nodes/large-preview/tile/{cache_id}/{level}/{tile}"),
|
||||
("GET", "/nh-nodes/large-preview/cache-info"),
|
||||
("POST", "/nh-nodes/large-preview/cleanup"),
|
||||
}
|
||||
assert pack.folder_routes is folder_routes and pack.large_preview_routes is preview_routes
|
||||
|
||||
|
||||
# --- /nh-nodes/load-images-folder/count --------------------------------------------------------------------------
|
||||
|
||||
@pytest.fixture
|
||||
def image_tree(tmp_path, monkeypatch):
|
||||
"""input/ holds 2 images, input/sub/ 2 more, input/sub/deeper/ 1 more; empty/ holds none."""
|
||||
root = tmp_path / "input"
|
||||
(root / "sub" / "deeper").mkdir(parents=True)
|
||||
for relative in ("top1.png", "top2.JPG", "notes.txt", "sub/deep1.png", "sub/deep2.jpg", "sub/deeper/deep3.webp"):
|
||||
(root / relative).write_bytes(b"")
|
||||
(tmp_path / "empty").mkdir()
|
||||
monkeypatch.setattr(folder_paths, "get_input_directory", lambda: str(root))
|
||||
return root
|
||||
|
||||
|
||||
def test_count_answers_with_the_count_and_nothing_else(image_tree):
|
||||
status, body = call("GET", COUNT, params={"folder_path": str(image_tree)})
|
||||
assert status == 200
|
||||
assert body == {"count": 2} # in particular no resolved "path": ComfyUI may be reachable from the network
|
||||
|
||||
|
||||
@pytest.mark.parametrize("recursive,expected", [("true", 5), ("TRUE", 5), ("1", 5), ("yes", 5), ("on", 5),
|
||||
("false", 2), ("0", 2), ("no", 2), ("off", 2), ("", 2)])
|
||||
def test_count_honours_the_recursive_parameter(image_tree, recursive, expected):
|
||||
status, body = call("GET", COUNT, params={"folder_path": str(image_tree), "recursive": recursive})
|
||||
assert (status, body) == (200, {"count": expected})
|
||||
|
||||
|
||||
def test_count_without_recursive_is_not_recursive(image_tree):
|
||||
assert call("GET", COUNT, params={"folder_path": str(image_tree)}) == (200, {"count": 2})
|
||||
|
||||
|
||||
def test_count_resolves_like_the_loader_node(image_tree):
|
||||
assert call("GET", COUNT, params={"folder_path": "sub"}) == (200, {"count": 2}) # relative to the input folder
|
||||
assert call("GET", COUNT, params={"folder_path": "sub", "recursive": "true"}) == (200, {"count": 3})
|
||||
assert call("GET", COUNT) == (200, {"count": 2}) # no folder: the input folder
|
||||
assert call("GET", COUNT, params={"folder_path": str(image_tree / "top1.png")}) == (200, {"count": 2}) # a file: its folder
|
||||
assert call("GET", COUNT, params={"folder_path": " sub "}) == (200, {"count": 2}) # surrounding blanks are ignored
|
||||
|
||||
|
||||
def test_count_of_a_folder_without_images_is_zero(image_tree):
|
||||
assert call("GET", COUNT, params={"folder_path": str(image_tree.parent / "empty")}) == (200, {"count": 0})
|
||||
|
||||
|
||||
def test_count_of_a_missing_folder_is_404_and_does_not_echo_the_path(image_tree):
|
||||
missing = image_tree.parent / "does-not-exist"
|
||||
status, body = call("GET", COUNT, params={"folder_path": str(missing)})
|
||||
assert status == 404
|
||||
assert body == {"error": "folder not found"}
|
||||
assert str(missing) not in json.dumps(body)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("params", [
|
||||
{"folder_path": "x", "recursive": "maybe"},
|
||||
{"folder_path": "x", "recursive": "2"},
|
||||
{"folder_path": "bad\0name"},
|
||||
{"folder_path": "a" * 5000},
|
||||
])
|
||||
def test_count_rejects_bad_input_with_400(image_tree, params):
|
||||
status, body = call("GET", COUNT, params=params)
|
||||
assert status == 400 and list(body) == ["error"]
|
||||
|
||||
|
||||
def test_count_errors_are_short_and_never_leak_internals(image_tree, monkeypatch):
|
||||
def denied(folder, recursive=False):
|
||||
raise PermissionError(f"[Errno 13] Permission denied: {folder!r}")
|
||||
|
||||
monkeypatch.setattr(folder_routes, "collect_image_files", denied)
|
||||
status, body = call("GET", COUNT, params={"folder_path": str(image_tree)})
|
||||
assert (status, body) == (403, {"error": "folder is not readable"})
|
||||
|
||||
def broken(folder, recursive=False):
|
||||
raise RuntimeError(f"secret detail about {folder}")
|
||||
|
||||
monkeypatch.setattr(folder_routes, "collect_image_files", broken)
|
||||
status, body = call("GET", COUNT, params={"folder_path": str(image_tree)})
|
||||
assert status == 500
|
||||
assert body == {"error": "could not read the folder"}
|
||||
|
||||
|
||||
def test_count_scans_in_a_worker_thread_not_on_the_event_loop(image_tree, monkeypatch):
|
||||
seen = []
|
||||
real = folder_routes.collect_image_files
|
||||
|
||||
def spy(folder, recursive=False):
|
||||
seen.append(threading.get_ident())
|
||||
return real(folder, recursive=recursive)
|
||||
|
||||
monkeypatch.setattr(folder_routes, "collect_image_files", spy)
|
||||
assert call("GET", COUNT, params={"folder_path": str(image_tree)})[0] == 200
|
||||
assert seen and seen[0] != threading.get_ident()
|
||||
|
||||
|
||||
# --- /nh-nodes/smart-resolution-picker/presets --------------------------------------------------------------------
|
||||
|
||||
def test_presets_payload_matches_the_resolution_data():
|
||||
status, body = call("GET", PRESETS)
|
||||
assert status == 200
|
||||
assert set(body) == {"models", "default_model", "default_preset", "presets_by_model", "resolution_levels",
|
||||
"ratios_by_model_level"}
|
||||
assert body["models"] == resolution_data.MODEL_LABELS
|
||||
assert body["default_model"] == resolution_data.DEFAULT_MODEL_LABEL
|
||||
assert body["default_preset"] == resolution_data.DEFAULT_PRESET
|
||||
assert body["presets_by_model"] == resolution_data.PRESET_LABELS_BY_MODEL
|
||||
assert body["resolution_levels"] == picker.TARGET_RESOLUTION_LEVELS
|
||||
for model in body["models"]:
|
||||
for level in body["resolution_levels"]:
|
||||
ratios = body["ratios_by_model_level"][model][level]
|
||||
assert ratios and len(ratios) == len(set(ratios))
|
||||
|
||||
|
||||
# --- /nh-nodes/large-preview/* -----------------------------------------------------------------------------------
|
||||
|
||||
@pytest.fixture
|
||||
def preview_temp(tmp_path, monkeypatch):
|
||||
temp = tmp_path / "temp"
|
||||
monkeypatch.setattr(folder_paths, "get_temp_directory", lambda: str(temp))
|
||||
return temp
|
||||
|
||||
|
||||
@pytest.mark.parametrize("body", ["[]", "[1, 2]", '"text"', "5", "null", "true"])
|
||||
def test_cleanup_rejects_a_json_body_that_is_not_an_object(preview_temp, body):
|
||||
status, payload = call("POST", CLEANUP, data=body, headers={"Content-Type": "application/json"})
|
||||
assert status == 400
|
||||
assert payload == {"error": "request body must be a JSON object"}
|
||||
|
||||
|
||||
def test_cleanup_rejects_a_body_that_is_not_json(preview_temp):
|
||||
status, payload = call("POST", CLEANUP, data="{not json", headers={"Content-Type": "application/json"})
|
||||
assert status == 400 and "JSON" in payload["error"]
|
||||
status, _ = call("POST", CLEANUP, data=b"\xff\xfe\x00", headers={"Content-Type": "application/json"})
|
||||
assert status == 400
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kwargs", [{}, {"data": ""}, {"data": " \n"}, {"json": {}}, {"json": {"max_age_hours": 24}}])
|
||||
def test_cleanup_accepts_an_empty_body_or_an_object(preview_temp, kwargs):
|
||||
status, payload = call("POST", CLEANUP, **kwargs)
|
||||
assert status == 200 and payload["removed"] == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("payload", [
|
||||
{"max_age_hours": [1]},
|
||||
{"max_age_hours": "abc"},
|
||||
{"max_total_mb": [1]},
|
||||
{"max_age_hours": 10 ** 400},
|
||||
{"cache_id": "../outside"},
|
||||
{"cache_id": 5},
|
||||
])
|
||||
def test_cleanup_answers_400_for_values_it_cannot_use(preview_temp, payload):
|
||||
status, body = call("POST", CLEANUP, json=payload)
|
||||
assert status == 400 and list(body) == ["error"]
|
||||
|
||||
|
||||
def test_cleanup_removes_old_entries_and_keeps_new_ones(preview_temp):
|
||||
root = preview_temp / preview_cache.CACHE_SUBDIR
|
||||
old, new = root / ("a" * 32), root / ("b" * 32)
|
||||
for entry in (old, new):
|
||||
entry.mkdir(parents=True)
|
||||
(entry / "manifest.json").write_text("{}", encoding="utf-8")
|
||||
long_ago = time.time() - 48 * 3600
|
||||
os.utime(old, (long_ago, long_ago))
|
||||
|
||||
status, body = call("POST", CLEANUP, json={"max_age_hours": 24})
|
||||
assert status == 200 and body["removed"] == 1
|
||||
assert not old.exists() and new.exists()
|
||||
|
||||
|
||||
def test_large_preview_routes_serve_and_delete_a_cache_entry(preview_temp):
|
||||
manifest = preview_cache.create_preview_cache(torch.rand(48, 64, 3))
|
||||
cache_id = manifest["cache_id"]
|
||||
base = "/nh-nodes/large-preview"
|
||||
|
||||
status, body = call("GET", f"{base}/manifest/{cache_id}")
|
||||
assert status == 200 and body["cache_id"] == cache_id
|
||||
for path in (f"{base}/thumb/{cache_id}", f"{base}/tile/{cache_id}/0/0_0.jpg"):
|
||||
status, data = call("GET", path)
|
||||
assert status == 200 and data[:2] == b"\xff\xd8" # a JPEG
|
||||
status, info = call("GET", f"{base}/cache-info")
|
||||
assert status == 200 and info["cache_count"] == 1
|
||||
assert "root" not in json.dumps(info) and str(preview_temp) not in json.dumps(info) # no absolute paths leak
|
||||
|
||||
status, body = call("POST", CLEANUP, json={"cache_id": cache_id})
|
||||
assert status == 200 and body["removed"] == 1
|
||||
assert call("GET", f"{base}/manifest/{cache_id}")[0] == 404
|
||||
|
||||
|
||||
@pytest.mark.parametrize("path,status", [
|
||||
(f"/nh-nodes/large-preview/manifest/{NO_CACHE_ID}", 404),
|
||||
("/nh-nodes/large-preview/manifest/not-an-id", 400),
|
||||
(f"/nh-nodes/large-preview/thumb/{NO_CACHE_ID}", 404),
|
||||
("/nh-nodes/large-preview/thumb/zz", 400),
|
||||
(f"/nh-nodes/large-preview/tile/{NO_CACHE_ID}/0/0_0.jpg", 404),
|
||||
(f"/nh-nodes/large-preview/tile/{NO_CACHE_ID}/x/0_0.jpg", 400),
|
||||
(f"/nh-nodes/large-preview/tile/{NO_CACHE_ID}/0/passwd", 400),
|
||||
])
|
||||
def test_large_preview_routes_answer_404_for_missing_and_400_for_malformed(preview_temp, path, status):
|
||||
assert call("GET", path)[0] == status
|
||||
|
||||
|
||||
def test_large_preview_tile_route_does_not_leave_the_cache_folder(preview_temp):
|
||||
cache_id = preview_cache.create_preview_cache(torch.rand(8, 8, 3))["cache_id"]
|
||||
status, _ = call("GET", f"/nh-nodes/large-preview/tile/{cache_id}/0/..%2F..%2Fmanifest.json")
|
||||
assert status in (400, 404)
|
||||
|
||||
|
||||
# --- web/*.js: syntax and structure ------------------------------------------------------------------------------
|
||||
|
||||
needs_node = pytest.mark.skipif(shutil.which("node") is None, reason="node is not installed")
|
||||
|
||||
|
||||
@needs_node
|
||||
@pytest.mark.parametrize("path", WEB_FILES, ids=lambda path: path.name)
|
||||
def test_every_js_file_is_valid_syntax(path):
|
||||
check = subprocess.run([shutil.which("node"), "--check", str(path)], capture_output=True, text=True,
|
||||
encoding="utf-8", timeout=60, check=False)
|
||||
assert check.returncode == 0, check.stderr
|
||||
|
||||
|
||||
def test_the_shared_module_imports_nothing_and_registers_no_extension():
|
||||
source = (PACK_DIR / "web" / SHARED_MODULE).read_text(encoding="utf-8")
|
||||
assert "registerExtension" not in source
|
||||
assert not re.search(r"^\s*import\b", source, re.MULTILINE)
|
||||
|
||||
|
||||
COPIED_HELPER = re.compile(
|
||||
r"^\s*(?:async\s+)?function\s+(findWidget|toInt|clampInt|hideWidget|showWidget|chainCallback|markDirty|refreshNode)\b",
|
||||
re.MULTILINE,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("path", [path for path in WEB_FILES if path.name != SHARED_MODULE], ids=lambda path: path.name)
|
||||
def test_extensions_import_the_shared_helpers_instead_of_copying_them(path):
|
||||
source = path.read_text(encoding="utf-8")
|
||||
assert not COPIED_HELPER.search(source)
|
||||
if re.search(r"\b(?:findWidget|toInt|clampInt|hideWidget|showWidget|chainCallback|markDirty|refreshNode)\(", source):
|
||||
assert f'from "./{SHARED_MODULE}"' in source
|
||||
|
||||
|
||||
# --- web/*.js: behaviour in Node ----------------------------------------------------------------------------------
|
||||
|
||||
JS_HARNESS = r"""
|
||||
import fs from "node:fs";
|
||||
import path from "node:path";
|
||||
import { pathToFileURL } from "node:url";
|
||||
|
||||
const webDir = process.argv[2];
|
||||
const sleep = (ms) => new Promise((resolve) => setTimeout(resolve, ms));
|
||||
const toUrl = (file) => pathToFileURL(path.join(webDir, file)).href;
|
||||
|
||||
// --- a fake ComfyUI application: it records which extension each file registers ---
|
||||
const registered = [];
|
||||
let loading = null;
|
||||
const appGraph = { dirty: 0, setDirtyCanvas() { this.dirty++; } };
|
||||
globalThis.__nhApp = { graph: appGraph, registerExtension(extension) { registered.push({ file: loading, extension }); } };
|
||||
globalThis.__nhApi = { apiURL: (url) => url };
|
||||
|
||||
const files = fs.readdirSync(webDir).filter((name) => name.endsWith(".js")).sort();
|
||||
for (const file of files) {
|
||||
loading = file;
|
||||
await import(toUrl(file));
|
||||
}
|
||||
loading = null;
|
||||
const utils = await import(toUrl("nh_widget_utils.js"));
|
||||
|
||||
// --- two kinds of widget ---
|
||||
// "modern": ComfyUI frontend 1.55.17 (litegraph BaseWidget). `hidden` is an accessor on the prototype, and a type that
|
||||
// contains "hidden" hides the widget as well. "legacy": a plain object of an older frontend, no `hidden` at all.
|
||||
class ModernWidget {
|
||||
constructor(name, type, value, options) {
|
||||
this.name = name;
|
||||
this._type = type;
|
||||
this.value = value;
|
||||
this.options = options ?? {};
|
||||
this._hidden = false;
|
||||
this._hiddenByType = false;
|
||||
}
|
||||
get type() { return this._type; }
|
||||
set type(value) { this._type = value; this._hiddenByType = /hidden/i.test(value); }
|
||||
get hidden() { return this._hidden || this._hiddenByType; }
|
||||
set hidden(value) { this._hidden = Boolean(value); }
|
||||
computeSize() { return [0, 20]; }
|
||||
}
|
||||
|
||||
function makeLegacyWidget(name, type, value, options) {
|
||||
return { name, type, value, options: options ?? {}, computeSize() { return [0, 20]; } };
|
||||
}
|
||||
|
||||
// A node with the parts of LGraphNode the extensions use, plus what the frontend does with its widgets: the
|
||||
// workflow keeps the value of every widget unless widget.serialize is false (litegraph serialiseWidgetValues), the
|
||||
// prompt gets every widget unless options.serialize is false (graphToPrompt), and a hidden widget takes no row.
|
||||
function makeNode(kind, specs) {
|
||||
const make = kind === "modern" ? (...args) => new ModernWidget(...args) : makeLegacyWidget;
|
||||
return {
|
||||
id: 1, graph: appGraph, size: [200, 100], inputs: [], outputs: [], sizeCalls: 0, dirtyCalls: 0,
|
||||
widgets: specs.map(([name, type, value, options]) => make(name, type, value, options)),
|
||||
addWidget(type, name, value, callback) {
|
||||
const widget = make(name, type, value, {});
|
||||
widget.callback = callback;
|
||||
this.widgets.push(widget);
|
||||
return widget;
|
||||
},
|
||||
addOutput(name, type) { this.outputs.push({ name, type }); },
|
||||
removeOutput(slot) { this.outputs.splice(slot, 1); },
|
||||
disconnectOutput() {},
|
||||
widget(name) { return this.widgets.find((widget) => widget.name === name); },
|
||||
isShown(widget) { return kind === "modern" ? !widget.hidden : widget.type !== "nh-hidden"; },
|
||||
shown() { return this.widgets.filter((widget) => this.isShown(widget)).map((widget) => widget.name); },
|
||||
hiddenNames() { return this.widgets.filter((widget) => !this.isShown(widget)).map((widget) => widget.name); },
|
||||
types() { return Object.fromEntries(this.widgets.map((widget) => [widget.name, widget.type])); },
|
||||
computeSize() { return [200, 40 + 24 * this.widgets.filter((widget) => this.isShown(widget)).length]; },
|
||||
setSize(size) { this.size = size; this.sizeCalls++; },
|
||||
setDirtyCanvas() { this.dirtyCalls++; },
|
||||
workflowValues() { return this.widgets.filter((widget) => widget.serialize !== false).map((widget) => widget.value); },
|
||||
promptInputs() {
|
||||
const inputs = {};
|
||||
for (const widget of this.widgets) {
|
||||
if (widget.name && widget.options?.serialize !== false) inputs[widget.name] = widget.value;
|
||||
}
|
||||
return inputs;
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
// What ComfyUI does for a node: beforeRegisterNodeDef hooks the class, the node is created, then configured.
|
||||
async function mountNode(file, nodeTypeName, kind, specs) {
|
||||
const { extension } = registered.find((entry) => entry.file === file);
|
||||
const NodeType = function () {};
|
||||
await extension.beforeRegisterNodeDef(NodeType, { name: nodeTypeName });
|
||||
const node = makeNode(kind, specs);
|
||||
NodeType.prototype.onNodeCreated.call(node);
|
||||
NodeType.prototype.onConfigure?.call(node, {});
|
||||
await sleep(30);
|
||||
return node;
|
||||
}
|
||||
|
||||
const results = {};
|
||||
async function scenario(name, run) {
|
||||
try {
|
||||
results[name] = await run();
|
||||
} catch (error) {
|
||||
results[name] = { error: String(error?.stack ?? error) };
|
||||
}
|
||||
}
|
||||
const numbered = (prefix, count) => Array.from({ length: count }, (_, index) => `${prefix}_${index + 1}`);
|
||||
|
||||
// --- the files ---
|
||||
results.registration = registered.map(({ file, extension }) => [file, extension.name]);
|
||||
results.utilsExports = Object.keys(utils).sort();
|
||||
|
||||
// --- the shared helpers ---
|
||||
await scenario("utils_basic", async () => {
|
||||
const node = { widgets: [{ name: "a", value: 1 }, { name: "b", value: 2 }] };
|
||||
return {
|
||||
found: utils.findWidget(node, "b")?.value,
|
||||
missing: utils.findWidget(node, "zzz") === undefined,
|
||||
noWidgets: utils.findWidget({}, "a") === undefined,
|
||||
nullNode: utils.findWidget(null, "a") === undefined,
|
||||
toInt: [["12", 0], [" 7", 0], [3.9, 0], ["12px", 0], ["abc", -1], [undefined, -1], [null, -1], [NaN, -1], ["", -1], [Infinity, -1]]
|
||||
.map(([value, fallback]) => utils.toInt(value, fallback)),
|
||||
clampInt: [[5, 0, 10, 1], [-3, 0, 10, 1], [99, 0, 10, 1], ["x", 2, 10, 5], ["x", 2, 10, 50], [NaN, 1, 64, 5], ["7", 1, 64, 5]]
|
||||
.map((args) => utils.clampInt(...args)),
|
||||
};
|
||||
});
|
||||
|
||||
await scenario("hide_show_modern", async () => {
|
||||
const widget = new ModernWidget("w", "text", "keep");
|
||||
const out = {};
|
||||
utils.hideWidget(widget);
|
||||
out.hidden = widget.hidden;
|
||||
out.type = widget.type;
|
||||
out.ownComputeSize = Object.hasOwn(widget, "computeSize");
|
||||
utils.hideWidget(widget); // twice
|
||||
utils.showWidget(widget);
|
||||
out.shownAgain = widget.hidden;
|
||||
out.typeAfter = widget.type;
|
||||
utils.showWidget(widget); // twice
|
||||
out.value = widget.value;
|
||||
out.serializeUntouched = widget.serialize === undefined;
|
||||
const hiddenByDefinition = new ModernWidget("d", "text", 1);
|
||||
hiddenByDefinition.hidden = true;
|
||||
utils.showWidget(hiddenByDefinition);
|
||||
out.hiddenByDefinitionStays = hiddenByDefinition.hidden;
|
||||
utils.hideWidget(undefined);
|
||||
utils.showWidget(undefined);
|
||||
utils.showWidget(null);
|
||||
return out;
|
||||
});
|
||||
|
||||
await scenario("hide_show_legacy", async () => {
|
||||
const widget = makeLegacyWidget("w", "text", "keep");
|
||||
const original = widget.computeSize;
|
||||
const out = {};
|
||||
utils.hideWidget(widget);
|
||||
out.hiddenType = widget.type;
|
||||
out.hiddenSize = widget.computeSize();
|
||||
out.gotHiddenProperty = "hidden" in widget;
|
||||
utils.hideWidget(widget); // a second call must not save "nh-hidden" as the type to restore
|
||||
utils.showWidget(widget);
|
||||
out.type = widget.type;
|
||||
out.sizeRestored = widget.computeSize === original;
|
||||
out.leftovers = Object.keys(widget).filter((key) => key.startsWith("__nh"));
|
||||
out.value = widget.value;
|
||||
const inherited = Object.assign(Object.create({ computeSize() { return [0, 33]; } }), { name: "i", type: "combo", value: 1, options: {} });
|
||||
utils.hideWidget(inherited);
|
||||
out.inheritedHiddenSize = inherited.computeSize();
|
||||
utils.showWidget(inherited);
|
||||
out.inheritedOwnSize = Object.hasOwn(inherited, "computeSize");
|
||||
out.inheritedSize = inherited.computeSize();
|
||||
out.inheritedType = inherited.type;
|
||||
return out;
|
||||
});
|
||||
|
||||
await scenario("chain_callback", async () => {
|
||||
const log = [];
|
||||
const plain = { name: "plain" };
|
||||
utils.chainCallback(plain, function (...args) { log.push(["after", this === plain, args]); });
|
||||
plain.callback(1, 2);
|
||||
const withOriginal = { name: "original", callback(...args) { log.push(["original", this === withOriginal, args]); } };
|
||||
utils.chainCallback(withOriginal, function (...args) { log.push(["after", this === withOriginal, args]); });
|
||||
withOriginal.callback("x", "y");
|
||||
return log;
|
||||
});
|
||||
|
||||
await scenario("refresh", async () => {
|
||||
const calls = [];
|
||||
const node = {
|
||||
size: [1, 1],
|
||||
computeSize() { calls.push("computeSize"); return [5, 6]; },
|
||||
setSize(size) { calls.push(["setSize", size]); },
|
||||
setDirtyCanvas(...args) { calls.push(["nodeDirty", ...args]); },
|
||||
};
|
||||
utils.refreshNode(node, { setDirtyCanvas(...args) { calls.push(["graphDirty", ...args]); } });
|
||||
const sizeless = { size: [9, 9], computeSize() { return null; }, setSize(size) { calls.push(["setSize", size]); } };
|
||||
utils.refreshNode(sizeless, null);
|
||||
utils.refreshNode({});
|
||||
utils.markDirty({}, undefined);
|
||||
return calls;
|
||||
});
|
||||
|
||||
// --- the extensions, on both kinds of frontend ---
|
||||
for (const kind of ["modern", "legacy"]) {
|
||||
await scenario(`value_match_index_${kind}`, async () => {
|
||||
const specs = [["value_count", "number", 3], ...numbered("value", 5).map((name) => [name, "text", `text ${name}`])];
|
||||
const node = await mountNode("nh_value_match_index.js", "NH_ValueMatchIndex", kind, specs);
|
||||
const out = { shown: node.shown(), hidden: node.hiddenNames(), types: node.types() };
|
||||
const count = node.widget("value_count");
|
||||
count.value = 5;
|
||||
count.callback();
|
||||
out.five = node.shown();
|
||||
out.typesAfterFive = node.types();
|
||||
count.value = 1;
|
||||
count.callback();
|
||||
out.one = node.shown();
|
||||
out.workflowValues = node.workflowValues();
|
||||
out.promptInputs = node.promptInputs();
|
||||
count.value = "abc";
|
||||
count.callback();
|
||||
out.notANumber = { count: count.value, shown: node.shown() };
|
||||
count.value = 1000;
|
||||
count.callback();
|
||||
out.tooMany = count.value;
|
||||
out.buttonSerialize = node.widget("update values").serialize;
|
||||
out.sized = node.sizeCalls > 0 && node.dirtyCalls > 0 && appGraph.dirty > 0;
|
||||
return out;
|
||||
});
|
||||
|
||||
await scenario(`save_image_path_${kind}`, async () => {
|
||||
const node = await mountNode("nh_save_image_path.js", "NH_SaveImagePath", kind,
|
||||
[["create_new_folder", "toggle", false], ["new_folder_name", "text", "renders"]]);
|
||||
const out = { off: node.shown() };
|
||||
const create = node.widget("create_new_folder");
|
||||
create.value = true;
|
||||
create.callback();
|
||||
out.on = node.shown();
|
||||
create.value = false;
|
||||
create.callback();
|
||||
out.offAgain = node.shown();
|
||||
out.workflowValues = node.workflowValues();
|
||||
out.promptInputs = node.promptInputs();
|
||||
return out;
|
||||
});
|
||||
|
||||
await scenario(`index_nodes_${kind}`, async () => {
|
||||
const specs = [
|
||||
["select_mode", "combo", "index"], ["index", "number", 0], ["lora_count", "number", 2],
|
||||
...numbered("lora", 3).map((name) => [name, "combo", name]),
|
||||
...numbered("enable", 3).map((name) => [name, "toggle", false]),
|
||||
];
|
||||
const node = await mountNode("nh_index_nodes.js", "NH_LoraModelIndex", kind, specs);
|
||||
const out = { indexMode: node.shown() };
|
||||
const mode = node.widget("select_mode");
|
||||
mode.value = "boolean";
|
||||
mode.callback();
|
||||
out.booleanMode = node.shown();
|
||||
const first = node.widget("enable_1");
|
||||
const second = node.widget("enable_2");
|
||||
first.value = true;
|
||||
second.value = true;
|
||||
second.callback();
|
||||
out.exclusive = [first.value, second.value];
|
||||
const count = node.widget("lora_count");
|
||||
count.value = 3;
|
||||
count.callback();
|
||||
out.threeLoras = node.shown();
|
||||
out.workflowValues = node.workflowValues();
|
||||
|
||||
const concat = await mountNode("nh_index_nodes.js", "NH_TextConcatenate", kind,
|
||||
[["string_count", "number", 1], ...numbered("string", 3).map((name) => [name, "text", ""])]);
|
||||
out.concat = { count: concat.widget("string_count").value, shown: concat.shown() };
|
||||
return out;
|
||||
});
|
||||
}
|
||||
|
||||
await scenario("any_list_split", async () => {
|
||||
const node = await mountNode("nh_any_list_split.js", "NH_AnyListSplit", "modern", [["index_count", "number", 3]]);
|
||||
const names = () => node.outputs.map((output) => output.name);
|
||||
const out = { three: names() };
|
||||
const count = node.widget("index_count");
|
||||
count.value = 1;
|
||||
count.callback();
|
||||
out.one = names();
|
||||
count.value = 100;
|
||||
count.callback();
|
||||
out.clamped = { count: count.value, names: names() };
|
||||
out.buttonSerialize = node.widget("update outputs").serialize;
|
||||
return out;
|
||||
});
|
||||
|
||||
await scenario("load_images_folder", async () => {
|
||||
globalThis.window = { location: { origin: "http://127.0.0.1:8188" } };
|
||||
const calls = [];
|
||||
let reply = (url) => ({ ok: true, json: async () => ({ count: url.searchParams.get("recursive") === "true" ? 5 : 2 }) });
|
||||
globalThis.fetch = async (input) => {
|
||||
const url = new URL(String(input));
|
||||
calls.push({ path: url.pathname, folder: url.searchParams.get("folder_path"), recursive: url.searchParams.get("recursive") });
|
||||
return reply(url);
|
||||
};
|
||||
const specs = [["folder_path", "text", "C:/imgs"], ["index", "number", 3], ["step", "number", 1],
|
||||
["seed_mode", "combo", "increment"], ["recursive", "toggle", true]];
|
||||
const node = await mountNode("nh_load_images_folder.js", "NH_LoadImagesFromFolder", "modern", specs);
|
||||
const index = node.widget("index");
|
||||
const recursive = node.widget("recursive");
|
||||
const out = {};
|
||||
const queue = async () => { const sent = await index.serializeValue(); await sleep(10); return [sent, index.value]; };
|
||||
|
||||
out.recursiveOn = await queue(); // 5 images: 3 is kept, next is 4
|
||||
recursive.value = false;
|
||||
out.recursiveOff = await queue(); // 2 images: 4 wraps to 0, next is 1
|
||||
recursive.value = true;
|
||||
out.recursiveOnAgain = await queue();
|
||||
out.calls = calls.map((call) => call.recursive);
|
||||
out.paths = [...new Set(calls.map((call) => call.path))];
|
||||
out.folders = [...new Set(calls.map((call) => call.folder))];
|
||||
|
||||
node.widget("seed_mode").value = "decrement";
|
||||
node.widget("step").value = 2;
|
||||
index.value = 0;
|
||||
out.decrement = await queue(); // 5 images: 0 stays, next is 3
|
||||
node.widget("seed_mode").value = "random";
|
||||
const random = Math.random;
|
||||
Math.random = () => 0.5;
|
||||
out.random = await queue(); // floor(0.5 * 5)
|
||||
Math.random = random;
|
||||
node.widget("seed_mode").value = "fixed";
|
||||
index.value = 7;
|
||||
out.fixed = await queue(); // 7 % 5 = 2
|
||||
reply = () => ({ ok: false, status: 404 });
|
||||
index.value = 7;
|
||||
out.notFound = await queue(); // no count: the index is left alone
|
||||
globalThis.fetch = async () => { throw new Error("offline"); };
|
||||
out.offline = await queue();
|
||||
|
||||
globalThis.fetch = async (input) => { calls.push({ recursive: new URL(String(input)).searchParams.get("recursive") }); return { ok: true, json: async () => ({ count: 4 }) }; };
|
||||
const without = await mountNode("nh_load_images_folder.js", "NH_LoadImagesFromFolder", "modern", specs.slice(0, 4));
|
||||
await without.widget("index").serializeValue();
|
||||
out.withoutRecursiveWidget = calls[calls.length - 1].recursive;
|
||||
return out;
|
||||
});
|
||||
|
||||
await scenario("smart_resolution_picker", async () => {
|
||||
const presets = {
|
||||
models: ["A", "B"], default_model: "A", default_preset: "A-1",
|
||||
presets_by_model: { A: ["A-1", "A-2"], B: ["B-1"] },
|
||||
resolution_levels: ["1 MP", "2 MP"],
|
||||
ratios_by_model_level: { A: { "1 MP": ["1:1", "3:2"], "2 MP": ["1:1"] }, B: { "1 MP": ["16:9"], "2 MP": ["16:9", "4:3"] } },
|
||||
};
|
||||
let fetches = 0;
|
||||
globalThis.fetch = async () => { fetches++; return { ok: true, status: 200, json: async () => presets }; };
|
||||
const out = {};
|
||||
|
||||
const pickerNode = await mountNode("nh_smart_resolution_picker.js", "NH_SmartResolutionPicker", "modern",
|
||||
[["model_family", "combo", "A"], ["preset", "combo", "stale"]]);
|
||||
const preset = () => [pickerNode.widget("preset").options.values, pickerNode.widget("preset").value];
|
||||
out.pickerStart = preset();
|
||||
pickerNode.widget("model_family").value = "B";
|
||||
pickerNode.widget("model_family").callback();
|
||||
out.pickerModelB = preset();
|
||||
|
||||
const ratioNode = await mountNode("nh_smart_resolution_picker.js", "NH_RatioPresetImageResize", "modern",
|
||||
[["model_family", "combo", "A"], ["resolution_level", "combo", "1 MP"], ["target_ratio", "combo", "9:9"]]);
|
||||
const ratio = () => [ratioNode.widget("target_ratio").options.values, ratioNode.widget("target_ratio").value];
|
||||
out.ratioStart = ratio();
|
||||
ratioNode.widget("resolution_level").value = "2 MP";
|
||||
ratioNode.widget("resolution_level").callback();
|
||||
out.ratioTwoMp = ratio();
|
||||
ratioNode.widget("model_family").value = "B";
|
||||
ratioNode.widget("model_family").callback();
|
||||
out.ratioModelB = ratio();
|
||||
out.fetches = fetches;
|
||||
return out;
|
||||
});
|
||||
|
||||
console.log("RESULT:" + JSON.stringify(results));
|
||||
"""
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def js(tmp_path_factory):
|
||||
"""The harness results. The files are copied unchanged into a folder shaped like ComfyUI's: the extensions'
|
||||
``../../../scripts/app.js`` finds the fake app, and ``./nh_widget_utils.js`` finds its sibling."""
|
||||
node = shutil.which("node")
|
||||
if node is None:
|
||||
pytest.skip("node is not installed")
|
||||
work = tmp_path_factory.mktemp("web_harness")
|
||||
(work / "package.json").write_text('{"type": "module"}', encoding="utf-8") # the copied .js files are ES modules
|
||||
(work / "scripts").mkdir()
|
||||
(work / "scripts" / "app.js").write_text("export const app = globalThis.__nhApp;\n", encoding="utf-8")
|
||||
(work / "scripts" / "api.js").write_text("export const api = globalThis.__nhApi;\n", encoding="utf-8")
|
||||
web_dir = work / "server" / "extensions" / "NH-Nodes"
|
||||
web_dir.mkdir(parents=True)
|
||||
for source in WEB_FILES:
|
||||
shutil.copy(source, web_dir / source.name)
|
||||
(work / "harness.mjs").write_text(JS_HARNESS, encoding="utf-8")
|
||||
run = subprocess.run([node, str(work / "harness.mjs"), str(web_dir)], capture_output=True, text=True,
|
||||
encoding="utf-8", timeout=120, check=False)
|
||||
assert run.returncode == 0, run.stderr
|
||||
line = next(line for line in run.stdout.splitlines() if line.startswith("RESULT:"))
|
||||
return json.loads(line[len("RESULT:"):])
|
||||
|
||||
|
||||
def scenario(js, name):
|
||||
result = js[name]
|
||||
assert "error" not in result, result.get("error")
|
||||
return result
|
||||
|
||||
|
||||
def test_each_extension_file_registers_exactly_one_uniquely_named_extension(js):
|
||||
registration = js["registration"]
|
||||
assert sorted(file for file, _ in registration) == sorted(path.name for path in WEB_FILES if path.name != SHARED_MODULE)
|
||||
names = [name for _, name in registration]
|
||||
assert len(names) == len(set(names)) and all(name.startswith("NH.Nodes.") for name in names)
|
||||
assert js["utilsExports"] == sorted(
|
||||
["chainCallback", "clampInt", "findWidget", "hideWidget", "markDirty", "refreshNode", "showWidget", "toInt"]
|
||||
)
|
||||
|
||||
|
||||
def test_findwidget_toint_and_clampint(js):
|
||||
out = scenario(js, "utils_basic")
|
||||
assert out["found"] == 2 and out["missing"] and out["noWidgets"] and out["nullNode"]
|
||||
assert out["toInt"] == [12, 7, 3, 12, -1, -1, -1, -1, -1, -1]
|
||||
assert out["clampInt"] == [5, 0, 10, 5, 10, 5, 7]
|
||||
|
||||
|
||||
def test_hidewidget_uses_the_hidden_flag_when_the_frontend_has_it(js):
|
||||
out = scenario(js, "hide_show_modern")
|
||||
assert out["hidden"] is True
|
||||
assert out["type"] == "text" and out["typeAfter"] == "text" # the type is not touched
|
||||
assert out["ownComputeSize"] is False
|
||||
assert out["shownAgain"] is False
|
||||
assert out["value"] == "keep" and out["serializeUntouched"] is True # still saved with the workflow
|
||||
assert out["hiddenByDefinitionStays"] is True # showWidget only undoes hideWidget
|
||||
|
||||
|
||||
def test_hidewidget_falls_back_to_the_type_swap_on_older_frontends(js):
|
||||
out = scenario(js, "hide_show_legacy")
|
||||
assert out["hiddenType"] == "nh-hidden" and out["hiddenSize"] == [0, -4] and out["gotHiddenProperty"] is False
|
||||
assert out["type"] == "text" and out["sizeRestored"] is True and out["leftovers"] == []
|
||||
assert out["value"] == "keep"
|
||||
assert out["inheritedHiddenSize"] == [0, -4]
|
||||
assert out["inheritedOwnSize"] is False and out["inheritedSize"] == [0, 33] and out["inheritedType"] == "combo"
|
||||
|
||||
|
||||
def test_chaincallback_runs_the_existing_callback_first_with_the_same_this_and_arguments(js):
|
||||
assert scenario(js, "chain_callback") == [
|
||||
["after", True, [1, 2]],
|
||||
["original", True, ["x", "y"]],
|
||||
["after", True, ["x", "y"]],
|
||||
]
|
||||
|
||||
|
||||
def test_refreshnode_resizes_then_repaints_and_tolerates_missing_methods(js):
|
||||
assert scenario(js, "refresh") == [
|
||||
"computeSize", ["setSize", [5, 6]], ["nodeDirty", True, True], ["graphDirty", True, True],
|
||||
["setSize", [9, 9]], # computeSize gave nothing: the current size is kept
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kind", ["modern", "legacy"])
|
||||
def test_value_match_index_shows_only_the_counted_values_and_still_saves_the_rest(js, kind):
|
||||
out = scenario(js, f"value_match_index_{kind}")
|
||||
assert out["shown"] == ["value_count", "value_1", "value_2", "value_3", "update values"]
|
||||
assert out["hidden"] == ["value_4", "value_5"]
|
||||
assert out["five"] == ["value_count", *[f"value_{i}" for i in range(1, 6)], "update values"]
|
||||
assert out["one"] == ["value_count", "value_1", "update values"]
|
||||
hidden_type = "text" if kind == "modern" else "nh-hidden"
|
||||
assert out["types"]["value_4"] == hidden_type and out["types"]["value_1"] == "text"
|
||||
assert set(out["typesAfterFive"].values()) <= {"text", "number", "button"} # shown again: the real type is back
|
||||
# the hidden widgets still travel with the workflow and the prompt
|
||||
assert out["workflowValues"] == [1, *[f"text value_{i}" for i in range(1, 6)]]
|
||||
assert {f"value_{i}" for i in range(1, 6)} <= set(out["promptInputs"])
|
||||
assert out["notANumber"] == {"count": 5, "shown": ["value_count", *[f"value_{i}" for i in range(1, 6)], "update values"]}
|
||||
assert out["tooMany"] == 64
|
||||
assert out["buttonSerialize"] is False and out["sized"] is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kind", ["modern", "legacy"])
|
||||
def test_save_image_path_shows_the_folder_name_only_when_creating_a_folder(js, kind):
|
||||
out = scenario(js, f"save_image_path_{kind}")
|
||||
assert out["off"] == ["create_new_folder"]
|
||||
assert out["on"] == ["create_new_folder", "new_folder_name"]
|
||||
assert out["offAgain"] == ["create_new_folder"]
|
||||
assert out["workflowValues"] == [False, "renders"]
|
||||
assert out["promptInputs"] == {"create_new_folder": False, "new_folder_name": "renders"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kind", ["modern", "legacy"])
|
||||
def test_index_nodes_switch_between_index_and_boolean_mode(js, kind):
|
||||
out = scenario(js, f"index_nodes_{kind}")
|
||||
assert out["indexMode"] == ["select_mode", "index", "lora_count", "lora_1", "lora_2", "update loras"]
|
||||
assert out["booleanMode"] == ["select_mode", "lora_count", "lora_1", "lora_2", "enable_1", "enable_2", "update loras"]
|
||||
assert out["exclusive"] == [False, True] # turning enable_2 on turns enable_1 off
|
||||
assert out["threeLoras"] == ["select_mode", "lora_count", "lora_1", "lora_2", "lora_3", "enable_1", "enable_2",
|
||||
"enable_3", "update loras"]
|
||||
assert len(out["workflowValues"]) == 9 # all 9 widgets, shown or not (the button is not saved)
|
||||
assert out["concat"] == {"count": 2, "shown": ["string_count", "string_1", "string_2", "update strings"]}
|
||||
|
||||
|
||||
def test_any_list_split_follows_index_count(js):
|
||||
out = scenario(js, "any_list_split")
|
||||
assert out["three"] == ["All", "index 1", "index 2", "index 3"]
|
||||
assert out["one"] == ["All", "index 1"]
|
||||
assert out["clamped"] == {"count": 10, "names": ["All", *[f"index {i}" for i in range(1, 11)]]}
|
||||
assert out["buttonSerialize"] is False
|
||||
|
||||
|
||||
def test_load_images_folder_counts_with_the_nodes_recursive_setting(js):
|
||||
out = scenario(js, "load_images_folder")
|
||||
assert out["calls"][:3] == ["true", "false", "true"] # follows the widget on every queue
|
||||
assert out["paths"] == ["/nh-nodes/load-images-folder/count"] and out["folders"] == ["C:/imgs"]
|
||||
assert out["recursiveOn"] == [3, 4] # 5 images: 3, then 4
|
||||
assert out["recursiveOff"] == [0, 1] # 2 images: 4 wraps to 0, then 1
|
||||
assert out["recursiveOnAgain"] == [1, 2]
|
||||
assert out["decrement"] == [0, 3] # step 2 back from 0 wraps to 3 of 5
|
||||
assert out["random"] == [2, 2]
|
||||
assert out["fixed"] == [2, 7] # 7 % 5 is sent, the widget keeps what the user typed
|
||||
assert out["notFound"] == [7, 7] and out["offline"][0] == 7 # unknown count: no wrapping
|
||||
assert out["withoutRecursiveWidget"] == "false"
|
||||
|
||||
|
||||
def test_smart_resolution_picker_fills_the_dropdowns_from_one_request(js):
|
||||
out = scenario(js, "smart_resolution_picker")
|
||||
assert out["pickerStart"] == [["A-1", "A-2"], "A-1"]
|
||||
assert out["pickerModelB"] == [["B-1"], "B-1"]
|
||||
assert out["ratioStart"] == [["1:1", "3:2"], "1:1"]
|
||||
assert out["ratioTwoMp"] == [["1:1"], "1:1"]
|
||||
assert out["ratioModelB"] == [["16:9", "4:3"], "16:9"]
|
||||
assert out["fetches"] == 1
|
||||
@@ -0,0 +1,903 @@
|
||||
"""VTON / SegFormer model handling: VRAM release (keep_model_loaded) and room-making that follow where
|
||||
the ONNX Runtime sessions really run, complete and pinned downloads, and model folders that follow
|
||||
extra_model_paths.yaml and list only model files.
|
||||
|
||||
Nothing here touches the network, a GPU or the real models folder: folder_paths points at a temp
|
||||
dir, huggingface_hub is a recorder, onnxruntime is a fake and the ONNX-backed classes are stubs.
|
||||
"""
|
||||
|
||||
import gc
|
||||
import importlib
|
||||
import os
|
||||
import sys
|
||||
import types
|
||||
import weakref
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
PACK_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
COMFY_ROOT = os.path.dirname(os.path.dirname(PACK_DIR))
|
||||
for path in (COMFY_ROOT, os.path.dirname(PACK_DIR)):
|
||||
if path not in sys.path:
|
||||
sys.path.insert(0, path)
|
||||
|
||||
folder_paths = importlib.import_module("folder_paths")
|
||||
backend = importlib.import_module("NH-Nodes.core.segformer_backend")
|
||||
model_files = importlib.import_module("NH-Nodes.core.model_files")
|
||||
ort_session = importlib.import_module("NH-Nodes.nodes.vton.preprocess.ort_session")
|
||||
vton = importlib.import_module("NH-Nodes.nodes.vton.vton_preprocessor_nodes")
|
||||
|
||||
VTON_FILES = (
|
||||
"dwpose/yolox_l.onnx", "dwpose/dw-ll_ucoco_384.onnx",
|
||||
"humanparsing/parsing_atr.onnx", "humanparsing/parsing_lip.onnx",
|
||||
)
|
||||
SEGFORMER_FILES = ("config.json", "model.safetensors", "preprocessor_config.json")
|
||||
GIB = 1024 ** 3
|
||||
|
||||
|
||||
# --- helpers -------------------------------------------------------------------------------------
|
||||
def _norm(path):
|
||||
return os.path.normcase(os.path.normpath(str(path)))
|
||||
|
||||
|
||||
def _touch(root, relative_files):
|
||||
for rel in relative_files:
|
||||
target = os.path.join(str(root), *rel.split("/"))
|
||||
os.makedirs(os.path.dirname(target), exist_ok=True)
|
||||
with open(target, "wb") as handle:
|
||||
handle.write(b"x")
|
||||
|
||||
|
||||
def _model_list(folder_name):
|
||||
"""ComfyUI's own file list for a registered folder name (what /models and the Model Library show)."""
|
||||
return sorted(name.replace("\\", "/") for name in folder_paths.get_filename_list(folder_name))
|
||||
|
||||
|
||||
class _FakeHub:
|
||||
"""huggingface_hub stand-in: records snapshot_download calls and 'downloads' `payload`."""
|
||||
|
||||
def __init__(self):
|
||||
self.calls = []
|
||||
self.payload = ()
|
||||
self.error = None
|
||||
|
||||
def snapshot_download(self, **kwargs):
|
||||
self.calls.append(kwargs)
|
||||
if self.error is not None:
|
||||
raise self.error
|
||||
_touch(kwargs["local_dir"], self.payload)
|
||||
return kwargs["local_dir"]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def hub(monkeypatch):
|
||||
fake = _FakeHub()
|
||||
monkeypatch.setitem(sys.modules, "huggingface_hub", types.SimpleNamespace(snapshot_download=fake.snapshot_download))
|
||||
return fake
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def models_root(tmp_path, monkeypatch):
|
||||
"""A fake ComfyUI install: empty models dir and a clean folder_paths registry under tmp_path."""
|
||||
root = tmp_path / "ComfyUI" / "models"
|
||||
root.mkdir(parents=True)
|
||||
monkeypatch.setattr(folder_paths, "models_dir", str(root))
|
||||
monkeypatch.setattr(folder_paths, "folder_names_and_paths", {})
|
||||
monkeypatch.setattr(folder_paths, "filename_list_cache", {}) # file lists of earlier tests point at deleted dirs
|
||||
monkeypatch.setattr(vton, "MODELS_DIR", str(root / "ComfyUI-Vton-Mask"))
|
||||
return root
|
||||
|
||||
|
||||
class _FakeComfyMemory:
|
||||
"""comfy.model_management stand-in that logs into a shared event list."""
|
||||
|
||||
reserved = 700 * 1024 ** 2 # what extra_reserved_memory() reports: VRAM ComfyUI keeps for itself
|
||||
|
||||
def __init__(self, events, device=None):
|
||||
self.events = events
|
||||
self.device = torch.device("cuda", 0) if device is None else device
|
||||
|
||||
def get_torch_device(self):
|
||||
return self.device
|
||||
|
||||
def extra_reserved_memory(self):
|
||||
return self.reserved
|
||||
|
||||
def free_memory(self, nbytes, device):
|
||||
self.events.append(("free_memory", nbytes, device))
|
||||
|
||||
def soft_empty_cache(self):
|
||||
self.events.append(("soft_empty_cache",))
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def cuda_spy(monkeypatch):
|
||||
"""Pretend CUDA exists and record the cache flushes the VTON node performs."""
|
||||
calls = []
|
||||
monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
|
||||
monkeypatch.setattr(torch.cuda, "empty_cache", lambda: calls.append("empty_cache"))
|
||||
real_collect = gc.collect
|
||||
|
||||
def collect(*args):
|
||||
calls.append("gc.collect")
|
||||
return real_collect(*args)
|
||||
|
||||
monkeypatch.setattr(gc, "collect", collect)
|
||||
return calls
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fake_ort(monkeypatch):
|
||||
"""onnxruntime stand-in. `available`: providers the installed wheel ships. `working`: providers that
|
||||
really start (a CUDA/cuDNN mismatch ships the CUDA provider but cannot start it).
|
||||
|
||||
Like the real one, a session is served by the first requested provider that works, else by the CPU.
|
||||
"""
|
||||
|
||||
class Session:
|
||||
def __init__(self, path, sess_options=None, providers=None):
|
||||
self.path = path
|
||||
self.active = [name for name in providers or () if name in fake.working] or ["CPUExecutionProvider"]
|
||||
|
||||
def get_providers(self):
|
||||
return list(self.active)
|
||||
|
||||
fake = types.SimpleNamespace(
|
||||
InferenceSession=Session,
|
||||
available=["CUDAExecutionProvider", "CPUExecutionProvider"],
|
||||
working={"CUDAExecutionProvider", "CPUExecutionProvider"},
|
||||
get_available_providers=lambda: list(fake.available),
|
||||
)
|
||||
monkeypatch.setitem(sys.modules, "onnxruntime", fake)
|
||||
return fake
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def isolated_ort_state(monkeypatch):
|
||||
"""Each test starts with no CPU fallback on record and no GPU session alive, whatever ran before it."""
|
||||
monkeypatch.setattr(ort_session, "_cpu_fallback_seen", False)
|
||||
monkeypatch.setattr(ort_session, "_warned_cpu_fallback", False)
|
||||
monkeypatch.setattr(ort_session, "_cuda_sessions", weakref.WeakSet())
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def stub_mask_location(monkeypatch):
|
||||
"""process() imports get_mask_location from utils_mask (needs scikit-image): use a trivial one."""
|
||||
monkeypatch.setitem(sys.modules, "NH-Nodes.core.utils_mask",
|
||||
types.SimpleNamespace(get_mask_location=_fake_get_mask_location))
|
||||
|
||||
|
||||
# --- SegFormer model folders ---------------------------------------------------------------------
|
||||
def test_pack_registers_its_folder_names_with_comfyui_at_import():
|
||||
assert _norm(os.path.join(folder_paths.models_dir, "NH-Nodes")) in map(_norm, folder_paths.get_folder_paths("nh_nodes"))
|
||||
assert _norm(vton.MODELS_DIR) in map(_norm, folder_paths.get_folder_paths("vton_mask"))
|
||||
|
||||
|
||||
def test_pack_limits_its_model_lists_to_model_files_at_import():
|
||||
registry = folder_paths.folder_names_and_paths
|
||||
assert ".onnx" in registry["vton_mask"][1]
|
||||
assert ".safetensors" in registry["nh_nodes"][1]
|
||||
|
||||
|
||||
def test_model_lists_show_model_files_only_not_hf_cache_metadata(models_root):
|
||||
_touch(vton.MODELS_DIR, [*VTON_FILES, ".gitattributes", "README.md",
|
||||
".cache/huggingface/download/dwpose/yolox_l.onnx.metadata",
|
||||
".cache/huggingface/download/dwpose/yolox_l.onnx.lock",
|
||||
".cache/huggingface/.gitignore"])
|
||||
_touch(models_root / "NH-Nodes", [
|
||||
"segformer_fashion/config.json", "segformer_fashion/model.safetensors", "segformer_fashion/preprocessor_config.json",
|
||||
"segformer_fashion/.cache/huggingface/download/model.safetensors.metadata",
|
||||
"segformer_fashion/.gitattributes", "face_parsing/model.safetensors"])
|
||||
vton.resolve_vton_models_dir(download=False) # registers vton_mask, as the pack does at start-up
|
||||
backend.model_search_dirs("segformer_fashion") # registers nh_nodes
|
||||
|
||||
assert _model_list("vton_mask") == sorted(VTON_FILES)
|
||||
assert _model_list("nh_nodes") == ["face_parsing/model.safetensors", "segformer_fashion/model.safetensors"]
|
||||
|
||||
|
||||
def test_extension_filter_joins_a_folder_name_the_user_already_registered(models_root, tmp_path):
|
||||
shared = tmp_path / "shared" / "ComfyUI-Vton-Mask"
|
||||
_touch(shared, [*VTON_FILES, ".gitattributes"])
|
||||
folder_paths.add_model_folder_path("vton_mask", str(shared)) # what extra_model_paths.yaml does: no extension filter
|
||||
assert folder_paths.folder_names_and_paths["vton_mask"][1] == set() # i.e. ComfyUI would list every file
|
||||
|
||||
vton.resolve_vton_models_dir(download=False)
|
||||
|
||||
paths, extensions = folder_paths.folder_names_and_paths["vton_mask"]
|
||||
assert [_norm(p) for p in paths] == [_norm(shared), _norm(vton.MODELS_DIR)] # the user's location is kept, first
|
||||
assert extensions == {".onnx"}
|
||||
assert _model_list("vton_mask") == sorted(VTON_FILES)
|
||||
|
||||
|
||||
def test_extension_filter_does_not_change_a_set_comfyui_shares_between_folders(models_root, tmp_path):
|
||||
shared_extensions = {".pt"} # like folder_paths.supported_pt_extensions, one set behind many folder names
|
||||
folder_paths.folder_names_and_paths["vton_mask"] = ([str(tmp_path / "other")], shared_extensions)
|
||||
|
||||
vton.resolve_vton_models_dir(download=False)
|
||||
|
||||
assert shared_extensions == {".pt"}
|
||||
assert folder_paths.folder_names_and_paths["vton_mask"][1] == {".pt", ".onnx"}
|
||||
|
||||
|
||||
def test_registering_the_extension_filter_twice_changes_nothing(models_root):
|
||||
default = str(models_root / "NH-Nodes")
|
||||
for _ in range(2):
|
||||
model_files.register_model_folder("nh_nodes", default, {".safetensors"})
|
||||
assert folder_paths.folder_names_and_paths["nh_nodes"] == ([default], {".safetensors"})
|
||||
|
||||
|
||||
def test_a_failing_extension_filter_does_not_lose_the_registered_locations(models_root, tmp_path, monkeypatch, capsys):
|
||||
class ReadOnlyRegistry(dict):
|
||||
def __setitem__(self, key, value):
|
||||
raise RuntimeError("registry is read-only")
|
||||
|
||||
user_dir = str(tmp_path / "shared" / "NH-Nodes")
|
||||
registry = ReadOnlyRegistry()
|
||||
dict.__setitem__(registry, "nh_nodes", ([user_dir], set()))
|
||||
monkeypatch.setattr(folder_paths, "folder_names_and_paths", registry)
|
||||
|
||||
default = str(models_root / "NH-Nodes")
|
||||
assert model_files.register_model_folder("nh_nodes", default, {".safetensors"}) == [user_dir, default]
|
||||
assert "read-only" in capsys.readouterr().out
|
||||
|
||||
|
||||
def test_default_search_order_is_unchanged_for_existing_installs(models_root):
|
||||
expected = {
|
||||
"segformer_fashion": ["NH-Nodes/segformer_fashion", "RMBG/segformer_fashion"],
|
||||
"face_parsing": ["NH-Nodes/face_parsing", "face_parsing", "RMBG/segformer_face"],
|
||||
}
|
||||
for model_name, relative in expected.items():
|
||||
assert [_norm(p) for p in backend.model_search_dirs(model_name)] == [_norm(models_root / r) for r in relative]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("legacy", [("segformer_fashion", "RMBG/segformer_fashion"),
|
||||
("face_parsing", "face_parsing"),
|
||||
("face_parsing", "RMBG/segformer_face"),
|
||||
("segformer_fashion", "NH-Nodes/segformer_fashion")])
|
||||
def test_models_already_on_disk_are_found_without_download(models_root, hub, legacy):
|
||||
model_name, relative = legacy
|
||||
_touch(models_root / relative, SEGFORMER_FILES)
|
||||
assert _norm(backend.resolve_model_dir(model_name)) == _norm(models_root / relative)
|
||||
assert hub.calls == []
|
||||
|
||||
|
||||
def test_extra_model_paths_location_is_found_and_wins_over_the_default(models_root, tmp_path, hub):
|
||||
shared = tmp_path / "shared" / "NH-Nodes"
|
||||
_touch(shared / "segformer_fashion", SEGFORMER_FILES)
|
||||
# ComfyUI applies extra_model_paths.yaml before custom nodes load, so the user's path is registered first.
|
||||
folder_paths.add_model_folder_path("nh_nodes", str(shared))
|
||||
_touch(models_root / "NH-Nodes" / "segformer_fashion", SEGFORMER_FILES)
|
||||
|
||||
assert _norm(backend.resolve_model_dir("segformer_fashion")) == _norm(shared / "segformer_fashion")
|
||||
assert hub.calls == []
|
||||
# The default location stays searched, so models that already live there are never orphaned.
|
||||
assert _norm(models_root / "NH-Nodes") in map(_norm, folder_paths.get_folder_paths("nh_nodes"))
|
||||
|
||||
|
||||
def test_registering_twice_does_not_duplicate_the_default(models_root):
|
||||
default = str(models_root / "NH-Nodes")
|
||||
assert model_files.register_model_folder("nh_nodes", default) == [default]
|
||||
assert model_files.register_model_folder("nh_nodes", default) == [default]
|
||||
assert folder_paths.get_folder_paths("nh_nodes") == [default]
|
||||
|
||||
|
||||
def test_unique_dirs_keeps_the_first_spelling_and_the_order(tmp_path):
|
||||
a, b = tmp_path / "a", tmp_path / "b"
|
||||
spellings = [str(a), str(b), str(a / ".." / "a"), str(b) + os.sep, str(tmp_path / "c")]
|
||||
assert model_files.unique_dirs(spellings) == [str(a), str(b), str(tmp_path / "c")]
|
||||
assert model_files.unique_dirs([]) == []
|
||||
|
||||
|
||||
def test_segformer_and_vton_use_the_same_model_helpers():
|
||||
for name in ("register_model_folder", "download_pinned", "make_room_on_gpu"):
|
||||
assert getattr(backend, name) is getattr(model_files, name) is getattr(vton, name)
|
||||
|
||||
|
||||
def test_incomplete_model_folder_is_skipped_and_the_next_candidate_is_used(models_root, hub):
|
||||
_touch(models_root / "NH-Nodes" / "segformer_fashion", ["config.json"]) # interrupted download
|
||||
_touch(models_root / "RMBG" / "segformer_fashion", SEGFORMER_FILES)
|
||||
assert _norm(backend.resolve_model_dir("segformer_fashion", download=False)) == _norm(models_root / "RMBG" / "segformer_fashion")
|
||||
(models_root / "RMBG" / "segformer_fashion" / "model.safetensors").unlink()
|
||||
assert backend.resolve_model_dir("segformer_fashion", download=False) is None
|
||||
assert hub.calls == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_name, repo_id, revision", [
|
||||
("segformer_fashion", "1038lab/segformer_fashion", "e5bae2d8e3333be9f12eb6aae454c0542c1371d3"),
|
||||
("face_parsing", "jonathandinu/face-parsing", "758b82e15a0178c9db39c1ff666a8b56e3a550c8"),
|
||||
])
|
||||
def test_segformer_download_is_pinned_and_lands_in_the_default_folder(models_root, tmp_path, hub, model_name, repo_id, revision):
|
||||
hub.payload = SEGFORMER_FILES
|
||||
# An extra location without the model must not change where new downloads go.
|
||||
folder_paths.add_model_folder_path("nh_nodes", str(tmp_path / "empty-shared"))
|
||||
|
||||
resolved = backend.resolve_model_dir(model_name)
|
||||
|
||||
assert _norm(resolved) == _norm(models_root / "NH-Nodes" / model_name)
|
||||
assert len(hub.calls) == 1
|
||||
call = hub.calls[0]
|
||||
assert call["repo_id"] == repo_id
|
||||
assert call["revision"] == revision
|
||||
assert call["allow_patterns"] == list(SEGFORMER_FILES) # safetensors only: no pytorch_model.bin / onnx
|
||||
assert _norm(call["local_dir"]) == _norm(resolved)
|
||||
assert "local_dir_use_symlinks" not in call
|
||||
|
||||
|
||||
def test_segformer_download_missing_files_and_failures_are_reported(models_root, hub):
|
||||
hub.payload = SEGFORMER_FILES[:2]
|
||||
with pytest.raises(RuntimeError, match="missing after download"):
|
||||
backend.resolve_model_dir("face_parsing")
|
||||
hub.error = OSError("network down")
|
||||
with pytest.raises(RuntimeError, match="1038lab/segformer_fashion") as info:
|
||||
backend.resolve_model_dir("segformer_fashion")
|
||||
assert "network down" in str(info.value) and isinstance(info.value.__cause__, OSError)
|
||||
|
||||
|
||||
def test_extra_model_paths_yaml_end_to_end(models_root, tmp_path, hub):
|
||||
"""Feed a real extra_model_paths.yaml through ComfyUI's own loader."""
|
||||
yaml = pytest.importorskip("yaml")
|
||||
extra_config = pytest.importorskip("utils.extra_config")
|
||||
if not hasattr(extra_config, "load_extra_path_config"):
|
||||
pytest.skip("another 'utils' module shadows ComfyUI's")
|
||||
shared = tmp_path / "shared"
|
||||
_touch(shared / "NH-Nodes" / "face_parsing", [*SEGFORMER_FILES, ".gitattributes"])
|
||||
_touch(shared / "ComfyUI-Vton-Mask", [*VTON_FILES, ".gitattributes"])
|
||||
config = tmp_path / "extra_model_paths.yaml"
|
||||
config.write_text(yaml.safe_dump({"my_drive": {
|
||||
"base_path": str(shared), "nh_nodes": "NH-Nodes/", "vton_mask": "ComfyUI-Vton-Mask/"}}), encoding="utf-8")
|
||||
extra_config.load_extra_path_config(str(config))
|
||||
|
||||
assert _norm(backend.resolve_model_dir("face_parsing")) == _norm(shared / "NH-Nodes" / "face_parsing")
|
||||
assert _norm(vton.resolve_vton_models_dir()) == _norm(shared / "ComfyUI-Vton-Mask")
|
||||
assert hub.calls == []
|
||||
# The user's folders are listed like ours: model files only.
|
||||
assert _model_list("vton_mask") == sorted(VTON_FILES)
|
||||
assert _model_list("nh_nodes") == ["face_parsing/model.safetensors"]
|
||||
|
||||
|
||||
# --- VTON model folders and downloads ------------------------------------------------------------
|
||||
def test_vton_folder_with_missing_files_is_not_complete(models_root, hub):
|
||||
_touch(vton.MODELS_DIR, VTON_FILES[:2]) # an interrupted first download leaves a half-filled folder
|
||||
assert vton.resolve_vton_models_dir(download=False) is None
|
||||
_touch(vton.MODELS_DIR, VTON_FILES)
|
||||
assert _norm(vton.resolve_vton_models_dir(download=False)) == _norm(vton.MODELS_DIR)
|
||||
assert hub.calls == []
|
||||
|
||||
|
||||
def test_interrupted_download_is_completed_instead_of_skipped(models_root, hub):
|
||||
_touch(vton.MODELS_DIR, VTON_FILES[:2])
|
||||
hub.payload = VTON_FILES
|
||||
resolved = vton.resolve_vton_models_dir()
|
||||
assert len(hub.calls) == 1
|
||||
assert all(os.path.isfile(os.path.join(resolved, *name.split("/"))) for name in VTON_FILES)
|
||||
|
||||
|
||||
def test_vton_download_is_pinned_filtered_and_symlink_free(models_root, hub):
|
||||
hub.payload = VTON_FILES
|
||||
vton.resolve_vton_models_dir()
|
||||
# An exact match also proves the deprecated local_dir_use_symlinks argument is gone.
|
||||
assert hub.calls == [{
|
||||
"repo_id": "kg-09/kg-vton-mask",
|
||||
"revision": "4648778286cac2dca9aab8da167ea5257462487d",
|
||||
"allow_patterns": ["dwpose/*.onnx", "humanparsing/*.onnx"],
|
||||
"local_dir": vton.MODELS_DIR,
|
||||
}]
|
||||
|
||||
|
||||
def test_vton_download_that_leaves_files_missing_is_an_error(models_root, hub):
|
||||
hub.payload = VTON_FILES[:3]
|
||||
with pytest.raises(RuntimeError, match="parsing_lip.onnx"):
|
||||
vton.resolve_vton_models_dir()
|
||||
|
||||
|
||||
def test_vton_download_failure_explains_where_to_get_the_files(models_root, hub):
|
||||
hub.error = OSError("network down")
|
||||
with pytest.raises(RuntimeError, match="kg-09/kg-vton-mask") as info:
|
||||
vton.resolve_vton_models_dir()
|
||||
assert "network down" in str(info.value) and isinstance(info.value.__cause__, OSError)
|
||||
|
||||
|
||||
def test_vton_models_in_extra_location_are_used_without_download(models_root, tmp_path, hub):
|
||||
shared = tmp_path / "shared" / "ComfyUI-Vton-Mask"
|
||||
_touch(shared, VTON_FILES)
|
||||
folder_paths.add_model_folder_path("vton_mask", str(shared)) # what an extra_model_paths.yaml entry does
|
||||
assert _norm(vton.resolve_vton_models_dir()) == _norm(shared)
|
||||
assert hub.calls == []
|
||||
|
||||
|
||||
def test_vton_download_goes_to_the_default_folder_even_with_extra_locations(models_root, tmp_path, hub):
|
||||
hub.payload = VTON_FILES
|
||||
folder_paths.add_model_folder_path("vton_mask", str(tmp_path / "empty-shared"))
|
||||
assert _norm(vton.resolve_vton_models_dir()) == _norm(vton.MODELS_DIR)
|
||||
|
||||
|
||||
def test_real_install_still_finds_its_models_in_the_old_places():
|
||||
"""Existing installs must not need to move anything (skipped where the weights are absent)."""
|
||||
models = folder_paths.models_dir
|
||||
legacy = {
|
||||
"segformer_fashion": ("NH-Nodes/segformer_fashion", "RMBG/segformer_fashion"),
|
||||
"face_parsing": ("NH-Nodes/face_parsing", "face_parsing", "RMBG/segformer_face"),
|
||||
}
|
||||
checked = 0
|
||||
for model_name, relative_dirs in legacy.items():
|
||||
if any(all(os.path.isfile(os.path.join(models, rel, f)) for f in SEGFORMER_FILES) for rel in relative_dirs):
|
||||
assert backend.resolve_model_dir(model_name, download=False) is not None
|
||||
checked += 1
|
||||
default_vton = os.path.join(models, "ComfyUI-Vton-Mask")
|
||||
if all(os.path.isfile(os.path.join(default_vton, *name.split("/"))) for name in VTON_FILES):
|
||||
assert vton.resolve_vton_models_dir(download=False) is not None
|
||||
checked += 1
|
||||
if not checked:
|
||||
pytest.skip("no model weights on disk")
|
||||
|
||||
|
||||
# --- keep_model_loaded: input definition ----------------------------------------------------------
|
||||
def test_keep_model_loaded_is_a_new_optional_boolean_and_nothing_else_changed():
|
||||
spec = vton.NH_VTonUltimateProcessor.INPUT_TYPES()
|
||||
assert list(spec["required"]) == ["human_image", "category", "mask_feathering", "cover_shoes",
|
||||
"refine_hands", "refine_hair", "device"]
|
||||
assert list(spec["optional"]) == ["offset_top", "offset_bottom", "offset_left", "offset_right", "keep_model_loaded"]
|
||||
kind, options = spec["optional"]["keep_model_loaded"]
|
||||
assert kind == "BOOLEAN" and options["default"] is False
|
||||
assert "VRAM" in options["tooltip"] and "cuda" in options["tooltip"]
|
||||
assert vton.NH_VTonUltimateProcessor.RETURN_TYPES == ("MASK", "MASK", "IMAGE", "IMAGE", "IMAGE", "MASK", "MASK", "MASK")
|
||||
assert vton.NH_VTonUltimateProcessor.RETURN_NAMES == (
|
||||
"final_mask", "agnostic_mask", "densepose_image", "parsing_image", "masked_image",
|
||||
"parsing_map_raw", "hair_mask", "hands_mask")
|
||||
|
||||
|
||||
def test_keep_model_loaded_tooltip_tells_list_users_why_to_turn_it_on():
|
||||
tooltip = vton.NH_VTonUltimateProcessor.INPUT_TYPES()["optional"]["keep_model_loaded"][1]["tooltip"]
|
||||
# ComfyUI calls a node once per item of an image list, and every call reloads the models while this is off.
|
||||
assert "list of images" in tooltip and "once per item" in tooltip
|
||||
assert "much faster" in tooltip and "batch" in tooltip # a batch tensor is one call; lists and repeated runs benefit
|
||||
assert "CPU" in tooltip # it has no effect on models that run on the CPU, including after a fallback
|
||||
|
||||
|
||||
# --- keep_model_loaded: releasing the sessions ----------------------------------------------------
|
||||
class _FakePose:
|
||||
"""DWposeDetector stand-in. It references itself so only a real gc.collect() can free it, and it owns
|
||||
a session made by create_session, so ort_session knows whether this model lives on the GPU."""
|
||||
|
||||
def __init__(self, device="cpu"):
|
||||
self.me = self
|
||||
self.session = ort_session.create_session("dwpose/yolox_l.onnx", device)
|
||||
|
||||
def __call__(self, bgr):
|
||||
height, width = bgr.shape[:2]
|
||||
return np.zeros((height, width, 3), np.uint8), None, None, np.full((1, 134, 2), 0.5, np.float32)
|
||||
|
||||
|
||||
class _ExplodingPose(_FakePose):
|
||||
def __call__(self, bgr):
|
||||
raise RuntimeError("boom")
|
||||
|
||||
|
||||
class _FakeParsing:
|
||||
def __call__(self, image):
|
||||
labels = np.zeros(np.array(image).shape[:2], np.uint8)
|
||||
labels[:, : labels.shape[1] // 2] = 5
|
||||
return Image.fromarray(labels), Image.fromarray(labels)
|
||||
|
||||
|
||||
def _fake_get_mask_location(category, parsing_map, keypoints, width, height, *offsets):
|
||||
return Image.fromarray(((np.array(parsing_map) == 5) * 255).astype(np.uint8)), None
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def cached_models(monkeypatch, fake_ort, stub_mask_location):
|
||||
"""Factory that installs fake cached sessions for a device and returns a weakref to the pose model.
|
||||
|
||||
The session is served by whatever `fake_ort.working` allows when it is installed: set that first to
|
||||
simulate a cuda request that fell back to the CPU. Only the weakref escapes, so a test never keeps
|
||||
the sessions alive by accident.
|
||||
"""
|
||||
|
||||
def install(device, pose_cls=_FakePose):
|
||||
pose = pose_cls(device)
|
||||
monkeypatch.setattr(vton, "LOADED_VTON_MODELS",
|
||||
{"dwprocessor": pose, "parsing_model": _FakeParsing(), "device": device})
|
||||
return weakref.ref(pose)
|
||||
|
||||
return install
|
||||
|
||||
|
||||
def _process(device, **kwargs):
|
||||
image = torch.full((1, 48, 64, 3), 0.5)
|
||||
return vton.NH_VTonUltimateProcessor().process(image, "Upper-body", 4, False, True, True, device, **kwargs)
|
||||
|
||||
|
||||
def test_cuda_run_releases_the_sessions_by_default(cached_models, cuda_spy):
|
||||
pose_ref = cached_models("cuda")
|
||||
assert ort_session.cuda_sessions_alive() # served by CUDAExecutionProvider: these pin VRAM
|
||||
outputs = _process("cuda")
|
||||
|
||||
assert len(outputs) == 8
|
||||
assert vton.LOADED_VTON_MODELS is None
|
||||
assert pose_ref() is None
|
||||
assert not ort_session.cuda_sessions_alive()
|
||||
assert "gc.collect" in cuda_spy and "empty_cache" in cuda_spy
|
||||
assert cuda_spy.index("gc.collect") < cuda_spy.index("empty_cache") # collect first, then hand VRAM back
|
||||
|
||||
|
||||
def test_keep_model_loaded_keeps_the_cache(cached_models, cuda_spy):
|
||||
pose_ref = cached_models("cuda")
|
||||
cached = vton.LOADED_VTON_MODELS
|
||||
_process("cuda", keep_model_loaded=True)
|
||||
assert vton.LOADED_VTON_MODELS is cached and pose_ref() is not None
|
||||
assert cuda_spy == []
|
||||
_process("cuda", keep_model_loaded=True) # still cached and usable
|
||||
assert vton.LOADED_VTON_MODELS is cached
|
||||
|
||||
|
||||
def test_cuda_run_with_keep_off_after_a_keep_on_run_releases(cached_models, cuda_spy):
|
||||
pose_ref = cached_models("cuda")
|
||||
_process("cuda", keep_model_loaded=True)
|
||||
assert pose_ref() is not None
|
||||
_process("cuda", keep_model_loaded=False)
|
||||
assert pose_ref() is None and vton.LOADED_VTON_MODELS is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("keep", [False, True])
|
||||
def test_cpu_behaviour_is_unchanged_models_stay_cached(cached_models, cuda_spy, keep):
|
||||
pose_ref = cached_models("cpu")
|
||||
cached = vton.LOADED_VTON_MODELS
|
||||
_process("cpu", keep_model_loaded=keep)
|
||||
assert vton.LOADED_VTON_MODELS is cached and pose_ref() is not None
|
||||
assert cuda_spy == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize("keep", [False, True])
|
||||
def test_cuda_request_that_onnxruntime_served_from_the_cpu_keeps_its_cache(cached_models, fake_ort, cuda_spy, keep):
|
||||
fake_ort.working = {"CPUExecutionProvider"} # onnxruntime-gpu is installed but CUDA/cuDNN do not start
|
||||
pose_ref = cached_models("cuda")
|
||||
cached = vton.LOADED_VTON_MODELS
|
||||
assert ort_session.cpu_fallback_seen() and not ort_session.cuda_sessions_alive()
|
||||
|
||||
_process("cuda", keep_model_loaded=keep)
|
||||
_process("cuda", keep_model_loaded=keep) # the next item of a list: still the same, usable sessions
|
||||
|
||||
# Nothing sits in VRAM: dropping the sessions would only make every list item reload four models.
|
||||
assert vton.LOADED_VTON_MODELS is cached and pose_ref() is not None
|
||||
assert cuda_spy == [] # no gc.collect(), no torch.cuda.empty_cache()
|
||||
|
||||
|
||||
def test_failed_cuda_run_served_from_the_cpu_keeps_the_cache(cached_models, fake_ort, cuda_spy):
|
||||
fake_ort.working = {"CPUExecutionProvider"}
|
||||
pose_ref = cached_models("cuda", pose_cls=_ExplodingPose)
|
||||
with pytest.raises(RuntimeError, match="boom"):
|
||||
_process("cuda")
|
||||
assert vton.LOADED_VTON_MODELS is not None and pose_ref() is not None
|
||||
assert cuda_spy == []
|
||||
|
||||
|
||||
def test_a_gpu_cache_is_still_released_once_cpu_sessions_exist_next_to_it(cached_models, fake_ort, cuda_spy):
|
||||
"""The decision follows the sessions that are alive, not a flag of the last session created."""
|
||||
pose_ref = cached_models("cuda") # on the GPU
|
||||
fake_ort.working = {"CPUExecutionProvider"}
|
||||
ort_session.create_session("other.onnx", "cuda") # a later session that fell back
|
||||
_process("cuda")
|
||||
assert vton.LOADED_VTON_MODELS is None and pose_ref() is None
|
||||
assert "empty_cache" in cuda_spy
|
||||
|
||||
|
||||
def test_sessions_are_released_even_when_processing_fails(cached_models, cuda_spy):
|
||||
pose_ref = cached_models("cuda", pose_cls=_ExplodingPose)
|
||||
with pytest.raises(RuntimeError, match="boom"):
|
||||
_process("cuda") # the traceback is still alive in this block: it must not pin the sessions
|
||||
assert vton.LOADED_VTON_MODELS is None
|
||||
assert pose_ref() is None
|
||||
assert "empty_cache" in cuda_spy
|
||||
|
||||
|
||||
def test_failed_run_on_cpu_keeps_the_cache(cached_models, cuda_spy):
|
||||
pose_ref = cached_models("cpu", pose_cls=_ExplodingPose)
|
||||
with pytest.raises(RuntimeError, match="boom"):
|
||||
_process("cpu")
|
||||
assert pose_ref() is not None and cuda_spy == []
|
||||
|
||||
|
||||
def test_a_failing_cache_flush_does_not_hide_the_original_error(cached_models, cuda_spy, monkeypatch, capsys):
|
||||
def broken_flush():
|
||||
raise RuntimeError("CUDA error: device-side assert")
|
||||
|
||||
pose_ref = cached_models("cuda", pose_cls=_ExplodingPose)
|
||||
monkeypatch.setattr(torch.cuda, "empty_cache", broken_flush)
|
||||
with pytest.raises(RuntimeError, match="boom"): # the run's own error, not the flush error
|
||||
_process("cuda")
|
||||
assert vton.LOADED_VTON_MODELS is None and pose_ref() is None
|
||||
assert "device-side assert" in capsys.readouterr().out
|
||||
|
||||
|
||||
def test_release_without_cuda_skips_empty_cache(monkeypatch):
|
||||
calls = []
|
||||
monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
|
||||
monkeypatch.setattr(torch.cuda, "empty_cache", lambda: calls.append("empty_cache"))
|
||||
monkeypatch.setattr(vton, "LOADED_VTON_MODELS", {"device": "cuda"})
|
||||
vton.release_vton_models()
|
||||
assert vton.LOADED_VTON_MODELS is None and calls == []
|
||||
|
||||
|
||||
def test_dropping_sessions_that_ran_on_the_cpu_leaves_torchs_cache_alone(cached_models, fake_ort, cuda_spy):
|
||||
fake_ort.working = {"CPUExecutionProvider"}
|
||||
pose_ref = cached_models("cuda") # a cuda request that onnxruntime served from the CPU
|
||||
vton.release_vton_models()
|
||||
assert vton.LOADED_VTON_MODELS is None and pose_ref() is None # still dropped (and collected) ...
|
||||
assert "gc.collect" in cuda_spy and "empty_cache" not in cuda_spy # ... but there is no VRAM to hand back
|
||||
|
||||
|
||||
def test_dropping_sessions_that_ran_on_the_gpu_hands_the_vram_back(cached_models, cuda_spy):
|
||||
pose_ref = cached_models("cuda")
|
||||
vton.release_vton_models()
|
||||
assert vton.LOADED_VTON_MODELS is None and pose_ref() is None
|
||||
assert cuda_spy == ["gc.collect", "empty_cache"]
|
||||
|
||||
|
||||
# --- making room before creating CUDA sessions ----------------------------------------------------
|
||||
@pytest.fixture
|
||||
def stub_loaders(monkeypatch, models_root, hub, fake_ort, stub_mask_location):
|
||||
"""load_model() with stub ONNX classes and no model files yet; returns the event log.
|
||||
|
||||
The stubs open their sessions with create_session, so the fake onnxruntime decides where they run.
|
||||
"""
|
||||
events = []
|
||||
monkeypatch.setattr(vton, "LOADED_VTON_MODELS", None)
|
||||
monkeypatch.setattr(ort_session, "ensure_onnxruntime", lambda: None)
|
||||
|
||||
class Detector(_FakePose):
|
||||
def __init__(self, model_root, device):
|
||||
super().__init__(device)
|
||||
events.append(("detector", device, model_root))
|
||||
|
||||
class Parsing(_FakeParsing):
|
||||
def __init__(self, model_root, device):
|
||||
self.session = ort_session.create_session("humanparsing/parsing_atr.onnx", device)
|
||||
events.append(("parsing", device, model_root))
|
||||
|
||||
monkeypatch.setitem(sys.modules, "NH-Nodes.nodes.vton.preprocess.dwpose", types.SimpleNamespace(DWposeDetector=Detector))
|
||||
monkeypatch.setitem(sys.modules, "NH-Nodes.nodes.vton.preprocess.humanparsing.run_parsing", types.SimpleNamespace(Parsing=Parsing))
|
||||
return events
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def ready_loaders(stub_loaders):
|
||||
"""stub_loaders plus a complete model folder in the default location."""
|
||||
_touch(vton.MODELS_DIR, VTON_FILES)
|
||||
return stub_loaders
|
||||
|
||||
|
||||
def test_cuda_load_asks_comfyui_for_room_before_creating_sessions(ready_loaders, monkeypatch, cuda_spy):
|
||||
comfy = _FakeComfyMemory(ready_loaders)
|
||||
monkeypatch.setattr(model_files, "_mm", comfy)
|
||||
vton.NH_VTonUltimateProcessor().load_model("cuda")
|
||||
|
||||
assert [event[0] for event in ready_loaders] == ["free_memory", "soft_empty_cache", "detector", "parsing"]
|
||||
_, nbytes, device = ready_loaders[0]
|
||||
assert vton._VTON_VRAM_ESTIMATE >= 2 * GIB # the four sessions hold about 2.9 GB after a run
|
||||
assert nbytes == vton._VTON_VRAM_ESTIMATE + comfy.reserved # on top of the reserve ComfyUI keeps for itself
|
||||
assert device == torch.device("cuda", 0) # ComfyUI's device, with the index its loaded models carry
|
||||
|
||||
|
||||
def test_cpu_load_never_touches_comfyui_memory(ready_loaders, monkeypatch, cuda_spy):
|
||||
monkeypatch.setattr(model_files, "_mm", _FakeComfyMemory(ready_loaders))
|
||||
vton.NH_VTonUltimateProcessor().load_model("cpu")
|
||||
assert [event[0] for event in ready_loaders] == ["detector", "parsing"]
|
||||
|
||||
|
||||
def test_cuda_request_without_the_cuda_provider_does_not_evict_comfyui_models(ready_loaders, fake_ort, monkeypatch, cuda_spy):
|
||||
monkeypatch.setattr(model_files, "_mm", _FakeComfyMemory(ready_loaders))
|
||||
fake_ort.available = ["CPUExecutionProvider"] # CPU-only onnxruntime wheel
|
||||
fake_ort.working = {"CPUExecutionProvider"}
|
||||
vton.NH_VTonUltimateProcessor().load_model("cuda")
|
||||
assert [event[0] for event in ready_loaders] == ["detector", "parsing"]
|
||||
|
||||
|
||||
def test_cuda_run_that_falls_back_asks_for_room_once_then_stops_evicting(ready_loaders, fake_ort, monkeypatch, cuda_spy):
|
||||
monkeypatch.setattr(model_files, "_mm", _FakeComfyMemory(ready_loaders))
|
||||
fake_ort.working = {"CPUExecutionProvider"} # onnxruntime-gpu is installed, but CUDA/cuDNN do not start
|
||||
|
||||
# First run: nothing says yet that the GPU is unusable, so ComfyUI is asked, once.
|
||||
assert len(_process("cuda")) == 8
|
||||
assert [event[0] for event in ready_loaders] == ["free_memory", "soft_empty_cache", "detector", "parsing"]
|
||||
assert ort_session.cpu_fallback_seen()
|
||||
assert vton.LOADED_VTON_MODELS is not None and cuda_spy == [] # CPU sessions are not released
|
||||
|
||||
# Every later load in this process skips it: whatever it would evict cannot help sessions on the CPU.
|
||||
ready_loaders.clear()
|
||||
monkeypatch.setattr(vton, "LOADED_VTON_MODELS", None)
|
||||
_process("cuda")
|
||||
vton.NH_VTonUltimateProcessor().load_model("cpu")
|
||||
monkeypatch.setattr(vton, "LOADED_VTON_MODELS", None)
|
||||
vton.NH_VTonUltimateProcessor().load_model("cuda")
|
||||
assert [event[0] for event in ready_loaders] == ["detector", "parsing"] * 3
|
||||
|
||||
|
||||
def test_cuda_runs_on_the_gpu_keep_asking_for_room_and_release(ready_loaders, monkeypatch, cuda_spy):
|
||||
monkeypatch.setattr(model_files, "_mm", _FakeComfyMemory(ready_loaders))
|
||||
for _ in range(2):
|
||||
ready_loaders.clear()
|
||||
_process("cuda")
|
||||
assert [event[0] for event in ready_loaders] == ["free_memory", "soft_empty_cache", "detector", "parsing"]
|
||||
assert vton.LOADED_VTON_MODELS is None # released after the run, so the next one reloads
|
||||
assert not ort_session.cpu_fallback_seen() and not ort_session.cuda_sessions_alive()
|
||||
assert "empty_cache" in cuda_spy
|
||||
|
||||
|
||||
def test_load_model_uses_the_resolved_folder_and_caches_per_device(stub_loaders, tmp_path):
|
||||
shared = tmp_path / "shared" / "ComfyUI-Vton-Mask"
|
||||
_touch(shared, VTON_FILES) # the default folder stays empty
|
||||
folder_paths.add_model_folder_path("vton_mask", str(shared))
|
||||
|
||||
node = vton.NH_VTonUltimateProcessor()
|
||||
first = node.load_model("cpu")
|
||||
assert [_norm(event[2]) for event in stub_loaders] == [_norm(shared)] * 2
|
||||
assert node.load_model("cpu") is first and len(stub_loaders) == 2 # cache hit: nothing is rebuilt
|
||||
|
||||
|
||||
def test_switching_device_drops_the_old_sessions_first(ready_loaders, monkeypatch, cuda_spy):
|
||||
monkeypatch.setattr(model_files, "_mm", _FakeComfyMemory(ready_loaders))
|
||||
node = vton.NH_VTonUltimateProcessor()
|
||||
cuda_ref = weakref.ref(node.load_model("cuda")["dwprocessor"])
|
||||
ready_loaders.clear()
|
||||
cpu_models = node.load_model("cpu")
|
||||
assert cpu_models["device"] == "cpu" and cuda_ref() is None
|
||||
assert [event[0] for event in ready_loaders] == ["detector", "parsing"]
|
||||
|
||||
|
||||
def test_switching_device_away_from_cpu_fallback_sessions_does_not_flush_torchs_cache(ready_loaders, fake_ort, monkeypatch, cuda_spy):
|
||||
monkeypatch.setattr(model_files, "_mm", _FakeComfyMemory(ready_loaders))
|
||||
fake_ort.working = {"CPUExecutionProvider"}
|
||||
node = vton.NH_VTonUltimateProcessor()
|
||||
cuda_ref = weakref.ref(node.load_model("cuda")["dwprocessor"])
|
||||
cuda_spy.clear()
|
||||
|
||||
assert node.load_model("cpu")["device"] == "cpu"
|
||||
assert cuda_ref() is None and "empty_cache" not in cuda_spy
|
||||
|
||||
|
||||
def test_missing_onnxruntime_still_fails_before_any_download(models_root, hub, monkeypatch):
|
||||
def not_installed():
|
||||
raise RuntimeError("ONNX Runtime missing")
|
||||
|
||||
monkeypatch.setattr(vton, "LOADED_VTON_MODELS", None)
|
||||
monkeypatch.setattr(ort_session, "ensure_onnxruntime", not_installed)
|
||||
with pytest.raises(RuntimeError, match="ONNX Runtime"):
|
||||
vton.NH_VTonUltimateProcessor().load_model("cuda")
|
||||
assert hub.calls == []
|
||||
|
||||
|
||||
def test_cuda_provider_available_reflects_the_onnxruntime_build(monkeypatch):
|
||||
fake = types.SimpleNamespace(get_available_providers=lambda: ["CPUExecutionProvider"])
|
||||
monkeypatch.setitem(sys.modules, "onnxruntime", fake)
|
||||
assert ort_session.cuda_provider_available() is False
|
||||
fake.get_available_providers = lambda: ["TensorrtExecutionProvider", "CUDAExecutionProvider", "CPUExecutionProvider"]
|
||||
assert ort_session.cuda_provider_available() is True
|
||||
|
||||
|
||||
# --- where the sessions really run -------------------------------------------------------------------
|
||||
def test_create_session_records_which_sessions_the_gpu_serves(fake_ort):
|
||||
gpu = ort_session.create_session("a.onnx", "cuda")
|
||||
cpu = ort_session.create_session("b.onnx", "cpu")
|
||||
|
||||
assert gpu.get_providers()[0] == "CUDAExecutionProvider" and cpu.get_providers() == ["CPUExecutionProvider"]
|
||||
assert ort_session.cuda_sessions_alive() and not ort_session.cpu_fallback_seen()
|
||||
del gpu
|
||||
gc.collect()
|
||||
assert not ort_session.cuda_sessions_alive() # the record never keeps a session (and its VRAM) alive
|
||||
|
||||
|
||||
def test_a_cuda_request_served_by_the_cpu_is_a_fallback_not_a_gpu_session(fake_ort, capsys):
|
||||
fake_ort.working = {"CPUExecutionProvider"}
|
||||
session = ort_session.create_session("a.onnx", "cuda")
|
||||
|
||||
assert session.get_providers() == ["CPUExecutionProvider"]
|
||||
assert ort_session.cpu_fallback_seen() and not ort_session.cuda_sessions_alive()
|
||||
assert "a.onnx on CPUExecutionProvider" in capsys.readouterr().out # the existing warning is still printed
|
||||
|
||||
|
||||
def test_the_fallback_is_remembered_even_after_its_warning_was_printed(fake_ort, monkeypatch):
|
||||
fake_ort.working = {"CPUExecutionProvider"}
|
||||
monkeypatch.setattr(ort_session, "_warned_cpu_fallback", True) # the one-time warning is already used up
|
||||
ort_session.create_session("a.onnx", "cuda")
|
||||
assert ort_session.cpu_fallback_seen()
|
||||
|
||||
|
||||
def test_a_session_that_reports_no_provider_counts_as_a_fallback(fake_ort):
|
||||
fake_ort.InferenceSession.get_providers = lambda self: []
|
||||
ort_session.create_session("a.onnx", "cuda")
|
||||
assert ort_session.cpu_fallback_seen() and not ort_session.cuda_sessions_alive()
|
||||
|
||||
|
||||
def test_cpu_requests_are_never_a_fallback(fake_ort):
|
||||
fake_ort.working = {"CPUExecutionProvider"}
|
||||
ort_session.create_session("a.onnx", "cpu")
|
||||
assert not ort_session.cpu_fallback_seen() and not ort_session.cuda_sessions_alive()
|
||||
|
||||
|
||||
# --- the shared GPU helper -------------------------------------------------------------------------
|
||||
def test_make_room_is_a_noop_without_cuda_in_comfy_cpu_mode_or_outside_comfyui(monkeypatch):
|
||||
events = []
|
||||
monkeypatch.setattr(model_files, "_mm", _FakeComfyMemory(events))
|
||||
monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
|
||||
model_files.make_room_on_gpu(GIB)
|
||||
|
||||
monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
|
||||
monkeypatch.setattr(model_files, "_mm", _FakeComfyMemory(events, device=torch.device("cpu"))) # ComfyUI --cpu
|
||||
model_files.make_room_on_gpu(GIB)
|
||||
|
||||
monkeypatch.setattr(model_files, "_mm", None)
|
||||
model_files.make_room_on_gpu(GIB)
|
||||
assert events == []
|
||||
|
||||
|
||||
def test_make_room_asks_for_the_size_plus_the_reserve_comfyui_keeps_for_itself(monkeypatch):
|
||||
events = []
|
||||
comfy = _FakeComfyMemory(events)
|
||||
monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
|
||||
monkeypatch.setattr(model_files, "_mm", comfy)
|
||||
model_files.make_room_on_gpu(GIB)
|
||||
# ComfyUI's own loaders ask for their size plus extra_reserved_memory(); asking for less would leave us short.
|
||||
assert events == [("free_memory", GIB + comfy.reserved, torch.device("cuda", 0)), ("soft_empty_cache",)]
|
||||
|
||||
|
||||
def test_make_room_on_a_comfyui_without_extra_reserved_memory_asks_for_the_size_alone(monkeypatch):
|
||||
events = []
|
||||
|
||||
class OlderComfy(_FakeComfyMemory):
|
||||
extra_reserved_memory = None # the function does not exist in older ComfyUI builds
|
||||
|
||||
monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
|
||||
monkeypatch.setattr(model_files, "_mm", OlderComfy(events))
|
||||
model_files.make_room_on_gpu(GIB)
|
||||
assert events[0] == ("free_memory", GIB, torch.device("cuda", 0))
|
||||
|
||||
|
||||
def test_make_room_reports_comfyui_errors_instead_of_failing_the_run(monkeypatch, capsys):
|
||||
class Broken(_FakeComfyMemory):
|
||||
def free_memory(self, nbytes, device):
|
||||
raise RuntimeError("model list changed")
|
||||
|
||||
monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
|
||||
monkeypatch.setattr(model_files, "_mm", Broken([]))
|
||||
model_files.make_room_on_gpu(GIB) # must not raise
|
||||
assert "model list changed" in capsys.readouterr().out
|
||||
|
||||
|
||||
class _FakeSegformer:
|
||||
def __init__(self, device):
|
||||
self.device = device
|
||||
self.moved_to = []
|
||||
|
||||
def parameters(self):
|
||||
return iter([torch.zeros(1000)]) # 4000 bytes
|
||||
|
||||
def to(self, device):
|
||||
self.moved_to.append(device)
|
||||
self.device = device
|
||||
return self
|
||||
|
||||
|
||||
def test_segformer_asks_comfyui_for_room_using_comfyuis_own_device(monkeypatch):
|
||||
events = []
|
||||
comfy = _FakeComfyMemory(events)
|
||||
monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
|
||||
monkeypatch.setattr(model_files, "_mm", comfy)
|
||||
model = _FakeSegformer(torch.device("cpu"))
|
||||
backend._ensure_on_device({"model": model, "device": torch.device("cuda")})
|
||||
|
||||
# ComfyUI's loaded models carry cuda:0; an index-less cuda device matches none of them.
|
||||
# The weights (4000 bytes) are asked for twice over, plus the reserve ComfyUI keeps for itself.
|
||||
assert events[0] == ("free_memory", 8000 + comfy.reserved, torch.device("cuda", 0))
|
||||
assert model.moved_to == [torch.device("cuda")]
|
||||
|
||||
|
||||
def test_segformer_already_on_the_gpu_is_not_moved_or_evicted_again(monkeypatch):
|
||||
events = []
|
||||
monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
|
||||
monkeypatch.setattr(model_files, "_mm", _FakeComfyMemory(events))
|
||||
model = _FakeSegformer(torch.device("cuda", 0)) # what .to("cuda") leaves behind
|
||||
backend._ensure_on_device({"model": model, "device": torch.device("cuda")})
|
||||
assert events == [] and model.moved_to == []
|
||||
@@ -1,19 +1,11 @@
|
||||
import { app } from "../../../scripts/app.js";
|
||||
import { chainCallback, clampInt, findWidget, refreshNode } from "./nh_widget_utils.js";
|
||||
|
||||
const NODE_NAME = "NH_AnyListSplit";
|
||||
const MAX_INDEX_OUTPUTS = 10;
|
||||
|
||||
function findWidget(node, name) {
|
||||
return node.widgets?.find((widget) => widget.name === name);
|
||||
}
|
||||
|
||||
function toInt(value, fallback) {
|
||||
const parsed = Number.parseInt(value, 10);
|
||||
return Number.isFinite(parsed) ? parsed : fallback;
|
||||
}
|
||||
|
||||
function clampIndexCount(value) {
|
||||
return Math.max(0, Math.min(MAX_INDEX_OUTPUTS, toInt(value, 4)));
|
||||
return clampInt(value, 0, MAX_INDEX_OUTPUTS, 4);
|
||||
}
|
||||
|
||||
function indexOutputNumber(output) {
|
||||
@@ -51,9 +43,7 @@ function updateOutputs(node) {
|
||||
}
|
||||
}
|
||||
|
||||
node.setSize?.(node.computeSize?.() || node.size);
|
||||
node.setDirtyCanvas?.(true, true);
|
||||
app.graph?.setDirtyCanvas?.(true, true);
|
||||
refreshNode(node, app.graph);
|
||||
}
|
||||
|
||||
function installControls(node) {
|
||||
@@ -69,11 +59,7 @@ function installControls(node) {
|
||||
|
||||
const countWidget = findWidget(node, "index_count");
|
||||
if (countWidget) {
|
||||
const oldCallback = countWidget.callback;
|
||||
countWidget.callback = function () {
|
||||
oldCallback?.apply(this, arguments);
|
||||
updateOutputs(node);
|
||||
};
|
||||
chainCallback(countWidget, () => updateOutputs(node));
|
||||
}
|
||||
|
||||
setTimeout(() => updateOutputs(node), 0);
|
||||
|
||||
+7
-52
@@ -1,4 +1,5 @@
|
||||
import { app } from "../../../scripts/app.js";
|
||||
import { chainCallback, clampInt, findWidget, hideWidget, refreshNode, showWidget } from "./nh_widget_utils.js";
|
||||
|
||||
const MAX_ITEMS = 64;
|
||||
|
||||
@@ -63,42 +64,8 @@ const NODE_CONFIGS = {
|
||||
},
|
||||
};
|
||||
|
||||
function findWidget(node, name) {
|
||||
return node.widgets?.find((widget) => widget.name === name);
|
||||
}
|
||||
|
||||
function toInt(value, fallback) {
|
||||
const parsed = Number.parseInt(value, 10);
|
||||
return Number.isFinite(parsed) ? parsed : fallback;
|
||||
}
|
||||
|
||||
function clampCount(value, minCount = 1) {
|
||||
return Math.max(minCount, Math.min(MAX_ITEMS, toInt(value, minCount)));
|
||||
}
|
||||
|
||||
function hideWidget(widget) {
|
||||
if (!widget || widget.__nhIndexHidden) {
|
||||
return;
|
||||
}
|
||||
|
||||
widget.__nhIndexHidden = true;
|
||||
widget.__nhIndexType = widget.type;
|
||||
widget.__nhIndexComputeSize = widget.computeSize;
|
||||
widget.computeSize = () => [0, -4];
|
||||
widget.type = "nh-hidden";
|
||||
}
|
||||
|
||||
function showWidget(widget) {
|
||||
if (!widget || !widget.__nhIndexHidden) {
|
||||
return;
|
||||
}
|
||||
|
||||
widget.type = widget.__nhIndexType;
|
||||
widget.computeSize = widget.__nhIndexComputeSize;
|
||||
|
||||
delete widget.__nhIndexHidden;
|
||||
delete widget.__nhIndexType;
|
||||
delete widget.__nhIndexComputeSize;
|
||||
return clampInt(value, minCount, MAX_ITEMS, minCount);
|
||||
}
|
||||
|
||||
function updateIndexedWidgets(node, config) {
|
||||
@@ -137,9 +104,7 @@ function updateIndexedWidgets(node, config) {
|
||||
}
|
||||
}
|
||||
|
||||
node.setSize?.(node.computeSize?.() || node.size);
|
||||
node.setDirtyCanvas?.(true, true);
|
||||
app.graph?.setDirtyCanvas?.(true, true);
|
||||
refreshNode(node, app.graph);
|
||||
}
|
||||
|
||||
function installIndexedControls(node, config) {
|
||||
@@ -155,21 +120,13 @@ function installIndexedControls(node, config) {
|
||||
|
||||
const countWidget = findWidget(node, config.countWidget);
|
||||
if (countWidget) {
|
||||
const oldCallback = countWidget.callback;
|
||||
countWidget.callback = function () {
|
||||
oldCallback?.apply(this, arguments);
|
||||
updateIndexedWidgets(node, config);
|
||||
};
|
||||
chainCallback(countWidget, () => updateIndexedWidgets(node, config));
|
||||
}
|
||||
|
||||
if (config.selectModeWidget) {
|
||||
const smWidget = findWidget(node, config.selectModeWidget);
|
||||
if (smWidget) {
|
||||
const smOld = smWidget.callback;
|
||||
smWidget.callback = function () {
|
||||
smOld?.apply(this, arguments);
|
||||
updateIndexedWidgets(node, config);
|
||||
};
|
||||
chainCallback(smWidget, () => updateIndexedWidgets(node, config));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -177,9 +134,7 @@ function installIndexedControls(node, config) {
|
||||
for (let i = 1; i <= MAX_ITEMS; i++) {
|
||||
const enableWidget = findWidget(node, `${config.enablePrefix}_${i}`);
|
||||
if (enableWidget) {
|
||||
const enOld = enableWidget.callback;
|
||||
enableWidget.callback = function () {
|
||||
enOld?.apply(this, arguments);
|
||||
chainCallback(enableWidget, function () {
|
||||
if (this.value) {
|
||||
for (let j = 1; j <= MAX_ITEMS; j++) {
|
||||
if (j !== i) {
|
||||
@@ -191,7 +146,7 @@ function installIndexedControls(node, config) {
|
||||
}
|
||||
updateIndexedWidgets(node, config);
|
||||
}
|
||||
};
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+209
-37
@@ -10,12 +10,12 @@ function ensurePreviewButton(node) {
|
||||
}
|
||||
|
||||
node.__nhLargePreviewButton = node.addWidget("button", "Open Viewer", "Open Viewer", () => {
|
||||
const manifest = node.__nhLargePreviewManifest?.[0];
|
||||
if (!manifest) {
|
||||
const manifests = node.__nhLargePreviewManifest;
|
||||
if (!manifests?.length) {
|
||||
alert("NH Large Image Preview has no viewer data yet. Queue the node first.");
|
||||
return;
|
||||
}
|
||||
openManifest(manifest);
|
||||
openManifest(manifests, 0);
|
||||
});
|
||||
}
|
||||
|
||||
@@ -25,12 +25,12 @@ function ensureCompareButton(node) {
|
||||
}
|
||||
|
||||
node.__nhLargeCompareButton = node.addWidget("button", "Open Compare", "Open Compare", () => {
|
||||
const manifest = node.__nhLargeCompareManifest?.[0];
|
||||
if (!manifest) {
|
||||
const manifests = node.__nhLargeCompareManifest;
|
||||
if (!manifests?.length) {
|
||||
alert("NH Large Image Compare has no viewer data yet. Queue the node first.");
|
||||
return;
|
||||
}
|
||||
openCompareManifest(manifest);
|
||||
openCompareManifest(manifests, 0);
|
||||
});
|
||||
}
|
||||
|
||||
@@ -63,9 +63,12 @@ function tileUrl(manifest, level, x, y) {
|
||||
.replace("{y}", `${y}`));
|
||||
}
|
||||
|
||||
function openManifest(manifest) {
|
||||
if (!manifest?.levels?.length) {
|
||||
const url = thumbUrl(manifest);
|
||||
// A run can hold a whole batch: the viewers get every manifest of it and the one to start on.
|
||||
function openManifest(manifests, startIndex = 0) {
|
||||
const list = Array.isArray(manifests) ? manifests : [manifests];
|
||||
const tiled = list.filter((item) => item?.levels?.length);
|
||||
if (!tiled.length) {
|
||||
const url = thumbUrl(list[startIndex] || list[0]);
|
||||
if (url) {
|
||||
window.open(url, "_blank", "noopener,noreferrer");
|
||||
return;
|
||||
@@ -73,15 +76,17 @@ function openManifest(manifest) {
|
||||
alert("NH Large Image Preview has no thumbnail or tile data.");
|
||||
return;
|
||||
}
|
||||
openTiledViewer(manifest);
|
||||
openTiledViewer(tiled, Math.max(0, tiled.indexOf(list[startIndex])));
|
||||
}
|
||||
|
||||
function openCompareManifest(compare) {
|
||||
if (!compare?.a || !compare?.b) {
|
||||
function openCompareManifest(compares, startIndex = 0) {
|
||||
const list = Array.isArray(compares) ? compares : [compares];
|
||||
const pairs = list.filter((item) => item?.a && item?.b);
|
||||
if (!pairs.length) {
|
||||
alert("NH Large Image Compare has no A/B tile data.");
|
||||
return;
|
||||
}
|
||||
openCompareViewer(compare);
|
||||
openCompareViewer(pairs, Math.max(0, pairs.indexOf(list[startIndex])));
|
||||
}
|
||||
|
||||
function makeButton(label, onClick) {
|
||||
@@ -100,15 +105,87 @@ function makeButton(label, onClick) {
|
||||
return button;
|
||||
}
|
||||
|
||||
function openTiledViewer(manifest) {
|
||||
// Cancels the tile requests still running and forgets the tiles: the image was left, or the viewer closed.
|
||||
function clearTileCache(tileCache) {
|
||||
for (const item of tileCache.values()) {
|
||||
item.image.onload = null;
|
||||
item.image.onerror = null;
|
||||
item.image.src = "";
|
||||
}
|
||||
tileCache.clear();
|
||||
}
|
||||
|
||||
function setButtonEnabled(button, enabled) {
|
||||
button.disabled = !enabled;
|
||||
button.style.opacity = enabled ? "1" : "0.4";
|
||||
button.style.cursor = enabled ? "pointer" : "default";
|
||||
}
|
||||
|
||||
// Batch navigation. The part that needs no DOM lives here so both viewers share it.
|
||||
function clampBatchIndex(index, count) {
|
||||
const last = Math.max(0, count - 1);
|
||||
const value = Number.isFinite(index) ? Math.trunc(index) : 0;
|
||||
return Math.min(last, Math.max(0, value));
|
||||
}
|
||||
|
||||
// -1 or +1 for an unmodified Left or Right arrow that was not pressed in a text field or slider, otherwise 0.
|
||||
function batchStepForKey(event) {
|
||||
if (event.altKey || event.ctrlKey || event.metaKey || event.shiftKey) {
|
||||
return 0;
|
||||
}
|
||||
const tag = (event.target?.tagName || "").toLowerCase();
|
||||
if (event.target?.isContentEditable || tag === "input" || tag === "textarea" || tag === "select") {
|
||||
return 0;
|
||||
}
|
||||
if (event.key === "ArrowLeft") {
|
||||
return -1;
|
||||
}
|
||||
return event.key === "ArrowRight" ? 1 : 0;
|
||||
}
|
||||
|
||||
// The "previous / n of N / next" group of the toolbar; update(index) refreshes the counter and the end buttons.
|
||||
function makeBatchNav(count, onStep) {
|
||||
const element = document.createElement("div");
|
||||
element.style.cssText = "display:flex;gap:6px;align-items:center";
|
||||
const previous = makeButton("\u25C0", () => onStep(-1));
|
||||
previous.title = "Previous image (Left arrow)";
|
||||
const next = makeButton("\u25B6", () => onStep(1));
|
||||
next.title = "Next image (Right arrow)";
|
||||
const label = document.createElement("span");
|
||||
label.style.cssText = "min-width:54px;text-align:center;color:#ddd;font:12px system-ui,sans-serif";
|
||||
element.appendChild(previous);
|
||||
element.appendChild(label);
|
||||
element.appendChild(next);
|
||||
return {
|
||||
element,
|
||||
update(index) {
|
||||
label.textContent = `${index + 1} / ${count}`;
|
||||
setButtonEnabled(previous, index > 0);
|
||||
setButtonEnabled(next, index < count - 1);
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
function openTiledViewer(items, startIndex = 0) {
|
||||
let index = clampBatchIndex(startIndex, items.length);
|
||||
let manifest = items[index];
|
||||
let levels = [];
|
||||
let fullLevel = null;
|
||||
let closed = false;
|
||||
const overlay = document.createElement("div");
|
||||
let closeOnEscape;
|
||||
let onKeyDown;
|
||||
let onMouseMove;
|
||||
let onMouseUp;
|
||||
const closeViewer = () => {
|
||||
closed = true;
|
||||
if (state.raf) {
|
||||
cancelAnimationFrame(state.raf);
|
||||
state.raf = 0;
|
||||
}
|
||||
clearTileCache(tileCache);
|
||||
overlay.remove();
|
||||
if (closeOnEscape) {
|
||||
document.removeEventListener("keydown", closeOnEscape);
|
||||
if (onKeyDown) {
|
||||
document.removeEventListener("keydown", onKeyDown);
|
||||
}
|
||||
if (onMouseMove) {
|
||||
window.removeEventListener("mousemove", onMouseMove);
|
||||
@@ -119,10 +196,12 @@ function openTiledViewer(manifest) {
|
||||
window.removeEventListener("resize", resizeCanvas);
|
||||
};
|
||||
|
||||
overlay.tabIndex = -1; // focusable, so the arrow keys reach the viewer even if a text field was focused before
|
||||
overlay.style.cssText = [
|
||||
"position:fixed",
|
||||
"inset:0",
|
||||
"z-index:100000",
|
||||
"outline:none",
|
||||
"background:rgba(0,0,0,0.82)",
|
||||
"display:flex",
|
||||
"align-items:center",
|
||||
@@ -156,7 +235,6 @@ function openTiledViewer(manifest) {
|
||||
].join(";");
|
||||
|
||||
const title = document.createElement("div");
|
||||
title.textContent = `${manifest.title || "NH Large Image Preview"} | ${manifest.width}x${manifest.height} | ${manifest.levels.length} levels`;
|
||||
title.style.cssText = "overflow:hidden;text-overflow:ellipsis;white-space:nowrap";
|
||||
|
||||
const buttons = document.createElement("div");
|
||||
@@ -198,8 +276,6 @@ function openTiledViewer(manifest) {
|
||||
document.body.appendChild(overlay);
|
||||
|
||||
const ctx = canvas.getContext("2d");
|
||||
const levels = [...manifest.levels].sort((a, b) => a.scale - b.scale);
|
||||
const fullLevel = levels[levels.length - 1];
|
||||
const tileCache = new Map();
|
||||
const state = {
|
||||
zoom: 1,
|
||||
@@ -211,6 +287,28 @@ function openTiledViewer(manifest) {
|
||||
raf: 0,
|
||||
};
|
||||
|
||||
function loadManifest() {
|
||||
manifest = items[index];
|
||||
levels = [...manifest.levels].sort((a, b) => a.scale - b.scale);
|
||||
fullLevel = levels[levels.length - 1];
|
||||
clearTileCache(tileCache);
|
||||
title.textContent = `${manifest.title || "NH Large Image Preview"} | ${manifest.width}x${manifest.height} | ${levels.length} levels`;
|
||||
if (nav) {
|
||||
nav.update(index);
|
||||
}
|
||||
}
|
||||
|
||||
// Every image starts fitted: the zoom and pan of the one before are not carried over.
|
||||
function showManifest(nextIndex) {
|
||||
const target = clampBatchIndex(nextIndex, items.length);
|
||||
if (target === index) {
|
||||
return;
|
||||
}
|
||||
index = target;
|
||||
loadManifest();
|
||||
fit();
|
||||
}
|
||||
|
||||
function getCanvasSize() {
|
||||
return {
|
||||
width: Math.max(1, canvas.clientWidth || canvasWrap.clientWidth || 1),
|
||||
@@ -340,11 +438,15 @@ function openTiledViewer(manifest) {
|
||||
}
|
||||
|
||||
function scheduleDraw() {
|
||||
if (!state.raf) {
|
||||
if (!state.raf && !closed) {
|
||||
state.raf = requestAnimationFrame(draw);
|
||||
}
|
||||
}
|
||||
|
||||
const nav = items.length > 1 ? makeBatchNav(items.length, (delta) => showManifest(index + delta)) : null;
|
||||
if (nav) {
|
||||
buttons.appendChild(nav.element);
|
||||
}
|
||||
buttons.appendChild(makeButton("Fit", fit));
|
||||
buttons.appendChild(makeButton("100%", setHundredPercent));
|
||||
buttons.appendChild(makeButton("-", () => zoomAt(canvas.clientWidth / 2, canvas.clientHeight / 2, 0.8)));
|
||||
@@ -389,31 +491,53 @@ function openTiledViewer(manifest) {
|
||||
closeViewer();
|
||||
}
|
||||
});
|
||||
closeOnEscape = function (event) {
|
||||
if (event.key === "Escape" && document.body.contains(overlay)) {
|
||||
onKeyDown = function (event) {
|
||||
if (!document.body.contains(overlay)) {
|
||||
return;
|
||||
}
|
||||
if (event.key === "Escape") {
|
||||
closeViewer();
|
||||
return;
|
||||
}
|
||||
const step = items.length > 1 ? batchStepForKey(event) : 0;
|
||||
if (step) {
|
||||
// The arrow keys belong to the viewer while it is open, not to the canvas behind it.
|
||||
event.preventDefault();
|
||||
event.stopPropagation();
|
||||
showManifest(index + step);
|
||||
}
|
||||
};
|
||||
document.addEventListener("keydown", closeOnEscape);
|
||||
document.addEventListener("keydown", onKeyDown);
|
||||
window.addEventListener("resize", resizeCanvas);
|
||||
resizeCanvas();
|
||||
loadManifest();
|
||||
fit();
|
||||
overlay.focus({ preventScroll: true });
|
||||
}
|
||||
|
||||
function openCompareViewer(compare) {
|
||||
const manifestA = compare.a;
|
||||
const manifestB = compare.b;
|
||||
const baseWidth = Math.max(manifestA.width || 1, manifestB.width || 1);
|
||||
const baseHeight = Math.max(manifestA.height || 1, manifestB.height || 1);
|
||||
function openCompareViewer(items, startIndex = 0) {
|
||||
let index = clampBatchIndex(startIndex, items.length);
|
||||
let compare = items[index];
|
||||
let manifestA = compare.a;
|
||||
let manifestB = compare.b;
|
||||
let baseWidth = 1;
|
||||
let baseHeight = 1;
|
||||
let closed = false;
|
||||
|
||||
const overlay = document.createElement("div");
|
||||
let closeOnEscape;
|
||||
let onKeyDown;
|
||||
let onMouseMove;
|
||||
let onMouseUp;
|
||||
const closeViewer = () => {
|
||||
closed = true;
|
||||
if (state.raf) {
|
||||
cancelAnimationFrame(state.raf);
|
||||
state.raf = 0;
|
||||
}
|
||||
clearTileCache(tileCache);
|
||||
overlay.remove();
|
||||
if (closeOnEscape) {
|
||||
document.removeEventListener("keydown", closeOnEscape);
|
||||
if (onKeyDown) {
|
||||
document.removeEventListener("keydown", onKeyDown);
|
||||
}
|
||||
if (onMouseMove) {
|
||||
window.removeEventListener("mousemove", onMouseMove);
|
||||
@@ -424,10 +548,12 @@ function openCompareViewer(compare) {
|
||||
window.removeEventListener("resize", resizeCanvas);
|
||||
};
|
||||
|
||||
overlay.tabIndex = -1; // focusable, so the arrow keys reach the viewer even if a text field was focused before
|
||||
overlay.style.cssText = [
|
||||
"position:fixed",
|
||||
"inset:0",
|
||||
"z-index:100000",
|
||||
"outline:none",
|
||||
"background:rgba(0,0,0,0.84)",
|
||||
"display:flex",
|
||||
"align-items:center",
|
||||
@@ -461,7 +587,6 @@ function openCompareViewer(compare) {
|
||||
].join(";");
|
||||
|
||||
const title = document.createElement("div");
|
||||
title.textContent = `${compare.title || "NH Large Image Compare"} | A ${manifestA.width}x${manifestA.height} | B ${manifestB.width}x${manifestB.height}`;
|
||||
title.style.cssText = "overflow:hidden;text-overflow:ellipsis;white-space:nowrap";
|
||||
|
||||
const buttons = document.createElement("div");
|
||||
@@ -525,6 +650,32 @@ function openCompareViewer(compare) {
|
||||
raf: 0,
|
||||
};
|
||||
|
||||
function loadPair() {
|
||||
compare = items[index];
|
||||
manifestA = compare.a;
|
||||
manifestB = compare.b;
|
||||
baseWidth = Math.max(manifestA.width || 1, manifestB.width || 1);
|
||||
baseHeight = Math.max(manifestA.height || 1, manifestB.height || 1);
|
||||
frameLayouts.clear();
|
||||
levelSets.clear();
|
||||
clearTileCache(tileCache);
|
||||
title.textContent = `${compare.title || "NH Large Image Compare"} | A ${manifestA.width}x${manifestA.height} | B ${manifestB.width}x${manifestB.height}`;
|
||||
if (nav) {
|
||||
nav.update(index);
|
||||
}
|
||||
}
|
||||
|
||||
// Every pair starts fitted; the compare mode and the split position stay as they are.
|
||||
function showPair(nextIndex) {
|
||||
const target = clampBatchIndex(nextIndex, items.length);
|
||||
if (target === index) {
|
||||
return;
|
||||
}
|
||||
index = target;
|
||||
loadPair();
|
||||
fit();
|
||||
}
|
||||
|
||||
function getCanvasSize() {
|
||||
return {
|
||||
width: Math.max(1, canvas.clientWidth || canvasWrap.clientWidth || 1),
|
||||
@@ -786,7 +937,7 @@ function openCompareViewer(compare) {
|
||||
}
|
||||
|
||||
function scheduleDraw() {
|
||||
if (!state.raf) {
|
||||
if (!state.raf && !closed) {
|
||||
state.raf = requestAnimationFrame(draw);
|
||||
}
|
||||
}
|
||||
@@ -802,6 +953,10 @@ function openCompareViewer(compare) {
|
||||
scheduleDraw();
|
||||
};
|
||||
|
||||
const nav = items.length > 1 ? makeBatchNav(items.length, (delta) => showPair(index + delta)) : null;
|
||||
if (nav) {
|
||||
buttons.appendChild(nav.element);
|
||||
}
|
||||
buttons.appendChild(makeButton("Fit", fit));
|
||||
buttons.appendChild(makeButton("100%", setHundredPercent));
|
||||
buttons.appendChild(makeButton("Slider", () => setMode("slider")));
|
||||
@@ -833,6 +988,10 @@ function openCompareViewer(compare) {
|
||||
onMouseUp = () => {
|
||||
state.dragging = false;
|
||||
canvasWrap.style.cursor = "grab";
|
||||
if (document.activeElement === splitInput) {
|
||||
// A slider that was just dragged would keep the arrow keys; hand them back to the batch navigation.
|
||||
splitInput.blur();
|
||||
}
|
||||
};
|
||||
window.addEventListener("mousemove", onMouseMove);
|
||||
window.addEventListener("mouseup", onMouseUp);
|
||||
@@ -849,15 +1008,28 @@ function openCompareViewer(compare) {
|
||||
closeViewer();
|
||||
}
|
||||
});
|
||||
closeOnEscape = function (event) {
|
||||
if (event.key === "Escape" && document.body.contains(overlay)) {
|
||||
onKeyDown = function (event) {
|
||||
if (!document.body.contains(overlay)) {
|
||||
return;
|
||||
}
|
||||
if (event.key === "Escape") {
|
||||
closeViewer();
|
||||
return;
|
||||
}
|
||||
const step = items.length > 1 ? batchStepForKey(event) : 0;
|
||||
if (step) {
|
||||
// The arrow keys belong to the viewer while it is open, not to the canvas behind it.
|
||||
event.preventDefault();
|
||||
event.stopPropagation();
|
||||
showPair(index + step);
|
||||
}
|
||||
};
|
||||
document.addEventListener("keydown", closeOnEscape);
|
||||
document.addEventListener("keydown", onKeyDown);
|
||||
window.addEventListener("resize", resizeCanvas);
|
||||
resizeCanvas();
|
||||
loadPair();
|
||||
setMode(state.mode);
|
||||
overlay.focus({ preventScroll: true });
|
||||
}
|
||||
|
||||
app.registerExtension({
|
||||
|
||||
@@ -1,19 +1,11 @@
|
||||
import { app } from "../../../scripts/app.js";
|
||||
import { chainCallback, clampInt, findWidget, refreshNode, toInt } from "./nh_widget_utils.js";
|
||||
|
||||
const NODE_NAME = "NH_LayerStackComposite";
|
||||
const MAX_LAYERS = 64;
|
||||
|
||||
function findWidget(node, name) {
|
||||
return node.widgets?.find((widget) => widget.name === name);
|
||||
}
|
||||
|
||||
function toInt(value, fallback) {
|
||||
const parsed = Number.parseInt(value, 10);
|
||||
return Number.isFinite(parsed) ? parsed : fallback;
|
||||
}
|
||||
|
||||
function clampLayerCount(value) {
|
||||
return Math.max(0, Math.min(MAX_LAYERS, toInt(value, 1)));
|
||||
return clampInt(value, 0, MAX_LAYERS, 1);
|
||||
}
|
||||
|
||||
function findInputIndex(node, name) {
|
||||
@@ -39,6 +31,46 @@ function addMissingLayerInputs(node, count) {
|
||||
}
|
||||
}
|
||||
|
||||
// graph.links is a Map in current frontends (a Proxy that still allows links[id]) and a plain object keyed by
|
||||
// link id in older ones.
|
||||
function findLink(graph, linkId) {
|
||||
if (linkId == null || !graph) {
|
||||
return null;
|
||||
}
|
||||
if (typeof graph.getLink === "function") {
|
||||
return graph.getLink(linkId) ?? null;
|
||||
}
|
||||
const links = graph.links;
|
||||
if (links instanceof Map) {
|
||||
return links.get(linkId) ?? null;
|
||||
}
|
||||
return links?.[linkId] ?? null;
|
||||
}
|
||||
|
||||
// The link attached to the input at `index`. Current frontends derive it from the index (input.link is only a
|
||||
// deprecated view of that), older ones keep the link id on the input.
|
||||
function getInputLink(node, index) {
|
||||
if (!node.graph) {
|
||||
return null; // removed from the graph before the timer fired: nothing can be linked to it
|
||||
}
|
||||
if (typeof node.getInputLink === "function") {
|
||||
return node.getInputLink(index) ?? null;
|
||||
}
|
||||
return findLink(node.graph, node.inputs[index]?.link);
|
||||
}
|
||||
|
||||
// A link remembers the index of the input it ends on (target_slot). Move every link onto its input's new index.
|
||||
function retargetLinks(moves, slotCount) {
|
||||
// Park the links on free slots first: current frontends refuse to move a link onto a slot that another link
|
||||
// still holds, which is what happens when two linked inputs change places.
|
||||
moves.forEach(({ link }, offset) => {
|
||||
link.target_slot = slotCount + offset;
|
||||
});
|
||||
moves.forEach(({ link, newIndex }) => {
|
||||
link.target_slot = newIndex;
|
||||
});
|
||||
}
|
||||
|
||||
function sortLayerInputs(node) {
|
||||
const inputs = node.inputs || [];
|
||||
const background = inputs.filter((input) => input.name === "background_image");
|
||||
@@ -48,7 +80,25 @@ function sortLayerInputs(node) {
|
||||
const others = inputs.filter(
|
||||
(input) => input.name !== "background_image" && !/^layer_\d+$/.test(input.name)
|
||||
);
|
||||
node.inputs = [...background, ...layers, ...others];
|
||||
const sorted = [...background, ...layers, ...others];
|
||||
if (sorted.every((input, index) => input === inputs[index])) {
|
||||
return;
|
||||
}
|
||||
|
||||
// Read the links before reordering: current frontends find an input's link through its index.
|
||||
const moves = [];
|
||||
sorted.forEach((input, newIndex) => {
|
||||
const oldIndex = inputs.indexOf(input);
|
||||
const link = oldIndex === newIndex ? null : getInputLink(node, oldIndex);
|
||||
if (link) {
|
||||
moves.push({ link, newIndex });
|
||||
}
|
||||
});
|
||||
|
||||
node.inputs = sorted;
|
||||
retargetLinks(moves, sorted.length);
|
||||
node.setDirtyCanvas?.(true, true);
|
||||
(node.graph ?? app.graph)?.setDirtyCanvas?.(true, true);
|
||||
}
|
||||
|
||||
function updateLayerInputs(node) {
|
||||
@@ -62,9 +112,7 @@ function updateLayerInputs(node) {
|
||||
addMissingLayerInputs(node, count);
|
||||
sortLayerInputs(node);
|
||||
|
||||
node.setSize?.(node.computeSize?.() || node.size);
|
||||
node.setDirtyCanvas?.(true, true);
|
||||
app.graph?.setDirtyCanvas?.(true, true);
|
||||
refreshNode(node, app.graph);
|
||||
}
|
||||
|
||||
function installLayerStackControls(node) {
|
||||
@@ -80,11 +128,7 @@ function installLayerStackControls(node) {
|
||||
|
||||
const layerCountWidget = findWidget(node, "layer_count");
|
||||
if (layerCountWidget) {
|
||||
const oldCallback = layerCountWidget.callback;
|
||||
layerCountWidget.callback = function () {
|
||||
oldCallback?.apply(this, arguments);
|
||||
updateLayerInputs(node);
|
||||
};
|
||||
chainCallback(layerCountWidget, () => updateLayerInputs(node));
|
||||
}
|
||||
|
||||
setTimeout(() => updateLayerInputs(node), 0);
|
||||
|
||||
@@ -1,17 +1,9 @@
|
||||
import { app } from "../../../scripts/app.js";
|
||||
import { findWidget, markDirty, toInt } from "./nh_widget_utils.js";
|
||||
|
||||
const NODE_NAME = "NH_LoadImagesFromFolder";
|
||||
const COUNT_ENDPOINT = "/nh-nodes/load-images-folder/count";
|
||||
|
||||
function findWidget(node, name) {
|
||||
return node.widgets?.find((widget) => widget.name === name);
|
||||
}
|
||||
|
||||
function toInt(value, fallback) {
|
||||
const parsed = Number.parseInt(value, 10);
|
||||
return Number.isFinite(parsed) ? parsed : fallback;
|
||||
}
|
||||
|
||||
function wrapIndex(value, count) {
|
||||
const index = Math.trunc(value);
|
||||
if (count > 0) {
|
||||
@@ -20,9 +12,13 @@ function wrapIndex(value, count) {
|
||||
return Math.max(0, index);
|
||||
}
|
||||
|
||||
async function getImageCount(folderPath) {
|
||||
// How many images the node sees in the folder. `recursive` has to be the node's own setting: a recursive node sees the
|
||||
// images of the sub-folders too, so counting only the top level would wrap the index too early. The count is asked for
|
||||
// on every queue, so a changed folder_path or recursive value is picked up without a cache to refresh.
|
||||
async function getImageCount(folderPath, recursive) {
|
||||
const url = new URL(COUNT_ENDPOINT, window.location.origin);
|
||||
url.searchParams.set("folder_path", folderPath || "");
|
||||
url.searchParams.set("recursive", recursive ? "true" : "false");
|
||||
|
||||
try {
|
||||
const response = await fetch(url);
|
||||
@@ -38,8 +34,7 @@ async function getImageCount(folderPath) {
|
||||
|
||||
function setWidgetValue(node, widget, value) {
|
||||
widget.value = value;
|
||||
node.setDirtyCanvas?.(true, true);
|
||||
app.graph?.setDirtyCanvas?.(true, true);
|
||||
markDirty(node, app.graph);
|
||||
}
|
||||
|
||||
function installIndexControl(node) {
|
||||
@@ -51,6 +46,7 @@ function installIndexControl(node) {
|
||||
const indexWidget = findWidget(node, "index");
|
||||
const stepWidget = findWidget(node, "step");
|
||||
const modeWidget = findWidget(node, "seed_mode");
|
||||
const recursiveWidget = findWidget(node, "recursive"); // optional: a node without it counts as not recursive
|
||||
|
||||
if (!folderWidget || !indexWidget || !stepWidget || !modeWidget) {
|
||||
return;
|
||||
@@ -61,7 +57,7 @@ function installIndexControl(node) {
|
||||
indexWidget.serializeValue = async function () {
|
||||
const mode = modeWidget.value || "fixed";
|
||||
const step = Math.max(1, toInt(stepWidget.value, 1));
|
||||
const imageCount = await getImageCount(folderWidget.value);
|
||||
const imageCount = await getImageCount(folderWidget.value, Boolean(recursiveWidget?.value));
|
||||
let currentIndex = wrapIndex(toInt(indexWidget.value, 0), imageCount);
|
||||
|
||||
if (mode === "random") {
|
||||
|
||||
@@ -1,36 +1,8 @@
|
||||
import { app } from "../../../scripts/app.js";
|
||||
import { chainCallback, findWidget, hideWidget, refreshNode, showWidget } from "./nh_widget_utils.js";
|
||||
|
||||
const NODE_NAME = "NH_SaveImagePath";
|
||||
|
||||
function findWidget(node, name) {
|
||||
return node.widgets?.find((widget) => widget.name === name);
|
||||
}
|
||||
|
||||
function hideWidget(widget) {
|
||||
if (!widget || widget.__nhSaveImagePathHidden) {
|
||||
return;
|
||||
}
|
||||
|
||||
widget.__nhSaveImagePathHidden = true;
|
||||
widget.__nhSaveImagePathType = widget.type;
|
||||
widget.__nhSaveImagePathComputeSize = widget.computeSize;
|
||||
widget.computeSize = () => [0, -4];
|
||||
widget.type = "nh-hidden";
|
||||
}
|
||||
|
||||
function showWidget(widget) {
|
||||
if (!widget || !widget.__nhSaveImagePathHidden) {
|
||||
return;
|
||||
}
|
||||
|
||||
widget.type = widget.__nhSaveImagePathType;
|
||||
widget.computeSize = widget.__nhSaveImagePathComputeSize;
|
||||
|
||||
delete widget.__nhSaveImagePathHidden;
|
||||
delete widget.__nhSaveImagePathType;
|
||||
delete widget.__nhSaveImagePathComputeSize;
|
||||
}
|
||||
|
||||
function updateNewFolderWidget(node) {
|
||||
const createWidget = findWidget(node, "create_new_folder");
|
||||
const folderNameWidget = findWidget(node, "new_folder_name");
|
||||
@@ -45,9 +17,7 @@ function updateNewFolderWidget(node) {
|
||||
hideWidget(folderNameWidget);
|
||||
}
|
||||
|
||||
node.setSize?.(node.computeSize?.() || node.size);
|
||||
node.setDirtyCanvas?.(true, true);
|
||||
app.graph?.setDirtyCanvas?.(true, true);
|
||||
refreshNode(node, app.graph);
|
||||
}
|
||||
|
||||
function installSaveImagePathControls(node) {
|
||||
@@ -58,11 +28,7 @@ function installSaveImagePathControls(node) {
|
||||
|
||||
const createWidget = findWidget(node, "create_new_folder");
|
||||
if (createWidget) {
|
||||
const oldCallback = createWidget.callback;
|
||||
createWidget.callback = function () {
|
||||
oldCallback?.apply(this, arguments);
|
||||
updateNewFolderWidget(node);
|
||||
};
|
||||
chainCallback(createWidget, () => updateNewFolderWidget(node));
|
||||
}
|
||||
|
||||
setTimeout(() => updateNewFolderWidget(node), 0);
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import { app } from "../../../scripts/app.js";
|
||||
import { chainCallback, findWidget, markDirty } from "./nh_widget_utils.js";
|
||||
|
||||
const NODE_NAME = "NH_SmartResolutionPicker";
|
||||
const RATIO_RESIZE_NODE_NAME = "NH_RatioPresetImageResize";
|
||||
@@ -6,10 +7,6 @@ const PRESETS_ENDPOINT = "/nh-nodes/smart-resolution-picker/presets";
|
||||
|
||||
let cachedPresetData = null;
|
||||
|
||||
function findWidget(node, name) {
|
||||
return node.widgets?.find((widget) => widget.name === name);
|
||||
}
|
||||
|
||||
async function getPresetData() {
|
||||
if (cachedPresetData) {
|
||||
return cachedPresetData;
|
||||
@@ -45,8 +42,7 @@ function updatePresetDropdown(node, presetData, keepCurrent = true) {
|
||||
presetWidget.value = model === presetData.default_model ? presetData.default_preset : presets[0];
|
||||
}
|
||||
|
||||
node.setDirtyCanvas?.(true, true);
|
||||
app.graph?.setDirtyCanvas?.(true, true);
|
||||
markDirty(node, app.graph);
|
||||
}
|
||||
|
||||
function updateRatioDropdown(node, presetData, keepCurrent = true) {
|
||||
@@ -69,8 +65,7 @@ function updateRatioDropdown(node, presetData, keepCurrent = true) {
|
||||
ratioWidget.value = ratios.includes("1:1") ? "1:1" : ratios[0];
|
||||
}
|
||||
|
||||
node.setDirtyCanvas?.(true, true);
|
||||
app.graph?.setDirtyCanvas?.(true, true);
|
||||
markDirty(node, app.graph);
|
||||
}
|
||||
|
||||
function installSmartResolutionPicker(node) {
|
||||
@@ -87,11 +82,7 @@ function installSmartResolutionPicker(node) {
|
||||
return;
|
||||
}
|
||||
|
||||
const oldCallback = modelWidget.callback;
|
||||
modelWidget.callback = function () {
|
||||
oldCallback?.apply(this, arguments);
|
||||
updatePresetDropdown(node, presetData, false);
|
||||
};
|
||||
chainCallback(modelWidget, () => updatePresetDropdown(node, presetData, false));
|
||||
}).catch((error) => {
|
||||
console.warn("[NH-Nodes] Smart Resolution Picker preset update failed:", error);
|
||||
});
|
||||
@@ -112,11 +103,7 @@ function installRatioPresetImageResize(node) {
|
||||
continue;
|
||||
}
|
||||
|
||||
const oldCallback = widget.callback;
|
||||
widget.callback = function () {
|
||||
oldCallback?.apply(this, arguments);
|
||||
updateRatioDropdown(node, presetData, false);
|
||||
};
|
||||
chainCallback(widget, () => updateRatioDropdown(node, presetData, false));
|
||||
}
|
||||
}).catch((error) => {
|
||||
console.warn("[NH-Nodes] Ratio Preset Image Resize ratio update failed:", error);
|
||||
|
||||
@@ -1,44 +1,11 @@
|
||||
import { app } from "../../../scripts/app.js";
|
||||
import { chainCallback, clampInt, findWidget, hideWidget, refreshNode, showWidget } from "./nh_widget_utils.js";
|
||||
|
||||
const NODE_NAME = "NH_ValueMatchIndex";
|
||||
const MAX_VALUES = 64;
|
||||
|
||||
function findWidget(node, name) {
|
||||
return node.widgets?.find((widget) => widget.name === name);
|
||||
}
|
||||
|
||||
function toInt(value, fallback) {
|
||||
const parsed = Number.parseInt(value, 10);
|
||||
return Number.isFinite(parsed) ? parsed : fallback;
|
||||
}
|
||||
|
||||
function clampValueCount(value) {
|
||||
return Math.max(1, Math.min(MAX_VALUES, toInt(value, 5)));
|
||||
}
|
||||
|
||||
function hideWidget(widget) {
|
||||
if (!widget || widget.__nhValueMatchHidden) {
|
||||
return;
|
||||
}
|
||||
|
||||
widget.__nhValueMatchHidden = true;
|
||||
widget.__nhValueMatchType = widget.type;
|
||||
widget.__nhValueMatchComputeSize = widget.computeSize;
|
||||
widget.computeSize = () => [0, -4];
|
||||
widget.type = "nh-hidden";
|
||||
}
|
||||
|
||||
function showWidget(widget) {
|
||||
if (!widget || !widget.__nhValueMatchHidden) {
|
||||
return;
|
||||
}
|
||||
|
||||
widget.type = widget.__nhValueMatchType;
|
||||
widget.computeSize = widget.__nhValueMatchComputeSize;
|
||||
|
||||
delete widget.__nhValueMatchHidden;
|
||||
delete widget.__nhValueMatchType;
|
||||
delete widget.__nhValueMatchComputeSize;
|
||||
return clampInt(value, 1, MAX_VALUES, 5);
|
||||
}
|
||||
|
||||
function updateValueWidgets(node) {
|
||||
@@ -58,9 +25,7 @@ function updateValueWidgets(node) {
|
||||
}
|
||||
}
|
||||
|
||||
node.setSize?.(node.computeSize?.() || node.size);
|
||||
node.setDirtyCanvas?.(true, true);
|
||||
app.graph?.setDirtyCanvas?.(true, true);
|
||||
refreshNode(node, app.graph);
|
||||
}
|
||||
|
||||
function installValueMatchControls(node) {
|
||||
@@ -76,11 +41,7 @@ function installValueMatchControls(node) {
|
||||
|
||||
const countWidget = findWidget(node, "value_count");
|
||||
if (countWidget) {
|
||||
const oldCallback = countWidget.callback;
|
||||
countWidget.callback = function () {
|
||||
oldCallback?.apply(this, arguments);
|
||||
updateValueWidgets(node);
|
||||
};
|
||||
chainCallback(countWidget, () => updateValueWidgets(node));
|
||||
}
|
||||
|
||||
setTimeout(() => updateValueWidgets(node), 0);
|
||||
|
||||
@@ -0,0 +1,91 @@
|
||||
// Helpers shared by the NH-Nodes frontend extensions. This file registers no extension and imports nothing: ComfyUI
|
||||
// loads every .js file of the web folder, the extensions import this one with a relative path, and the browser keeps a
|
||||
// single copy of it.
|
||||
|
||||
export function findWidget(node, name) {
|
||||
return node?.widgets?.find((widget) => widget.name === name);
|
||||
}
|
||||
|
||||
export function toInt(value, fallback) {
|
||||
const parsed = Number.parseInt(value, 10);
|
||||
return Number.isFinite(parsed) ? parsed : fallback;
|
||||
}
|
||||
|
||||
// `value` as an integer between min and max; `fallback` (also clamped) when it is not a number.
|
||||
export function clampInt(value, min, max, fallback) {
|
||||
return Math.max(min, Math.min(max, toInt(value, fallback)));
|
||||
}
|
||||
|
||||
// Hiding a widget.
|
||||
//
|
||||
// ComfyUI frontend 1.55.17 (litegraph BaseWidget) has `widget.hidden`: a hidden widget is not drawn, takes no row, cannot
|
||||
// be clicked, and is still saved with the workflow and sent with the prompt. Older frontends have no such property, so
|
||||
// there a widget is hidden the way this pack always did it: a type no renderer knows and a computeSize that takes no
|
||||
// room. `serialize` is never touched, so the value of a hidden widget is saved either way.
|
||||
//
|
||||
// showWidget only undoes what hideWidget did: a widget the node definition hides stays hidden.
|
||||
const HIDDEN_STATE = "__nhHiddenState";
|
||||
const HIDDEN_TYPE = "nh-hidden";
|
||||
|
||||
export function hideWidget(widget) {
|
||||
if (!widget || widget[HIDDEN_STATE]) {
|
||||
return;
|
||||
}
|
||||
|
||||
if ("hidden" in widget) {
|
||||
widget[HIDDEN_STATE] = { native: true };
|
||||
widget.hidden = true;
|
||||
return;
|
||||
}
|
||||
|
||||
widget[HIDDEN_STATE] = {
|
||||
native: false,
|
||||
type: widget.type,
|
||||
computeSize: widget.computeSize,
|
||||
ownComputeSize: Object.prototype.hasOwnProperty.call(widget, "computeSize"),
|
||||
};
|
||||
widget.computeSize = () => [0, -4];
|
||||
widget.type = HIDDEN_TYPE;
|
||||
}
|
||||
|
||||
export function showWidget(widget) {
|
||||
const state = widget?.[HIDDEN_STATE];
|
||||
if (!state) {
|
||||
return;
|
||||
}
|
||||
delete widget[HIDDEN_STATE];
|
||||
|
||||
if (state.native) {
|
||||
widget.hidden = false;
|
||||
return;
|
||||
}
|
||||
|
||||
widget.type = state.type;
|
||||
if (state.ownComputeSize) {
|
||||
widget.computeSize = state.computeSize;
|
||||
} else {
|
||||
delete widget.computeSize;
|
||||
}
|
||||
}
|
||||
|
||||
// Calls `after` each time the widget's callback fires, once the callback that was already there (the frontend's, or
|
||||
// another extension's) has run. `after` gets the same `this` and arguments.
|
||||
export function chainCallback(widget, after) {
|
||||
const original = widget.callback;
|
||||
widget.callback = function () {
|
||||
original?.apply(this, arguments);
|
||||
after.apply(this, arguments);
|
||||
};
|
||||
}
|
||||
|
||||
// Repaints the canvas. `graph` is the application's graph (app.graph), which can differ from node.graph.
|
||||
export function markDirty(node, graph) {
|
||||
node.setDirtyCanvas?.(true, true);
|
||||
graph?.setDirtyCanvas?.(true, true);
|
||||
}
|
||||
|
||||
// Fits the node to the widgets and slots it has now, then repaints.
|
||||
export function refreshNode(node, graph) {
|
||||
node.setSize?.(node.computeSize?.() || node.size);
|
||||
markDirty(node, graph);
|
||||
}
|
||||
Reference in New Issue
Block a user