From 7a9e7c61eaee88afc2eec3cce44afacc38154f52 Mon Sep 17 00:00:00 2001 From: StableLlama Date: Fri, 18 Sep 2026 22:04:36 +0200 Subject: [PATCH 1/3] fix typo --- src/basic_data_handling/string_nodes.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/basic_data_handling/string_nodes.py b/src/basic_data_handling/string_nodes.py index f7b3714..4bafda3 100644 --- a/src/basic_data_handling/string_nodes.py +++ b/src/basic_data_handling/string_nodes.py @@ -1697,7 +1697,7 @@ def zfill(self, string, width): "Basic data handling: StringRstrip": "rstrip", "Basic data handling: StringSplitDataList": "split (to data list)", "Basic data handling: StringSplitList": "split (to LIST)", - "Basic data handling: StringSplitlinesDataList": "splitlines (from data list)", + "Basic data handling: StringSplitlinesDataList": "splitlines (to data list)", "Basic data handling: StringSplitlinesList": "splitlines (to LIST)", "Basic data handling: StringStartswith": "startswith", "Basic data handling: StringStrip": "strip", From 1cdbd3044c4666ff955a8f490aadca7619662c96 Mon Sep 17 00:00:00 2001 From: StableLlama Date: Fri, 18 Sep 2026 22:06:43 +0200 Subject: [PATCH 2/3] Save workflow in the image and show preview --- src/basic_data_handling/path_nodes.py | 335 ++++++++++++++++++++------ tests/test_path_nodes.py | 246 +++++++++++++++++-- 2 files changed, 487 insertions(+), 94 deletions(-) diff --git a/src/basic_data_handling/path_nodes.py b/src/basic_data_handling/path_nodes.py index d556129..31a9737 100644 --- a/src/basic_data_handling/path_nodes.py +++ b/src/basic_data_handling/path_nodes.py @@ -1,4 +1,6 @@ from inspect import cleandoc +from typing import Any, Optional +import json import os import glob @@ -19,13 +21,14 @@ class IO: 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:`` and the remaining entries (``workflow``, ...) in ``Make`` + (0x010F) and downwards as ``:``. ``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 @@ def INPUT_TYPES(cls): }, ), }, + "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 @@ def save_image( 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 @@ def save_image( "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 @@ def save_image( 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 @@ def INPUT_TYPES(cls): }, ), }, + "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 @@ def save_image_with_mask( 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 @@ def save_image_with_mask( "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 @@ def save_image_with_mask( # 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 @@ def save_image_with_mask( 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): diff --git a/tests/test_path_nodes.py b/tests/test_path_nodes.py index 89dcd1c..d023380 100644 --- a/tests/test_path_nodes.py +++ b/tests/test_path_nodes.py @@ -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 mock_extract_mask_from_alpha(img): 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 mock_load_image_helper(path): 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 " 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"),) From 20c6e9618516e18caf86f3b23d0bddf93edc8272 Mon Sep 17 00:00:00 2001 From: StableLlama Date: Sat, 19 Sep 2026 00:17:10 +0200 Subject: [PATCH 3/3] bump to version 2.1.0 --- CHANGELOG.md | 25 +++++++++++++++++++++++++ pyproject.toml | 2 +- 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0158609..00fcd13 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,31 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [2.1.0] - 2026-09-19 + +Backward compatible changes: +- "save IMAGE to file" and "save IMAGE+MASK to file" can now be used like + ComfyUI's built-in "Save Image" node: images written into the 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. In addition + both nodes got a second `images` output that passes the images through + unchanged, like the built-in node does. + - files written somewhere else (plain `path` mode) are still saved, they just + cannot be previewed, as ComfyUI can only serve files from its own folders +- Both save nodes now embed ComfyUI's workflow metadata (the API prompt and the + full workflow) in addition to the existing prompt metadata. Dragging a saved + file back onto the canvas therefore restores the workflow, while the explicitly + given `prompt` / `negative_prompt` remain in the `parameters` text that image + viewers display and that is easy to extract afterwards. This is what makes the + two channels worth having when a whole batch is saved from a data list: the + workflow then contains every prompt of the graph, whereas `parameters` always + holds the prompt of the image that was actually written. + - ComfyUI's `--disable-metadata` option suppresses both channels, exactly like + it does for the built-in save node +- No widget was added, removed or reordered, so existing workflows keep their + saved widget values. + + ## [2.0.0] - 2026-09-11 Critical change that ensures reliablity and **should** have no bad side diff --git a/pyproject.toml b/pyproject.toml index 2a9334a..e71d132 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "basic_data_handling" -version = "2.0.0" +version = "2.1.0" description = """Basic Python functions for manipulating data that every programmer is used to, lightweight with no additional dependencies. Supported data types: