From 8ff49ec89588d61f8cc4678f15d69b9062f64c5b Mon Sep 17 00:00:00 2001 From: StableLlama Date: Wed, 9 Sep 2026 14:22:21 +0200 Subject: [PATCH] Enhance: - "save IMAGE to file" and "save IMAGE+MASK to file": `format` is now a dropdown (still connectable as a STRING) whose entries are derived from the formats the installed Pillow/image plugins can actually write (RGB node offers them all, IMAGE+MASK only alpha-capable ones); every batch frame is written instead of only the first; and `path` can be switched to ComfyUI-style `filename_prefix` mode via a toggle, resolved through ComfyUI's own folder/file helpers (auto-numbered under the output dir); both modes accept ComfyUI templates such as `%date:yyyy-MM-dd%` --- CHANGELOG.md | 8 + src/basic_data_handling/path_nodes.py | 536 ++++++++++++++++++++------ tests/test_path_nodes.py | 193 ++++++++++ 3 files changed, 624 insertions(+), 113 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 84a6a2d..9199ea7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Add "create DICT from JSON string" and "create LIST from JSON string" nodes to parse a multiline JSON object/array into a DICT/LIST +- "save IMAGE to file" and "save IMAGE+MASK to file": `format` is now a dropdown + (still connectable as a STRING) whose entries are derived from the formats the + installed Pillow/image plugins can actually write (RGB node offers them all, + IMAGE+MASK only alpha-capable ones); every batch frame is written instead of + only the first; and `path` can be switched to ComfyUI-style `filename_prefix` + mode via a toggle, resolved through ComfyUI's own folder/file helpers + (auto-numbered under the output dir); both modes accept ComfyUI templates such + as `%date:yyyy-MM-dd%` ## [1.8.1] - 2026-09-06 diff --git a/src/basic_data_handling/path_nodes.py b/src/basic_data_handling/path_nodes.py index 6badb60..92358bd 100644 --- a/src/basic_data_handling/path_nodes.py +++ b/src/basic_data_handling/path_nodes.py @@ -24,6 +24,16 @@ except: get_output_directory = get_input_directory +# ComfyUI's official prefix->path resolver is reused when running under ComfyUI +# so the nodes follow upstream naming exactly (folder layout, %width%/%height%/ +# date token expansion, auto-incrementing counters and containment checks). +# Outside ComfyUI (e.g. the standalone test-suite) `folder_paths` is not +# importable, so a minimal local equivalent is used as a fallback only. +try: + from folder_paths import get_save_image_path as _official_get_save_image_path +except Exception: + _official_get_save_image_path = None + def _require_numpy(): try: @@ -1148,6 +1158,339 @@ def metadata_save_kwargs(metadata_text: str, fmt: str) -> dict: return {} +# --- shared helpers for the image-save nodes -------------------------------- +# +# The save-format dropdown is not hard-coded: it is derived from the installed +# Pillow (+ any image plugins such as pillow-jxl it can save with), so it always +# matches what the running library supports and automatically picks up new +# formats/plugins. Only "formats" that are not meaningful photographic raster +# output (icons / GPU textures / vector / multi-image containers, e.g. DDS, +# ICNS, EPS, MPO) are never offered even when Pillow can technically write them. +# +# A save *token* is the short, user-facing format name and equals the file +# extension we write (jpg/tif/jp2/...). Each token maps to the Pillow format-id +# used with ``Image.save(format=...)``. +_PIL_FORMAT_BY_TOKEN: dict[str, str] = { + "png": "PNG", + "jpg": "JPEG", + "jpeg": "JPEG", # alias; normalised to "jpg" + "webp": "WEBP", + "jxl": "JXL", + "bmp": "BMP", + "tif": "TIFF", + "tiff": "TIFF", # alias; normalised to "tif" + "gif": "GIF", + "jp2": "JPEG2000", + "jpeg2000": "JPEG2000", # alias; normalised to "jp2" + "ppm": "PPM", + "pcx": "PCX", + "tga": "TGA", + "qoi": "QOI", + "avif": "AVIF", +} + +# friendly ordering for the dropdown: the well-known set first, then whatever +# else is discovered in the installed library/plugins. +_IMAGE_SAVE_ORDER = ["png", "jpg", "webp", "jxl", "bmp", "tif", "gif", "jp2", "ppm", "pcx", "tga", "qoi", "avif"] + +# Pillow format-ids that are not meaningful photographic raster *stills* for this +# save node (icons/cursors, GPU textures, vector/postscript, single-multi-image +# containers, scientific dumps). Never offered even if Pillow can write them. +_NON_PHOTOGRAPHIC_RASTER_IDS = { + "ICO", "CUR", "ICNS", "DDS", "DIB", "EPS", "IM", "MPO", "SGI", + "PDF", "WMF", "EMF", "BUFR", "GRIB", "HDF5", "FITS", +} + +# canonical token for a couple of common user-spelled aliases +_FORMAT_ALIASES = {"jpeg": "jpg", "tiff": "tif", "jpeg2000": "jp2"} + + +def _load_image_plugins() -> None: + """Import optional Pillow plugins (e.g. pillow-jxl) so Pillow knows about + every format the library can actually save.""" + try: + import pillow_jxl # noqa: F401 - imported but unused, registers the JXL plugin + except ModuleNotFoundError: + pass + + +def _writable_pillow_formats(mode: str) -> set[str]: + """ + The set of Pillow format-ids that can successfully encode a small image in + ``mode`` (probed at runtime against the actually installed library/plugins). + """ + import io as _io + from PIL import Image as _PIL + _load_image_plugins() + ids = set(_PIL.SAVE.keys()) + ids.update(str(v).upper() for v in _PIL.registered_extensions().values()) + writable: set[str] = set() + for fid in sorted(ids): + try: + buf = _io.BytesIO() + _PIL.new(mode, (4, 4)).save(buf, format=fid) + buf.seek(0) + _PIL.open(buf).load() + writable.add(str(fid).upper()) + except Exception: + continue + return writable + + +def _discover_save_formats() -> tuple[list[str], list[str]]: + """ + Derive ``(all_formats, alpha_formats)`` dropdown lists from Pillow/plugins. + ``alpha_formats`` (surfaced on the IMAGE+MASK node) is the subset that can + also hold an alpha channel. + """ + writable_rgb = _writable_pillow_formats("RGB") + writable_rgba = _writable_pillow_formats("RGBA") + + # 1) tokens we already know about and that are actually writable + offered = [ + t for t in _IMAGE_SAVE_ORDER + if _PIL_FORMAT_BY_TOKEN[t] in writable_rgb + ] + known_ids = {_PIL_FORMAT_BY_TOKEN[t] for t in offered} + + # 2) any additional writable photographic format in the library/plugins + extras = [] + for pid in sorted(writable_rgb): + if pid in known_ids or pid in _NON_PHOTOGRAPHIC_RASTER_IDS: + continue + token = pid.lower() + # register it so it can be written through the same machinery + _PIL_FORMAT_BY_TOKEN[token] = pid + if token not in _IMAGE_SAVE_ORDER: + _IMAGE_SAVE_ORDER.append(token) + extras.append(token) + + all_formats = offered + extras + alpha_formats = [t for t in all_formats if _PIL_FORMAT_BY_TOKEN[t] in writable_rgba] + return all_formats, alpha_formats + + +(_IMAGE_SAVE_FORMATS, _IMAGE_SAVE_ALPHA_FORMATS) = _discover_save_formats() + +# Always make sure the most common formats are present when the backend supports +# them (guards against an odd Pillow not probing its own core formats). +for _token in ("png", "jpg", "webp"): + if _token in _PIL_FORMAT_BY_TOKEN and _PIL_FORMAT_BY_TOKEN[_token] and _token not in _IMAGE_SAVE_FORMATS: + _IMAGE_SAVE_FORMATS.insert(0, _token) + +# Log the supported save formats once at node-load time, so users can see what +# the running Pillow/image plugins can write. +print("Basic data handling: supported image formats: " + ", ".join(_IMAGE_SAVE_FORMATS)) +print("Basic data handling: supported IMAGE+MASK (alpha) formats: " + ", ".join(_IMAGE_SAVE_ALPHA_FORMATS)) + + +def _normalize_image_format(format: str) -> str: + """ + Normalise a user-supplied format to its canonical save token (lower-cased), + mapping common aliases (jpeg->jpg, tiff->tif, jpeg2000->jp2) so both the + dropdown spellings and values wired in from elsewhere are accepted. Unknown + tokens are passed through unchanged (they may name a format the writer will + then attempt, or reject gracefully). + """ + token = str(format or "png").strip().lower() + return _FORMAT_ALIASES.get(token, token) + + +def _comfy_format_date(fmt_text: str, when) -> str: + """ + Format ``when`` using ComfyUI's own ``formatDate`` semantics (used by + ``%date:%``). Supported tokens: ``yyyy``/``yy`` for the year, and + ``M``/``MM``, ``d``/``dd``, ``h``/``hh``, ``m``/``mm``, ``s``/``ss`` + (lower-case ``h`` for hours, exactly like ComfyUI). Single letters are + unpadded, doubled letters are zero-padded. + """ + import re + parts = { + "d": when.day, + "M": when.month, + "h": when.hour, + "m": when.minute, + "s": when.second, + } + # Mirrors ComfyUI's regex: dd?|MM?|hh?|mm?|ss?|yyy?y? (yyy?y? last so a + # four-digit year is matched as a single token, not two "yy" groups). + pattern = re.compile(r"dd?|MM?|hh?|mm?|ss?|yyy?y?") + + def replace(token: str) -> str: + if token == "yy": + return str(when.year)[-2:] + if token == "yyyy": + return str(when.year) + if token and token[0] in parts: + # padStart(token.length, "0") + return str(parts[token[0]]).zfill(len(token)) + return token + + return pattern.sub(lambda m: replace(m.group(0)), fmt_text) + + +def _expand_filename_tokens(text: str, width: int = 0, height: int = 0) -> str: + """ + Expand ComfyUI-style filename templates in ``text`` for both the plain path + and the ``filename_prefix`` modes. + + Handles ``%date:%`` (using ComfyUI's own date formatting) as well as + the standard built-in tokens ``%width%``, ``%height%``, ``%year%``, + ``%month%``, ``%day%``, ``%hour%``, ``%minute%`` and ``%second%``. Node + reference tokens (``%Node.widget%``) can only be resolved by the ComfyUI + frontend, so they are left untouched here. + """ + import re + from datetime import datetime + + if not text: + return text + + now = datetime.now() + text = re.sub(r"%date:(.*?)%", lambda m: _comfy_format_date(m.group(1), now), text) + text = text.replace("%width%", str(width)) + text = text.replace("%height%", str(height)) + text = text.replace("%year%", str(now.year)) + text = text.replace("%month%", str(now.month).zfill(2)) + text = text.replace("%day%", str(now.day).zfill(2)) + text = text.replace("%hour%", str(now.hour).zfill(2)) + text = text.replace("%minute%", str(now.minute).zfill(2)) + text = text.replace("%second%", str(now.second).zfill(2)) + return text + + + +def _official_save_image_path(filename_prefix, output_dir, image_width=0, image_height=0): + """ + Resolve ``filename_prefix`` to concrete output folder/filename/counter using + ComfyUI's own ``folder_paths.get_save_image_path`` wherever possible. + + Returns ``(full_output_folder, filename, counter, subfolder, resolved_prefix)`` + exactly like the official function. When ``folder_paths`` is unavailable + (standalone test-suite / non-ComfyUI host), a minimal local equivalent that + mirrors the same subfolder + ``__`` scheme is used instead. + """ + if _official_get_save_image_path is not None: + return _official_get_save_image_path(filename_prefix, output_dir, image_width, image_height) + + # --- fallback (only reached outside ComfyUI) ------------------------- + def expand(text, w, h): + import time + now = time.localtime() + text = text.replace("%width%", str(w)).replace("%height%", str(h)) + text = text.replace("%year%", str(now.tm_year)) + text = text.replace("%month%", str(now.tm_mon).zfill(2)) + text = text.replace("%day%", str(now.tm_mday).zfill(2)) + text = text.replace("%hour%", str(now.tm_hour).zfill(2)) + text = text.replace("%minute%", str(now.tm_min).zfill(2)) + text = text.replace("%second%", str(now.tm_sec).zfill(2)) + return text + + if "%" in filename_prefix: + filename_prefix = expand(filename_prefix, image_width, image_height) + subfolder = os.path.dirname(os.path.normpath(filename_prefix)) + filename = os.path.basename(os.path.normpath(filename_prefix)) + full_output_folder = os.path.join(output_dir, subfolder) + + highest = 0 + try: + for entry in os.listdir(full_output_folder): + stem, _, _ = entry.rpartition(".") + if stem.startswith(filename + "_") and stem[len(filename) + 1:].rstrip("_").isdigit(): + digits = stem[len(filename) + 1:].rstrip("_") + if digits.isdigit(): + highest = max(highest, int(digits)) + except FileNotFoundError: + os.makedirs(full_output_folder, exist_ok=True) + return full_output_folder, filename, highest + 1, subfolder, filename_prefix + + +def _plan_save_paths(path: str, format: str, use_prefix_mode: bool, + frame_count: int, width: int, height: int) -> tuple[list[str], bool]: + """ + Compute the absolute destination path(s) for ``frame_count`` images. + + Returns ``(paths, create_dirs)``: + - ``use_prefix_mode`` False: ``path`` is a concrete file location. A single + frame writes exactly to ``path`` (+ extension); several frames get an + incrementing ``_00000``-style suffix so they do not overwrite each other. + - ``use_prefix_mode`` True: ``path`` is a ComfyUI ``filename_prefix`` under + the output folder, resolved through ComfyUI's own filename helpers and + auto-numbered like the built-in "Save Image" node. + """ + fmt = _normalize_image_format(format) + # Expand ComfyUI-style templates (%date:...%, %width%, %height%, ...) up + # front so they apply in both the plain-path and the prefix modes. + path = _expand_filename_tokens(path, width, height) + + if use_prefix_mode: + full_output_folder, filename, counter, _, _ = _official_save_image_path( + path, get_output_directory(), width, height + ) + paths = [] + for _ in range(frame_count): + name = f"{filename}_{counter:05}_.{fmt}" + paths.append(os.path.join(full_output_folder, name)) + counter += 1 + return paths, False + + if not path.endswith(f".{fmt}"): + base = f"{path}.{fmt}" + else: + base = path + + if frame_count <= 1: + return [base], True + + root, ext = os.path.splitext(base) + paths = [f"{root}_{i:05}{ext}" for i in range(frame_count)] + return paths, True + + +def _ensure_parent_directories(path: str, create_dirs: bool) -> None: + """Create the parent directory(ies) of ``path`` when requested.""" + directory = os.path.dirname(path) + if directory and create_dirs and not os.path.exists(directory): + os.makedirs(directory) + + +def _has_jxl_support() -> bool: + """Return True when the pillow-jxl plugin is importable.""" + try: + import pillow_jxl # noqa: F401 - imported but unused, kept for JPEG XL support + return True + except ModuleNotFoundError: + return False + + +def _write_pil_image(pil_img, path: str, fmt: str, quality: int, metadata_text: str) -> bool: + """ + Write ``pil_img`` to ``path``. ``fmt`` is a canonical save token (see + ``_PIL_FORMAT_BY_TOKEN``). For the lossy/metadata-capable set (png/jpg/webp/ + jxl) the prompt text is embedded as ``parameters`` metadata and ``quality`` + is honoured; all other discovered formats are written plainly. Returns True + on success. + """ + try: + if fmt == "png": + pil_img.save(path, format="PNG", **metadata_save_kwargs(metadata_text, fmt)) + elif fmt == "jpg": + # JPEG cannot carry alpha; drop it so an RGBA input still saves. + pil_img.convert("RGB").save(path, format="JPEG", quality=quality, **metadata_save_kwargs(metadata_text, fmt)) + elif fmt == "webp": + pil_img.save(path, format="WEBP", quality=quality, **metadata_save_kwargs(metadata_text, fmt)) + elif fmt == "jxl": + pil_img.save(path, format="JXL", quality=quality, **metadata_save_kwargs(metadata_text, fmt)) + else: + if metadata_text: + print("Basic data handling: Prompt metadata is not supported for this format; skipping it.") + pil_img.save(path, format=_PIL_FORMAT_BY_TOKEN.get(fmt, fmt.upper())) + return True + except Exception as e: + print(f"Basic data handling: Error saving image: {e}") + return False + + class PathSaveImageRGB(ComfyNodeABC): """ Saves an image to a file. @@ -1155,6 +1498,12 @@ class PathSaveImageRGB(ComfyNodeABC): This node takes an image tensor and saves it to the specified path. Supports various image formats like PNG, JPG, WEBP, JXL (if pillow-jxl is installed), etc. + By default ``path`` is a concrete absolute/relative file location and every + image frame of the batch is written (several frames get an incrementing + suffix). When ``use_prefix_mode`` is enabled ``path`` is instead treated as + a ComfyUI ``filename_prefix`` under the output folder, named and + auto-numbered exactly like the built-in "Save Image" node. + When ``prompt`` and/or ``negative_prompt`` are provided, they are embedded into the saved image as ``parameters`` metadata: in the PNG text chunk, in the EXIF ``UserComment`` (and ``ImageDescription`` for JPEG) fields, and in @@ -1166,14 +1515,18 @@ class PathSaveImageRGB(ComfyNodeABC): return { "required": { "images": (IO.IMAGE,), - "path": (IO.STRING, {"default": "", "tooltip": "Destination file path (an extension is added from the format when missing)."}), + "path": (IO.STRING, {"default": "", "tooltip": "Destination file path (an extension is added from the format when missing), or a ComfyUI filename_prefix under the output folder when \"use prefix mode\" is enabled."}), }, "optional": { - "format": (IO.STRING, {"default": "png", "tooltip": "Image format: png, jpg, webp or jxl (jxl needs pillow-jxl installed)."}), + "format": (_IMAGE_SAVE_FORMATS, {"default": "png", "tooltip": "Image save format (png, jpg, webp, bmp, ...). The list is derived from what the installed Pillow/plugins can write. Drag a STRING onto this to override it."}), "quality": (IO.INT, {"default": 95, "min": 1, "max": 100, "tooltip": "Quality for lossy formats (jpg/webp/jxl)."}), - "create_dirs": (IO.BOOLEAN, {"default": True, "tooltip": "Create missing parent directories."}), + "create_dirs": (IO.BOOLEAN, {"default": True, "tooltip": "Create missing parent directories (plain path mode only)."}), "prompt": (IO.STRING, {"default": "", "tooltip": "Optional positive prompt embedded as parameters metadata."}), "negative_prompt": (IO.STRING, {"default": "", "tooltip": "Optional negative prompt embedded as parameters metadata."}), + # Appended last on purpose: the node stores widget values + # positionally, so adding before the existing widgets would shift + # old workflows and mis-assign their saved values. + "use_prefix_mode": (IO.BOOLEAN, {"default": False, "tooltip": "When True, 'path' is treated as a ComfyUI filename_prefix under the output folder and files are named/auto-numbered exactly like ComfyUI's \"Save Image\" node."}), } } @@ -1185,72 +1538,42 @@ class PathSaveImageRGB(ComfyNodeABC): FUNCTION = "save_image" OUTPUT_NODE = True - def save_image(self, images, path: str, format: str = "png", quality: int = 95, - create_dirs: bool = True, prompt: str = "", negative_prompt: str = ""): + def save_image(self, images, path: str = "", use_prefix_mode: bool = False, format: str = "png", + quality: int = 95, create_dirs: bool = True, + prompt: str = "", negative_prompt: str = ""): if not path: print("Basic data handling: Save failed - no path specified") return (False,) - # If the path doesn't have an extension or it doesn't match the format, add it - if not path.lower().endswith(f".{format.lower()}"): - path = f"{path}.{format.lower()}" + import numpy as np + from PIL import Image - try: - import numpy as np - from PIL import Image + fmt = _normalize_image_format(format) + if fmt == "jxl" and not _has_jxl_support(): + print("Basic data handling: JPEG XL format requested but pillow_jxl module is not installed. " + "Please install it with 'pip install pillow-jxl-plugin'.") + return (False,) - # Check if pillow_jxl is available for JXL support - has_jxl_support = False - try: - import pillow_jxl # noqa: F401 - imported but unused, kept for JPEG XL support - has_jxl_support = True - except ModuleNotFoundError: - # pillow_jxl is not installed - if format.lower() == "jxl": - print("Basic data handling: JPEG XL format requested but pillow_jxl module is not installed. " - "Please install it with 'pip install pillow-jxl-plugin'.") - return (False,) + batch = len(images) + height, width = images.shape[1], images.shape[2] + paths, create_dirs = _plan_save_paths(path, fmt, use_prefix_mode, batch, width, height) - # Create directories if needed - directory = os.path.dirname(path) - if directory and create_dirs and not os.path.exists(directory): - os.makedirs(directory) + # Compose the prompt metadata to embed into the saved file once. + metadata_text = compose_prompt_text(prompt, negative_prompt) + + for index, target_path in enumerate(paths): + _ensure_parent_directories(target_path, create_dirs) # Convert from tensor format back to PIL Image - # Extract the first image from the batch - i = 0 - img_tensor = images[i].cpu().numpy() - - # Convert to uint8 format for PIL + img_tensor = images[index].cpu().numpy() img_np = (img_tensor * 255).astype(np.uint8) - - # Create PIL image pil_img = Image.fromarray(img_np) - # Compose the prompt metadata to embed into the saved file - metadata_text = compose_prompt_text(prompt, negative_prompt) - fmt = format.lower() + if not _write_pil_image(pil_img, target_path, fmt, quality, metadata_text): + return (False,) + print(f"Basic data handling: Successfully saved image to {target_path}") - # Save the image, embedding prompt metadata where the format supports it - if fmt == "jpg" or fmt == "jpeg": - pil_img.save(path, format="JPEG", quality=quality, **metadata_save_kwargs(metadata_text, fmt)) - elif fmt == "webp": - pil_img.save(path, format="WEBP", quality=quality, **metadata_save_kwargs(metadata_text, fmt)) - elif fmt == "jxl" and has_jxl_support: - # JPEG XL specific options - pil_img.save(path, format="JXL", quality=quality, **metadata_save_kwargs(metadata_text, fmt)) - elif fmt == "png": - pil_img.save(path, format="PNG", **metadata_save_kwargs(metadata_text, fmt)) - else: - if metadata_text: - print("Basic data handling: Prompt metadata is not supported for this format; skipping it.") - pil_img.save(path, format=format.upper()) - - print(f"Basic data handling: Successfully saved image to {path}") - return (True,) - except Exception as e: - print(f"Basic data handling: Error saving image: {e}") - return (False,) + return (True,) class PathSaveImageRGBA(ComfyNodeABC): @@ -1261,6 +1584,12 @@ class PathSaveImageRGBA(ComfyNodeABC): specified path as an image with transparency, where the mask defines the alpha channel. + By default ``path`` is a concrete absolute/relative file location and every + image frame of the batch is written (several frames get an incrementing + suffix). When ``use_prefix_mode`` is enabled ``path`` is instead treated as + a ComfyUI ``filename_prefix`` under the output folder, named and + auto-numbered exactly like the built-in "Save Image" node. + When ``prompt`` and/or ``negative_prompt`` are provided, they are embedded into the saved image as ``parameters`` metadata: in the PNG text chunk, in the EXIF ``UserComment`` (and ``ImageDescription`` for JPEG) fields, and in @@ -1273,15 +1602,19 @@ class PathSaveImageRGBA(ComfyNodeABC): "required": { "images": (IO.IMAGE,), "mask": (IO.MASK,), - "path": (IO.STRING, {"default": "", "tooltip": "Destination file path (an extension is added from the format when missing)."}), + "path": (IO.STRING, {"default": "", "tooltip": "Destination file path (an extension is added from the format when missing), or a ComfyUI filename_prefix under the output folder when \"use prefix mode\" is enabled."}), }, "optional": { - "format": (IO.STRING, {"default": "png", "tooltip": "Image format supporting alpha: png, webp or jxl; jpg is coerced to png."}), + "format": (_IMAGE_SAVE_ALPHA_FORMATS, {"default": "png", "tooltip": "Image save format that can hold an alpha channel (png, webp, ...; jpg is coerced to png). Derived from what the installed Pillow/plugins can write. Drag a STRING onto this to override it."}), "quality": (IO.INT, {"default": 95, "min": 1, "max": 100, "tooltip": "Quality for lossy formats (webp/jxl)."}), "invert_mask": (IO.BOOLEAN, {"default": False, "tooltip": "Invert the mask before using it as the alpha channel."}), - "create_dirs": (IO.BOOLEAN, {"default": True, "tooltip": "Create missing parent directories."}), + "create_dirs": (IO.BOOLEAN, {"default": True, "tooltip": "Create missing parent directories (plain path mode only)."}), "prompt": (IO.STRING, {"default": "", "tooltip": "Optional positive prompt embedded as parameters metadata."}), "negative_prompt": (IO.STRING, {"default": "", "tooltip": "Optional negative prompt embedded as parameters metadata."}), + # Appended last on purpose: the node stores widget values + # positionally, so adding before the existing widgets would shift + # old workflows and mis-assign their saved values. + "use_prefix_mode": (IO.BOOLEAN, {"default": False, "tooltip": "When True, 'path' is treated as a ComfyUI filename_prefix under the output folder and files are named/auto-numbered exactly like ComfyUI's \"Save Image\" node."}), } } @@ -1293,49 +1626,43 @@ class PathSaveImageRGBA(ComfyNodeABC): FUNCTION = "save_image_with_mask" OUTPUT_NODE = True - def save_image_with_mask(self, images, mask, path: str, format: str = "png", - quality: int = 95, invert_mask: bool = False, - create_dirs: bool = True, prompt: str = "", - negative_prompt: str = ""): + def save_image_with_mask(self, images, mask, path: str = "", use_prefix_mode: bool = False, + format: str = "png", quality: int = 95, + invert_mask: bool = False, create_dirs: bool = True, + prompt: str = "", negative_prompt: str = ""): if not path: print("Basic data handling: Save failed - no path specified") return (False,) - # Check format compatibility - needs to support alpha channel - if format.lower() in ["jpg", "jpeg"]: + import numpy as np + from PIL import Image + + # JPEG doesn't support transparency -> coerce to PNG, as before. + fmt = _normalize_image_format(format) + if fmt in ("jpg", "jpeg"): print("Basic data handling: JPEG format doesn't support transparency. Using PNG instead.") - format = "png" + fmt = "png" + if fmt == "jxl" and not _has_jxl_support(): + print("Basic data handling: JPEG XL format requested but pillow_jxl module is not installed. " + "Please install it with 'pip install pillow-jxl-plugin'.") + return (False,) - # If the path doesn't have an extension or it doesn't match the format, add it - if not path.lower().endswith(f".{format.lower()}"): - path = f"{path}.{format.lower()}" + batch = len(images) + height, width = images.shape[1], images.shape[2] + paths, create_dirs = _plan_save_paths(path, fmt, use_prefix_mode, batch, width, height) - try: - import numpy as np - from PIL import Image + mask_batch = len(mask) + # Mask tensor may be padded to the image batches; round-robin when the + # mask has fewer frames (e.g. a single mask applied to every frame). + mask_frames = mask.cpu() - # Check if pillow_jxl is available for JXL support - has_jxl_support = False - try: - import pillow_jxl # noqa: F401 - imported but unused, kept for JPEG XL support - has_jxl_support = True - except ModuleNotFoundError: - # pillow_jxl is not installed - if format.lower() == "jxl": - print("Basic data handling: JPEG XL format requested but pillow_jxl module is not installed. " - "Please install it with 'pip install pillow-jxl-plugin'.") - return (False,) + metadata_text = compose_prompt_text(prompt, negative_prompt) - # Create directories if needed - directory = os.path.dirname(path) - if directory and create_dirs and not os.path.exists(directory): - os.makedirs(directory) + for index, target_path in enumerate(paths): + _ensure_parent_directories(target_path, create_dirs) - # Convert from tensor format back to PIL Image - # Extract the first image from the batch - i = 0 - img_tensor = images[i].cpu().numpy() - mask_tensor = mask[i].cpu() + img_tensor = images[index].cpu().numpy() + mask_tensor = mask_frames[index % mask_batch] # Invert the mask if needed (1.0 becomes transparent, 0.0 becomes opaque) if invert_mask: @@ -1357,28 +1684,11 @@ class PathSaveImageRGBA(ComfyNodeABC): pil_img_rgba = pil_img.convert("RGBA") pil_img_rgba.putalpha(alpha_img) - # Compose the prompt metadata to embed into the saved file - metadata_text = compose_prompt_text(prompt, negative_prompt) - fmt = format.lower() + if not _write_pil_image(pil_img_rgba, target_path, fmt, quality, metadata_text): + return (False,) + print(f"Basic data handling: Successfully saved image with mask to {target_path}") - # Save the image, embedding prompt metadata where the format supports it - if fmt == "webp": - pil_img_rgba.save(path, format="WEBP", quality=quality, **metadata_save_kwargs(metadata_text, fmt)) - elif fmt == "jxl" and has_jxl_support: - # JPEG XL supports alpha channel - pil_img_rgba.save(path, format="JXL", quality=quality, **metadata_save_kwargs(metadata_text, fmt)) - elif fmt == "png": - pil_img_rgba.save(path, format="PNG", **metadata_save_kwargs(metadata_text, fmt)) - else: - if metadata_text: - print("Basic data handling: Prompt metadata is not supported for this format; skipping it.") - pil_img_rgba.save(path, format=format.upper()) - - print(f"Basic data handling: Successfully saved image with mask to {path}") - return (True,) - except Exception as e: - print(f"Basic data handling: Error saving image with mask: {e}") - return (False,) + return (True,) class PathInputDir(ComfyNodeABC): diff --git a/tests/test_path_nodes.py b/tests/test_path_nodes.py index 3ed6fb6..7edebeb 100644 --- a/tests/test_path_nodes.py +++ b/tests/test_path_nodes.py @@ -589,6 +589,199 @@ def test_path_save_image_jxl_prompt_metadata(tmp_path): assert escaped.encode("utf-8") in raw +def _batch_rgb(n, size=(16, 16)): + """Build an n-frame RGB tensor, frame i filled with a distinct colour.""" + frames = torch.zeros(n, size[1], size[0], 3) + for i in range(n): + frames[i, :, :, i % 3] = (i + 1) / float(n) + return frames + + +def test_path_save_image_format_is_dropdown(): + """The format input is a combo/dropdown whose entries are derived from what + the installed Pillow/plugins can actually save (not a hard-coded constant).""" + import src.basic_data_handling.path_nodes as pn + + # Must include the classic formats and be reflect the live library/plugins. + assert pn._IMAGE_SAVE_FORMATS + assert pn._IMAGE_SAVE_FORMATS[0] == "png" + + # png & alpha-capable core are always surfaced on the IMAGE+MASK node. + assert pn._IMAGE_SAVE_ALPHA_FORMATS + assert pn._IMAGE_SAVE_ALPHA_FORMATS[0] == "png" + assert "png" in pn._IMAGE_SAVE_ALPHA_FORMATS + + rgb_types = PathSaveImageRGB.INPUT_TYPES() + formats = rgb_types["optional"]["format"] + assert isinstance(formats[0], list) and formats[0] == pn._IMAGE_SAVE_FORMATS + assert formats[1]["default"] == "png" + + rgba_types = PathSaveImageRGBA.INPUT_TYPES() + formats_rgba = rgba_types["optional"]["format"] + # the IMAGE+MASK node offers only alpha-capable formats (jpg is absent, png + # present) but its runtime value is still a STRING-typed text. + assert isinstance(formats_rgba[0], list) + assert formats_rgba[0] == pn._IMAGE_SAVE_ALPHA_FORMATS + assert "jpg" not in formats_rgba[0] + assert formats_rgba[1]["default"] == "png" + + # Downstream aliases still accepted: jpg -> canonical jpg token. + assert pn._normalize_image_format("jpeg") == "jpg" + + # The format is still text at runtime, so it can be driven by a STRING link. + for f in formats[0]: + assert isinstance(f, str) + + # Both nodes expose a use_prefix_mode toggle. + assert rgba_types["optional"]["use_prefix_mode"][1]["default"] is False + + +def test_path_save_image_formats_are_library_derived(): + """The dropdown is derived from the live Pillow/plugin library, so it must + contain only formats that really save, exclude non-photographic ones, and + grow/shrink with the environment.""" + import src.basic_data_handling.path_nodes as pn + + # Non-photographic raster ids are never surfaced. + assert not any( + t in pn._NON_PHOTOGRAPHIC_RASTER_IDS for t in map(str.upper, pn._IMAGE_SAVE_FORMATS) + ) + + save_node = PathSaveImageRGB() + img = _batch_rgb(1) + for fmt in pn._IMAGE_SAVE_FORMATS: + out = f"/tmp/_bdh_fmt_{fmt}" + # RGB node writes any offered format + assert save_node.save_image(img, out, format=fmt) == (True,), fmt + # correct file extension appended + expected = f"{out}.{fmt}" + assert os.path.exists(expected), fmt + os.remove(expected) + + # Some classic formats are always present when Pillow can write them. + for fmt in ("png", "jpg", "webp"): + assert fmt in pn._IMAGE_SAVE_FORMATS + + +def test_path_save_image_rgb_multi_frame_suffix(tmp_path): + save_node = PathSaveImageRGB() + frames = _batch_rgb(3) + out = str(tmp_path / "frames") + + assert save_node.save_image(frames, out) == (True,) + expected = [ + str(tmp_path / "frames_00000.png"), + str(tmp_path / "frames_00001.png"), + str(tmp_path / "frames_00002.png"), + ] + for path in expected: + assert os.path.exists(path) + + +def test_path_save_image_prefix_mode_mirrors_comfy(tmp_path, monkeypatch): + """use_prefix_mode writes into the output dir with ComfyUI-style + ``prefix_00001_.png`` naming and auto-increments on repeated saves.""" + monkeypatch.setattr("src.basic_data_handling.path_nodes.get_output_directory", lambda: str(tmp_path)) + save_node = PathSaveImageRGB() + img = _batch_rgb(1) + prefix = "my_subdir/shot" + + assert save_node.save_image(img, prefix, use_prefix_mode=True) == (True,) + first = tmp_path / "my_subdir" / "shot_00001_.png" + assert first.exists() + + # A second save must not overwrite the first file. + assert save_node.save_image(img, prefix, use_prefix_mode=True) == (True,) + assert (tmp_path / "my_subdir" / "shot_00002_.png").exists() + + +def test_path_save_image_prefix_mode_batch(tmp_path, monkeypatch): + monkeypatch.setattr("src.basic_data_handling.path_nodes.get_output_directory", lambda: str(tmp_path)) + save_node = PathSaveImageRGB() + frames = _batch_rgb(2) + + assert save_node.save_image(frames, "batch", use_prefix_mode=True) == (True,) + assert (tmp_path / "batch_00001_.png").exists() + assert (tmp_path / "batch_00002_.png").exists() + + +def test_path_save_image_prefix_mode_template_tokens(tmp_path, monkeypatch): + """Int/width/height tokens are substituted like ComfyUI's backend before use.""" + monkeypatch.setattr("src.basic_data_handling.path_nodes.get_output_directory", lambda: str(tmp_path)) + save_node = PathSaveImageRGB() + img = _batch_rgb(1, size=(8, 4)) # width=8, height=4 + + assert save_node.save_image(img, "w%width%_h%height%", use_prefix_mode=True) == (True,) + assert (tmp_path / "w8_h4_00001_.png").exists() + + +def test_path_save_image_prefix_mode_date_token(tmp_path, monkeypatch): + """``%date:yyyy-MM-dd%`` is expanded to the real date folder (not a literal + ``%date:...%`` directory) in prefix mode.""" + from datetime import datetime + monkeypatch.setattr("src.basic_data_handling.path_nodes.get_output_directory", lambda: str(tmp_path)) + save_node = PathSaveImageRGB() + img = _batch_rgb(1) + + assert save_node.save_image(img, "%date:yyyy-MM-dd%/dated_test", use_prefix_mode=True) == (True,) + # The literal token must never be used as a folder name. + assert not (tmp_path / "%date:yyyy-MM-dd%").exists() + # A folder named after today's date holds the numbered file. + today = datetime.now().strftime("%Y-%m-%d") + assert (tmp_path / today / "dated_test_00001_.png").exists() + + +def test_path_save_image_path_mode_date_token(tmp_path): + """``%date:...%`` templates also work in the plain path mode.""" + from datetime import datetime + save_node = PathSaveImageRGB() + img = _batch_rgb(1) + today = datetime.now().strftime("%Y-%m-%d") + out = os.path.join(str(tmp_path), "%date:yyyy-MM-dd%", "dated") + + assert save_node.save_image(img, out) == (True,) + assert not (tmp_path / "%date:yyyy-MM-dd%").exists() + assert (tmp_path / today / "dated.png").exists() + + +def test_path_save_image_empty_path_returns_false(): + save_node = PathSaveImageRGB() + img = _batch_rgb(1) + assert save_node.save_image(img, "") == (False,) + + +def test_path_save_image_rgba_prefix_mode_and_transparency(tmp_path, monkeypatch): + monkeypatch.setattr("src.basic_data_handling.path_nodes.get_output_directory", lambda: str(tmp_path)) + save_node = PathSaveImageRGBA() + size = (8, 8) + img = torch.zeros(1, size[1], size[0], 3) + img[0, :, :, 0] = 1.0 + mask = torch.zeros(1, size[1], size[0]) + + assert save_node.save_image_with_mask(img, mask, "alpha", use_prefix_mode=True) == (True,) + path = tmp_path / "alpha_00001_.png" + assert path.exists() + + # The right half is masked (alpha 255/opaque where mask==0); check alpha via RGBA. + with Image.open(path) as loaded: + loaded.load() + # PNG metadata must be empty (no prompt given). + assert loaded.text.get("parameters") is None + + +def test_path_save_image_rgba_multi_frame(tmp_path): + save_node = PathSaveImageRGBA() + size = (8, 8) + frames = torch.zeros(2, size[1], size[0], 3) + frames[0, :, :, 0] = 1.0 + frames[1, :, :, 1] = 1.0 + mask = torch.zeros(1, size[1], size[0]) # single mask shared across frames + + assert save_node.save_image_with_mask(frames, mask, str(tmp_path / "rgba_batch")) == (True,) + assert (tmp_path / "rgba_batch_00000.png").exists() + assert (tmp_path / "rgba_batch_00001.png").exists() + + def test_path_normalize(): node = PathNormalize() assert node.normalize_path("folder/../file.txt") == (os.path.normpath("folder/../file.txt"),)