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

Save workflow in the image and show preview

Browse Source
This commit is contained in:
StableLlama
2026-09-18 22:06:43 +02:00
parent 7a9e7c61ea
commit 1cdbd3044c
2 changed files with 487 additions and 94 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/basic_data_handling/path_nodes.py
+267 -68
View File
@@ -1,4 +1,6 @@
from inspect import cleandoc
from typing import Any, Optional
import json
import os
import glob
@@ -19,13 +21,14 @@ except:
ComfyNodeABC = object
try:
from folder_paths import get_input_directory, get_output_directory
from folder_paths import get_input_directory, get_output_directory, get_temp_directory
except:
def get_input_directory():
def get_input_directory() -> str:
return "./"
get_output_directory = get_input_directory
get_temp_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%/
@@ -1103,41 +1106,116 @@ def compose_prompt_text(prompt: str, negative_prompt: str) -> str:
return "\n".join(lines)
def build_png_info(metadata_text: str):
def _metadata_disabled() -> bool:
"""
Wrap ``metadata_text`` in a Pillow ``PngInfo`` container under the standard
``parameters`` text-chunk key so it can be embedded in a PNG file.
True when ComfyUI was started with ``--disable-metadata``.
Returns ``None`` when there is no text to embed (or Pillow's PNG metadata
support is unavailable), in which case the image should be saved without
extra metadata.
The flag is honoured exactly like the built-in save node does. ``comfy``
only exists inside ComfyUI, hence the guarded runtime import.
"""
if not metadata_text:
try:
from comfy.cli_args import args
except Exception:
return False
return bool(getattr(args, "disable_metadata", False))
def build_workflow_metadata(workflow_prompt: Any, workflow_extra_pnginfo: Any) -> dict[str, str]:
"""
Serialise ComfyUI's hidden ``PROMPT`` / ``EXTRA_PNGINFO`` inputs into the
``{key: json}`` mapping the built-in save node writes into a file.
``EXTRA_PNGINFO`` carries the complete workflow - and with it *every* prompt
of the graph. When a whole batch is saved from a data list there is no single
prompt to point at, so this full workflow is what makes the saved file a
usable workflow again when it is dragged back onto the canvas. The prompt of
the image currently being written is kept *separately* in the
``parameters`` text (see ``compose_prompt_text``) so that it stays easy to
extract afterwards and is still displayed by image viewers.
Returns an empty mapping when nothing is available.
"""
metadata: dict[str, str] = {}
if workflow_prompt is not None:
metadata["prompt"] = json.dumps(workflow_prompt)
if workflow_extra_pnginfo is not None:
for key, value in workflow_extra_pnginfo.items():
metadata[key] = json.dumps(value)
return metadata
def compose_save_metadata(
prompt: str, negative_prompt: str, workflow_prompt: Any, workflow_extra_pnginfo: Any
) -> tuple[str, dict[str, str]]:
"""
Build both metadata channels of an image save.
Returns ``(metadata_text, workflow_metadata)``: the A1111-style
``parameters`` text of the explicit prompt(s) and the ``{key: json}`` mapping
of ComfyUI's workflow prompt info.
ComfyUI's ``--disable-metadata`` flag suppresses *all* metadata, exactly like
the built-in save node does, so a run started with it writes files without
either channel.
"""
if _metadata_disabled():
return "", {}
return compose_prompt_text(prompt, negative_prompt), build_workflow_metadata(workflow_prompt, workflow_extra_pnginfo)
def build_png_info(metadata_text: str, workflow_metadata: Optional[dict[str, str]] = None):
"""
Wrap the metadata in a Pillow ``PngInfo`` container so it can be embedded in
a PNG file.
``metadata_text`` is written under the standard ``parameters`` text-chunk key
(the A1111 convention: shown by image viewers and easy to extract), while
every entry of ``workflow_metadata`` (``prompt``, ``workflow``, ...) gets its
own text chunk - the same keys ComfyUI's built-in "Save Image" node writes,
which is what makes dragging the file back onto the canvas restore the
workflow.
Returns ``None`` when there is nothing to embed (or Pillow's PNG metadata
support is unavailable), in which case the image is saved without metadata.
"""
workflow_metadata = workflow_metadata or {}
if not metadata_text and not workflow_metadata:
return None
try:
from PIL import PngImagePlugin
except ModuleNotFoundError:
return None
pnginfo = PngImagePlugin.PngInfo()
pnginfo.add_text("parameters", metadata_text)
if metadata_text:
pnginfo.add_text("parameters", metadata_text)
for key, value in workflow_metadata.items():
pnginfo.add_text(key, value)
return pnginfo
def build_image_exif(metadata_text: str, include_description: bool = True):
def build_image_exif(metadata_text: str, include_description: bool = True, workflow_metadata: Optional[dict[str, str]] = None):
"""
Build an EXIF block that stores ``metadata_text`` for formats without a
native text chunk (JPEG, WEBP, JXL).
Build an EXIF block for formats without a native text chunk (JPEG, WEBP,
JXL).
The payload is written into the EXIF ``UserComment`` field (tag 0x9286) of
the Exif IFD as ``UNICODE\0`` + UTF-16-BE, which matches what Stable
``metadata_text`` is written into the EXIF ``UserComment`` field (tag 0x9286)
of the Exif IFD as ``UNICODE\0`` + UTF-16-BE, which matches what Stable
Diffusion WebUI / piexif based readers expect. When ``include_description``
is true, the payload is also written as UTF-8 into the EXIF
``ImageDescription`` field (tag 0x010E) of IFD0.
is true, it is also written as UTF-8 into the EXIF ``ImageDescription`` field
(tag 0x010E) of IFD0.
``workflow_metadata`` is stored the same way ComfyUI's own writer does it for
these formats: the workflow prompt in the EXIF ``Model`` tag (0x0110) as
``prompt:<json>`` and the remaining entries (``workflow``, ...) in ``Make``
(0x010F) and downwards as ``<key>:<json>``. ``parameters`` and the workflow
info therefore coexist, and readers looking for either one find it in the
field they expect.
Returns the EXIF bytes (starting with the ``Exif\0\0`` marker), or ``None``
when there is no text to embed.
when there is nothing to embed.
"""
if not metadata_text:
workflow_metadata = workflow_metadata or {}
if not metadata_text and not workflow_metadata:
return None
try:
from PIL import ExifTags
@@ -1145,9 +1223,30 @@ def build_image_exif(metadata_text: str, include_description: bool = True):
return None
Image, _ = _require_pillow()
exif = Image.Exif()
if include_description:
exif[0x010E] = metadata_text.encode("utf-8")
exif.get_ifd(ExifTags.IFD.Exif)[0x9286] = b"UNICODE\x00" + metadata_text.encode("utf-16-be")
used_tags = set()
if metadata_text:
if include_description:
exif[0x010E] = metadata_text.encode("utf-8")
used_tags.add(0x010E)
exif.get_ifd(ExifTags.IFD.Exif)[0x9286] = b"UNICODE\x00" + metadata_text.encode("utf-16-be")
if workflow_metadata:
workflow_json = workflow_metadata.get("prompt")
if workflow_json is not None:
exif[0x0110] = f"prompt:{workflow_json}"
used_tags.add(0x0110)
# Same descending-tag scheme as ComfyUI's own EXIF/WebP metadata writer,
# but skipping the tags the explicit prompt already occupies.
tag = 0x010F
for key, value in workflow_metadata.items():
if key == "prompt":
continue
while tag in used_tags and tag > 0x0100:
tag -= 1
if tag <= 0x0100:
break
exif[tag] = f"{key}:{value}"
used_tags.add(tag)
tag -= 1
return exif.tobytes()
@@ -1173,28 +1272,32 @@ def build_xmp_packet(metadata_text: str) -> bytes:
return packet.encode("utf-8")
def metadata_save_kwargs(metadata_text: str, fmt: str) -> dict:
def metadata_save_kwargs(metadata_text: str, fmt: str, workflow_metadata: Optional[dict[str, str]] = None) -> dict:
"""
Return the extra keyword arguments that embed ``metadata_text`` when saving
an image in the (lower-case) format ``fmt``.
Return the extra keyword arguments that embed ``metadata_text`` (the explicit
prompt) and ``workflow_metadata`` (ComfyUI's workflow prompt info) when
saving an image in the (lower-case) format ``fmt``.
Returns an empty dict when there is no text to embed or when the format
Returns an empty dict when there is nothing to embed or when the format
cannot carry text metadata.
"""
if not metadata_text:
workflow_metadata = workflow_metadata or {}
if not metadata_text and not workflow_metadata:
return {}
if fmt == "png":
return {"pnginfo": build_png_info(metadata_text)}
pnginfo = build_png_info(metadata_text, workflow_metadata)
return {"pnginfo": pnginfo} if pnginfo is not None else {}
if fmt in ("jpg", "jpeg"):
exif = build_image_exif(metadata_text, include_description=True)
exif = build_image_exif(metadata_text, include_description=True, workflow_metadata=workflow_metadata)
return {"exif": exif} if exif is not None else {}
if fmt in ("webp", "jxl"):
exif = build_image_exif(metadata_text, include_description=False)
exif = build_image_exif(metadata_text, include_description=False, workflow_metadata=workflow_metadata)
kwargs = {"exif": exif} if exif is not None else {}
if fmt == "jxl":
# EXIF and XMP boxes are only available in the JXL container format
kwargs["use_container"] = True
kwargs["xmp"] = build_xmp_packet(metadata_text)
if metadata_text:
kwargs["xmp"] = build_xmp_packet(metadata_text)
return kwargs
return {}
@@ -1523,27 +1626,29 @@ def _has_jxl_support() -> bool:
return False
def _write_pil_image(pil_img, path: str, fmt: str, quality: int, metadata_text: str) -> bool:
def _write_pil_image(
pil_img, path: str, fmt: str, quality: int, metadata_text: str, workflow_metadata: Optional[dict[str, str]] = None
) -> 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``
jxl) the prompt text and the workflow metadata are embedded 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))
pil_img.save(path, format="PNG", **metadata_save_kwargs(metadata_text, fmt, workflow_metadata))
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))
pil_img.convert("RGB").save(path, format="JPEG", quality=quality, **metadata_save_kwargs(metadata_text, fmt, workflow_metadata))
elif fmt == "webp":
pil_img.save(path, format="WEBP", quality=quality, **metadata_save_kwargs(metadata_text, fmt))
pil_img.save(path, format="WEBP", quality=quality, **metadata_save_kwargs(metadata_text, fmt, workflow_metadata))
elif fmt == "jxl":
pil_img.save(path, format="JXL", quality=quality, **metadata_save_kwargs(metadata_text, fmt))
pil_img.save(path, format="JXL", quality=quality, **metadata_save_kwargs(metadata_text, fmt, workflow_metadata))
else:
if metadata_text:
print("Basic data handling: Prompt metadata is not supported for this format; skipping it.")
if metadata_text or workflow_metadata:
print("Basic data handling: Prompt/workflow 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:
@@ -1551,6 +1656,37 @@ def _write_pil_image(pil_img, path: str, fmt: str, quality: int, metadata_text:
return False
def _ui_image_entry(path: str) -> Optional[dict[str, str]]:
"""
Build the ComfyUI ``ui.images`` entry that makes a saved file show up.
ComfyUI displays previews and lists generated images by resolving
``/view?filename=..&subfolder=..&type=..``, where ``type`` selects the
output/temp (or input) directory and any path outside it is refused. A file
written somewhere else (plain ``path`` mode pointing outside the ComfyUI
folders) is still saved, it simply cannot be previewed, so ``None`` is
returned for it. ``input`` is deliberately never reported: the frontend
treats an output whose images are all ``type: "input"`` as the upload
preview of a load node.
Returns ``{"filename": ..., "subfolder": ..., "type": ...}`` for a file inside
one of the folders ComfyUI serves, or ``None`` when there is no such folder.
"""
directory = os.path.dirname(os.path.abspath(path))
for folder_type, folder in (("output", get_output_directory()), ("temp", get_temp_directory())):
try:
real_folder = os.path.realpath(folder)
if os.path.commonpath([real_folder, os.path.realpath(directory)]) != real_folder:
continue
except ValueError: # e.g. different drives on Windows
continue
subfolder = os.path.relpath(os.path.realpath(directory), real_folder).replace(os.sep, "/")
if subfolder == ".":
subfolder = ""
return {"filename": os.path.basename(path), "subfolder": subfolder, "type": folder_type}
return None
class PathSaveImageRGB(ComfyNodeABC):
"""
Saves an image to a file.
@@ -1564,11 +1700,21 @@ class PathSaveImageRGB(ComfyNodeABC):
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
the EXIF + XMP boxes for JPEG XL. Formats that cannot carry text metadata
ignore the prompts.
Files written into ComfyUI's output (or temp) folder are reported back to the
frontend, so the node shows a preview of what it saved and the images appear
in ComfyUI's generated-images list, exactly like the built-in "Save Image"
node does. Files written anywhere else (plain ``path`` mode) are still saved,
they just cannot be previewed.
Two independent metadata channels are embedded, so both stay usable: with
``prompt``/``negative_prompt`` the A1111-style ``parameters`` text is written
(PNG text chunk, EXIF ``UserComment``, XMP for JPEG XL) which image viewers
display and which is easy to extract afterwards, and in addition ComfyUI's
hidden workflow prompt info (``prompt``/``workflow`` chunks) is stored, which
is what makes a saved file drag-and-drop back into ComfyUI as a workflow.
When a batch is saved from a data list the workflow necessarily contains
every prompt of the graph; the prompt of the individual image is the one in
``parameters``. Formats that cannot carry text metadata ignore both.
"""
@classmethod
@@ -1607,11 +1753,21 @@ class PathSaveImageRGB(ComfyNodeABC):
},
),
},
"hidden": {
# Deliberately not called "prompt"/"extra_pnginfo": hidden input
# names become function-argument names and "prompt" is already the
# explicit-prompt widget of this node.
"workflow_prompt": "PROMPT",
"workflow_extra_pnginfo": "EXTRA_PNGINFO",
},
}
RETURN_TYPES = (IO.BOOLEAN,)
RETURN_NAMES = ("success",)
OUTPUT_TOOLTIPS = ("True when the image was saved successfully.",)
RETURN_TYPES = (IO.BOOLEAN, IO.IMAGE)
RETURN_NAMES = ("success", "images")
OUTPUT_TOOLTIPS = (
"True when the image was saved successfully.",
"The saved images, passed through unchanged so they can be used further down the graph.",
)
CATEGORY = "Basic/Path"
DESCRIPTION = cleandoc(__doc__ or "")
FUNCTION = "save_image"
@@ -1627,10 +1783,12 @@ class PathSaveImageRGB(ComfyNodeABC):
create_dirs: bool = True,
prompt: str = "",
negative_prompt: str = "",
workflow_prompt: Any = None,
workflow_extra_pnginfo: Any = None,
):
if not path:
print("Basic data handling: Save failed - no path specified")
return (False,)
return {"ui": {"images": []}, "result": (False, images)}
import numpy as np
from PIL import Image
@@ -1641,14 +1799,19 @@ class PathSaveImageRGB(ComfyNodeABC):
"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,)
return {"ui": {"images": []}, "result": (False, images)}
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)
# Compose the prompt metadata to embed into the saved file once.
metadata_text = compose_prompt_text(prompt, negative_prompt)
# Compose both metadata channels: the explicit prompt(s) of this save
# (the A1111-style "parameters" text) and ComfyUI's workflow prompt info.
metadata_text, workflow_metadata = compose_save_metadata(prompt, negative_prompt, workflow_prompt, workflow_extra_pnginfo)
# Reporting the written files makes the frontend preview them and list
# them with the generated images (same payload as the built-in node).
ui_images = []
for index, target_path in enumerate(paths):
_ensure_parent_directories(target_path, create_dirs)
@@ -1658,11 +1821,15 @@ class PathSaveImageRGB(ComfyNodeABC):
img_np = (img_tensor * 255).astype(np.uint8)
pil_img = Image.fromarray(img_np)
if not _write_pil_image(pil_img, target_path, fmt, quality, metadata_text):
return (False,)
if not _write_pil_image(pil_img, target_path, fmt, quality, metadata_text, workflow_metadata):
return {"ui": {"images": ui_images}, "result": (False, images)}
entry = _ui_image_entry(target_path)
if entry is not None:
ui_images.append(entry)
print(f"Basic data handling: Successfully saved image to {target_path}")
return (True,)
return {"ui": {"images": ui_images}, "result": (True, images)}
class PathSaveImageRGBA(ComfyNodeABC):
@@ -1679,11 +1846,21 @@ class PathSaveImageRGBA(ComfyNodeABC):
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
the EXIF + XMP boxes for JPEG XL. Formats that cannot carry text metadata
ignore the prompts.
Files written into ComfyUI's output (or temp) folder are reported back to the
frontend, so the node shows a preview of what it saved and the images appear
in ComfyUI's generated-images list, exactly like the built-in "Save Image"
node does. Files written anywhere else (plain ``path`` mode) are still saved,
they just cannot be previewed.
Two independent metadata channels are embedded, so both stay usable: with
``prompt``/``negative_prompt`` the A1111-style ``parameters`` text is written
(PNG text chunk, EXIF ``UserComment``, XMP for JPEG XL) which image viewers
display and which is easy to extract afterwards, and in addition ComfyUI's
hidden workflow prompt info (``prompt``/``workflow`` chunks) is stored, which
is what makes a saved file drag-and-drop back into ComfyUI as a workflow.
When a batch is saved from a data list the workflow necessarily contains
every prompt of the graph; the prompt of the individual image is the one in
``parameters``. Formats that cannot carry text metadata ignore both.
"""
@classmethod
@@ -1724,11 +1901,21 @@ class PathSaveImageRGBA(ComfyNodeABC):
},
),
},
"hidden": {
# Deliberately not called "prompt"/"extra_pnginfo": hidden input
# names become function-argument names and "prompt" is already the
# explicit-prompt widget of this node.
"workflow_prompt": "PROMPT",
"workflow_extra_pnginfo": "EXTRA_PNGINFO",
},
}
RETURN_TYPES = (IO.BOOLEAN,)
RETURN_NAMES = ("success",)
OUTPUT_TOOLTIPS = ("True when the image with alpha was saved successfully.",)
RETURN_TYPES = (IO.BOOLEAN, IO.IMAGE)
RETURN_NAMES = ("success", "images")
OUTPUT_TOOLTIPS = (
"True when the image with alpha was saved successfully.",
"The saved images, passed through unchanged so they can be used further down the graph.",
)
CATEGORY = "Basic/Path"
DESCRIPTION = cleandoc(__doc__ or "")
FUNCTION = "save_image_with_mask"
@@ -1746,10 +1933,12 @@ class PathSaveImageRGBA(ComfyNodeABC):
create_dirs: bool = True,
prompt: str = "",
negative_prompt: str = "",
workflow_prompt: Any = None,
workflow_extra_pnginfo: Any = None,
):
if not path:
print("Basic data handling: Save failed - no path specified")
return (False,)
return {"ui": {"images": []}, "result": (False, images)}
import numpy as np
from PIL import Image
@@ -1764,7 +1953,7 @@ class PathSaveImageRGBA(ComfyNodeABC):
"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,)
return {"ui": {"images": []}, "result": (False, images)}
batch = len(images)
height, width = images.shape[1], images.shape[2]
@@ -1775,7 +1964,13 @@ class PathSaveImageRGBA(ComfyNodeABC):
# mask has fewer frames (e.g. a single mask applied to every frame).
mask_frames = mask.cpu()
metadata_text = compose_prompt_text(prompt, negative_prompt)
# Compose both metadata channels: the explicit prompt(s) of this save
# (the A1111-style "parameters" text) and ComfyUI's workflow prompt info.
metadata_text, workflow_metadata = compose_save_metadata(prompt, negative_prompt, workflow_prompt, workflow_extra_pnginfo)
# Reporting the written files makes the frontend preview them and list
# them with the generated images (same payload as the built-in node).
ui_images = []
for index, target_path in enumerate(paths):
_ensure_parent_directories(target_path, create_dirs)
@@ -1803,11 +1998,15 @@ class PathSaveImageRGBA(ComfyNodeABC):
pil_img_rgba = pil_img.convert("RGBA")
pil_img_rgba.putalpha(alpha_img)
if not _write_pil_image(pil_img_rgba, target_path, fmt, quality, metadata_text):
return (False,)
if not _write_pil_image(pil_img_rgba, target_path, fmt, quality, metadata_text, workflow_metadata):
return {"ui": {"images": ui_images}, "result": (False, images)}
entry = _ui_image_entry(target_path)
if entry is not None:
ui_images.append(entry)
print(f"Basic data handling: Successfully saved image with mask to {target_path}")
return (True,)
return {"ui": {"images": ui_images}, "result": (True, images)}
class PathInputDir(ComfyNodeABC):
tests/test_path_nodes.py
+220 -26
View File
@@ -1,4 +1,5 @@
import os
import json
import pytest
import platform
import numpy as np
@@ -50,6 +51,29 @@ def decode_user_comment(value):
return value
def save_success(result):
"""Success flag of a save-node result.
The save nodes return ComfyUI's ``{"ui": ..., "result": ...}`` form so that
the frontend previews/lists the written image, so the socket outputs live
under ``result``.
"""
assert isinstance(result, dict), result
return result["result"][0]
def save_passthrough_images(result):
"""The IMAGE passthrough output of a save-node result."""
assert isinstance(result, dict), result
return result["result"][1]
def save_ui_images(result):
"""The ``ui.images`` entries a save-node result reports for the frontend."""
assert isinstance(result, dict), result
return result["ui"]["images"]
def test_path_join():
node = PathJoin()
assert node.join_paths("folder", "file.txt") == (os.path.join("folder", "file.txt"),)
@@ -429,17 +453,17 @@ def test_path_load_save_image_rgba(tmp_path, monkeypatch):
mask[0, :, i] = i / img_size[0]
# Save the image with mask
assert save_node.save_image_with_mask(red_img, mask, output_path) == (True,)
assert save_success(save_node.save_image_with_mask(red_img, mask, output_path)) is True
# Verify the image was saved with the correct extension
assert os.path.exists(output_path + ".png")
# Test with invert_mask option
assert save_node.save_image_with_mask(red_img, mask, str(tmp_path / "inverted"), invert_mask=True) == (True,)
assert save_success(save_node.save_image_with_mask(red_img, mask, str(tmp_path / "inverted"), invert_mask=True)) is True
assert os.path.exists(str(tmp_path / "inverted.png"))
# Test with JPEG format (should switch to PNG for transparency)
assert save_node.save_image_with_mask(red_img, mask, str(tmp_path / "jpeg_test"), format="jpg") == (True,)
assert save_success(save_node.save_image_with_mask(red_img, mask, str(tmp_path / "jpeg_test"), format="jpg")) is True
# Should be saved as PNG despite the request for JPEG
assert os.path.exists(str(tmp_path / "jpeg_test.png"))
@@ -478,18 +502,18 @@ def test_path_load_save_image_rgb(tmp_path, monkeypatch):
red_img[0, :, :, 0] = 1.0 # Red channel set to 1
# Save the image
assert save_node.save_image(red_img, output_path) == (True,)
assert save_success(save_node.save_image(red_img, output_path)) is True
# Verify the image was saved with the correct extension
assert os.path.exists(output_path + ".png")
# Test different formats
assert save_node.save_image(red_img, str(tmp_path / "jpeg_test"), format="jpg") == (True,)
assert save_success(save_node.save_image(red_img, str(tmp_path / "jpeg_test"), format="jpg")) is True
assert os.path.exists(str(tmp_path / "jpeg_test.jpg"))
# Test with directories that don't exist
nested_path = str(tmp_path / "nested" / "images" / "test_rgb")
assert save_node.save_image(red_img, nested_path) == (True,)
assert save_success(save_node.save_image(red_img, nested_path)) is True
assert os.path.exists(nested_path + ".png")
# Test error handling for loading
@@ -505,35 +529,38 @@ def test_path_save_image_rgb_prompt_metadata(tmp_path):
# No metadata when neither prompt nor negative prompt is provided
plain_path = str(tmp_path / "plain")
assert save_node.save_image(red_img, plain_path) == (True,)
assert save_success(save_node.save_image(red_img, plain_path)) is True
with Image.open(plain_path + ".png") as img:
img.load()
assert img.text.get("parameters") is None
# Both prompt and negative prompt provided
full_path = str(tmp_path / "full")
assert save_node.save_image(red_img, full_path, prompt="a red square", negative_prompt="blurry, low quality") == (True,)
assert save_success(save_node.save_image(red_img, full_path, prompt="a red square", negative_prompt="blurry, low quality")) is True
with Image.open(full_path + ".png") as img:
img.load()
assert img.text.get("parameters") == "a red square\nNegative prompt: blurry, low quality"
# Only the prompt is provided
pos_path = str(tmp_path / "pos_only")
assert save_node.save_image(red_img, pos_path, prompt="only positive") == (True,)
assert save_success(save_node.save_image(red_img, pos_path, prompt="only positive")) is True
with Image.open(pos_path + ".png") as img:
img.load()
assert img.text.get("parameters") == "only positive"
# Only the negative prompt is provided
neg_path = str(tmp_path / "neg_only")
assert save_node.save_image(red_img, neg_path, negative_prompt="only negative") == (True,)
assert save_success(save_node.save_image(red_img, neg_path, negative_prompt="only negative")) is True
with Image.open(neg_path + ".png") as img:
img.load()
assert img.text.get("parameters") == "Negative prompt: only negative"
# JPEG embeds the prompt into the EXIF UserComment and ImageDescription fields
jpg_path = str(tmp_path / "jpeg_meta")
assert save_node.save_image(red_img, jpg_path, format="jpg", prompt="a red square", negative_prompt="blurry, low quality") == (True,)
assert (
save_success(save_node.save_image(red_img, jpg_path, format="jpg", prompt="a red square", negative_prompt="blurry, low quality"))
is True
)
assert os.path.exists(jpg_path + ".jpg")
with Image.open(jpg_path + ".jpg") as img:
img.load()
@@ -544,7 +571,10 @@ def test_path_save_image_rgb_prompt_metadata(tmp_path):
# WEBP embeds the prompt into the EXIF UserComment field (no ImageDescription)
webp_path = str(tmp_path / "webp_meta")
assert save_node.save_image(red_img, webp_path, format="webp", prompt="a red square", negative_prompt="blurry, low quality") == (True,)
assert (
save_success(save_node.save_image(red_img, webp_path, format="webp", prompt="a red square", negative_prompt="blurry, low quality"))
is True
)
with Image.open(webp_path + ".webp") as img:
img.load()
exif_bytes = img.info.get("exif", b"")
@@ -565,14 +595,17 @@ def test_path_save_image_rgba_prompt_metadata(tmp_path):
# No metadata when neither prompt nor negative prompt is provided
plain_path = str(tmp_path / "plain")
assert save_node.save_image_with_mask(red_img, mask, plain_path) == (True,)
assert save_success(save_node.save_image_with_mask(red_img, mask, plain_path)) is True
with Image.open(plain_path + ".png") as img:
img.load()
assert img.text.get("parameters") is None
# Both prompt and negative prompt provided
full_path = str(tmp_path / "full")
assert save_node.save_image_with_mask(red_img, mask, full_path, prompt="a red square", negative_prompt="blurry, low quality") == (True,)
assert (
save_success(save_node.save_image_with_mask(red_img, mask, full_path, prompt="a red square", negative_prompt="blurry, low quality"))
is True
)
with Image.open(full_path + ".png") as img:
img.load()
assert img.text.get("parameters") == "a red square\nNegative prompt: blurry, low quality"
@@ -588,7 +621,7 @@ def test_path_save_image_jxl_prompt_metadata(tmp_path):
prompt = "a red square <lora:x:1.0>"
negative_prompt = "blurry, low quality"
jxl_path = str(tmp_path / "jxl_meta")
assert save_node.save_image(red_img, jxl_path, format="jxl", prompt=prompt, negative_prompt=negative_prompt) == (True,)
assert save_success(save_node.save_image(red_img, jxl_path, format="jxl", prompt=prompt, negative_prompt=negative_prompt)) is True
assert os.path.exists(jxl_path + ".jxl")
# JXL embeds the prompt into the EXIF UserComment field
@@ -671,7 +704,7 @@ def test_path_save_image_formats_are_library_derived():
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
assert save_success(save_node.save_image(img, out, format=fmt)) is True, fmt
# correct file extension appended
expected = f"{out}.{fmt}"
assert os.path.exists(expected), fmt
@@ -687,7 +720,7 @@ def test_path_save_image_rgb_multi_frame_suffix(tmp_path):
frames = _batch_rgb(3)
out = str(tmp_path / "frames")
assert save_node.save_image(frames, out) == (True,)
assert save_success(save_node.save_image(frames, out)) is True
expected = [
str(tmp_path / "frames_00000.png"),
str(tmp_path / "frames_00001.png"),
@@ -705,12 +738,12 @@ def test_path_save_image_prefix_mode_mirrors_comfy(tmp_path, monkeypatch):
img = _batch_rgb(1)
prefix = "my_subdir/shot"
assert save_node.save_image(img, prefix, use_prefix_mode=True) == (True,)
assert save_success(save_node.save_image(img, prefix, use_prefix_mode=True)) is 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 save_success(save_node.save_image(img, prefix, use_prefix_mode=True)) is True
assert (tmp_path / "my_subdir" / "shot_00002_.png").exists()
@@ -719,7 +752,7 @@ def test_path_save_image_prefix_mode_batch(tmp_path, monkeypatch):
save_node = PathSaveImageRGB()
frames = _batch_rgb(2)
assert save_node.save_image(frames, "batch", use_prefix_mode=True) == (True,)
assert save_success(save_node.save_image(frames, "batch", use_prefix_mode=True)) is True
assert (tmp_path / "batch_00001_.png").exists()
assert (tmp_path / "batch_00002_.png").exists()
@@ -730,7 +763,7 @@ def test_path_save_image_prefix_mode_template_tokens(tmp_path, monkeypatch):
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 save_success(save_node.save_image(img, "w%width%_h%height%", use_prefix_mode=True)) is True
assert (tmp_path / "w8_h4_00001_.png").exists()
@@ -743,7 +776,7 @@ def test_path_save_image_prefix_mode_date_token(tmp_path, monkeypatch):
save_node = PathSaveImageRGB()
img = _batch_rgb(1)
assert save_node.save_image(img, "%date:yyyy-MM-dd%/dated_test", use_prefix_mode=True) == (True,)
assert save_success(save_node.save_image(img, "%date:yyyy-MM-dd%/dated_test", use_prefix_mode=True)) is 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.
@@ -760,7 +793,7 @@ def test_path_save_image_path_mode_date_token(tmp_path):
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 save_success(save_node.save_image(img, out)) is True
assert not (tmp_path / "%date:yyyy-MM-dd%").exists()
assert (tmp_path / today / "dated.png").exists()
@@ -768,7 +801,11 @@ def test_path_save_image_path_mode_date_token(tmp_path):
def test_path_save_image_empty_path_returns_false():
save_node = PathSaveImageRGB()
img = _batch_rgb(1)
assert save_node.save_image(img, "") == (False,)
result = save_node.save_image(img, "")
assert save_success(result) is False
# The images are passed through even when nothing was written.
assert save_passthrough_images(result) is img
assert save_ui_images(result) == []
def test_path_save_image_rgba_prefix_mode_and_transparency(tmp_path, monkeypatch):
@@ -779,7 +816,7 @@ def test_path_save_image_rgba_prefix_mode_and_transparency(tmp_path, monkeypatch
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,)
assert save_success(save_node.save_image_with_mask(img, mask, "alpha", use_prefix_mode=True)) is True
path = tmp_path / "alpha_00001_.png"
assert path.exists()
@@ -798,11 +835,168 @@ def test_path_save_image_rgba_multi_frame(tmp_path):
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 save_success(save_node.save_image_with_mask(frames, mask, str(tmp_path / "rgba_batch"))) is True
assert (tmp_path / "rgba_batch_00000.png").exists()
assert (tmp_path / "rgba_batch_00001.png").exists()
def test_path_save_image_reports_preview_entries(tmp_path, monkeypatch):
"""Files saved into ComfyUI's output folder are reported as ``ui.images``,
which is what makes the frontend preview them and list them with the
generated images (same payload as the built-in "Save Image" node)."""
monkeypatch.setattr("src.basic_data_handling.path_nodes.get_output_directory", lambda: str(tmp_path))
save_node = PathSaveImageRGB()
frames = _batch_rgb(2)
result = save_node.save_image(frames, "my_subdir/shot", use_prefix_mode=True)
assert save_ui_images(result) == [
{"filename": "shot_00001_.png", "subfolder": "my_subdir", "type": "output"},
{"filename": "shot_00002_.png", "subfolder": "my_subdir", "type": "output"},
]
# The images are passed through unchanged, like the built-in node does.
assert save_passthrough_images(result) is frames
def test_path_save_image_rgba_reports_preview_entries(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)
mask = torch.zeros(1, size[1], size[0])
result = save_node.save_image_with_mask(img, mask, "alpha", use_prefix_mode=True)
assert save_ui_images(result) == [{"filename": "alpha_00001_.png", "subfolder": "", "type": "output"}]
def test_path_save_image_plain_path_inside_output_dir_is_previewed(tmp_path, monkeypatch):
"""Plain path mode is previewable as long as the file lands in the output folder."""
monkeypatch.setattr("src.basic_data_handling.path_nodes.get_output_directory", lambda: str(tmp_path))
save_node = PathSaveImageRGB()
img = _batch_rgb(1)
result = save_node.save_image(img, str(tmp_path / "sub" / "plain"))
assert save_ui_images(result) == [{"filename": "plain.png", "subfolder": "sub", "type": "output"}]
def test_path_save_image_temp_dir_is_previewed_as_temp(tmp_path, monkeypatch):
"""Files in ComfyUI's temp folder are reported with ``type: "temp"``."""
monkeypatch.setattr("src.basic_data_handling.path_nodes.get_output_directory", lambda: str(tmp_path / "out"))
monkeypatch.setattr("src.basic_data_handling.path_nodes.get_temp_directory", lambda: str(tmp_path))
save_node = PathSaveImageRGB()
img = _batch_rgb(1)
result = save_node.save_image(img, str(tmp_path / "scratch"))
assert save_ui_images(result) == [{"filename": "scratch.png", "subfolder": "", "type": "temp"}]
def test_path_save_image_outside_comfy_folders_has_no_preview(tmp_path, monkeypatch):
"""The file is still written, but nothing is reported for it: ComfyUI could
not serve it through /view, so a preview entry would be a broken image."""
monkeypatch.setattr("src.basic_data_handling.path_nodes.get_output_directory", lambda: str(tmp_path / "out"))
monkeypatch.setattr("src.basic_data_handling.path_nodes.get_temp_directory", lambda: str(tmp_path / "tmp"))
save_node = PathSaveImageRGB()
img = _batch_rgb(1)
out = str(tmp_path / "elsewhere" / "image")
result = save_node.save_image(img, out)
assert save_success(result) is True
assert save_ui_images(result) == []
assert os.path.exists(out + ".png")
def test_path_save_image_exposes_workflow_inputs_and_image_passthrough():
"""Drop-in superset of the built-in node: it takes ComfyUI's hidden workflow
prompt info and passes the images through as a second output."""
for cls in (PathSaveImageRGB, PathSaveImageRGBA):
assert cls.INPUT_TYPES()["hidden"] == {"workflow_prompt": "PROMPT", "workflow_extra_pnginfo": "EXTRA_PNGINFO"}
assert cls.RETURN_TYPES == ("BOOLEAN", "IMAGE")
assert cls.RETURN_NAMES == ("success", "images")
assert cls.OUTPUT_NODE is True
def test_path_save_image_embeds_workflow_and_explicit_prompt(tmp_path, monkeypatch):
"""Both metadata channels coexist in a PNG: the workflow (all prompts of the
graph, needed so dragging the file back restores the workflow) and the
explicit prompt of this save as ``parameters`` (easy to extract, shown by
image viewers)."""
monkeypatch.setattr("src.basic_data_handling.path_nodes.get_output_directory", lambda: str(tmp_path))
save_node = PathSaveImageRGB()
frames = _batch_rgb(2)
workflow_prompt = {
"3": {"class_type": "KSampler", "inputs": {"seed": 1}},
"6": {"class_type": "CLIPTextEncode", "inputs": {"text": "prompt of a different node"}},
}
extra_pnginfo = {"workflow": {"nodes": [{"id": 3}, {"id": 6}]}}
result = save_node.save_image(
frames,
"meta",
use_prefix_mode=True,
prompt="the prompt of this image",
negative_prompt="bad",
workflow_prompt=workflow_prompt,
workflow_extra_pnginfo=extra_pnginfo,
)
assert save_success(result) is True
# Every frame of the batch carries both channels.
for name in ("meta_00001_.png", "meta_00002_.png"):
with Image.open(tmp_path / name) as img:
img.load()
assert img.text["parameters"] == "the prompt of this image\nNegative prompt: bad"
assert json.loads(img.text["prompt"]) == workflow_prompt
assert json.loads(img.text["workflow"]) == extra_pnginfo["workflow"]
def test_path_save_image_workflow_metadata_in_exif(tmp_path, monkeypatch):
"""Formats without a native text chunk keep the explicit prompt in
UserComment and store the workflow in the fields ComfyUI uses for them."""
monkeypatch.setattr("src.basic_data_handling.path_nodes.get_output_directory", lambda: str(tmp_path))
save_node = PathSaveImageRGB()
img = _batch_rgb(1)
workflow_prompt = {"3": {"class_type": "KSampler", "inputs": {}}}
extra_pnginfo = {"workflow": {"nodes": [{"id": 3}]}}
out = str(tmp_path / "exif_meta")
result = save_node.save_image(
img, out, format="jpg", prompt="explicit", workflow_prompt=workflow_prompt, workflow_extra_pnginfo=extra_pnginfo
)
assert save_success(result) is True
with Image.open(out + ".jpg") as loaded:
loaded.load()
exif = loaded.getexif()
assert decode_user_comment(exif.get_ifd(0x8769).get(0x9286)) == "explicit"
# ComfyUI's own convention for these formats.
assert exif.get(0x0110) == "prompt:" + json.dumps(workflow_prompt)
assert exif.get(0x010F) == "workflow:" + json.dumps(extra_pnginfo["workflow"])
def test_path_save_image_metadata_can_be_disabled(tmp_path, monkeypatch):
"""ComfyUI's ``--disable-metadata`` suppresses both metadata channels, like
it does in the built-in save node."""
import src.basic_data_handling.path_nodes as pn
monkeypatch.setattr(pn, "_metadata_disabled", lambda: True)
save_node = PathSaveImageRGB()
img = _batch_rgb(1)
out = str(tmp_path / "disabled")
result = save_node.save_image(img, out, prompt="explicit", workflow_prompt={"1": {}}, workflow_extra_pnginfo={"workflow": {}})
assert save_success(result) is True
with Image.open(out + ".png") as loaded:
loaded.load()
assert loaded.text.get("parameters") is None
assert loaded.text.get("prompt") is None
assert loaded.text.get("workflow") is None
def test_path_normalize():
node = PathNormalize()
assert node.normalize_path("folder/../file.txt") == (os.path.normpath("folder/../file.txt"),)