nodes/StableLlama-ComfyUI-basic_data_handling
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Enhance:

Browse Source 
- "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%`
This commit is contained in:
StableLlama
2026-09-09 14:22:21 +02:00
parent b46f4ba2bf
commit 8ff49ec895
3 changed files with 624 additions and 113 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CHANGELOG.md
+8
View File
@@ -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
src/basic_data_handling/path_nodes.py
+423 -113
View File
@@ -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:<format>%``). 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:<format>%`` (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 + ``_<counter>_`` 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):
tests/test_path_nodes.py
+193
View File
@@ -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"),)