From 16e30ea6894606cd1227fd0df77ec04b0169f9d8 Mon Sep 17 00:00:00 2001 From: Will Miao Date: Sat, 28 Mar 2026 11:17:36 +0800 Subject: [PATCH] fix(nodes): add save_with_metadata toggle to save image --- py/nodes/save_image.py | 16 ++++- tests/nodes/test_save_image.py | 120 +++++++++++++++++++++++++++++++++ 2 files changed, 133 insertions(+), 3 deletions(-) create mode 100644 tests/nodes/test_save_image.py diff --git a/py/nodes/save_image.py b/py/nodes/save_image.py index 78b70394..b4022667 100644 --- a/py/nodes/save_image.py +++ b/py/nodes/save_image.py @@ -72,6 +72,13 @@ class SaveImageLM: "tooltip": "Embeds the complete workflow data into the image metadata. Only works with PNG and WebP formats.", }, ), + "save_with_metadata": ( + "BOOLEAN", + { + "default": True, + "tooltip": "When enabled, embeds generation parameters into the saved image metadata. Disable to skip writing generation metadata.", + }, + ), "add_counter_to_filename": ( "BOOLEAN", { @@ -350,6 +357,7 @@ class SaveImageLM: lossless_webp=True, quality=100, embed_workflow=False, + save_with_metadata=True, add_counter_to_filename=True, ): """Save images with metadata""" @@ -421,7 +429,7 @@ class SaveImageLM: try: if file_format == "png": assert pnginfo is not None - if metadata: + if save_with_metadata and metadata: pnginfo.add_text("parameters", metadata) if embed_workflow and extra_pnginfo is not None: workflow_json = json.dumps(extra_pnginfo["workflow"]) @@ -430,7 +438,7 @@ class SaveImageLM: img.save(file_path, format="PNG", **save_kwargs) elif file_format == "jpeg": # For JPEG, use piexif - if metadata: + if save_with_metadata and metadata: try: exif_dict = { "Exif": { @@ -448,7 +456,7 @@ class SaveImageLM: # For WebP, use piexif for metadata exif_dict = {} - if metadata: + if save_with_metadata and metadata: exif_dict["Exif"] = { piexif.ExifIFD.UserComment: b"UNICODE\0" + metadata.encode("utf-16be") @@ -489,6 +497,7 @@ class SaveImageLM: lossless_webp=True, quality=100, embed_workflow=False, + save_with_metadata=True, add_counter_to_filename=True, ): """Process and save image with metadata""" @@ -516,6 +525,7 @@ class SaveImageLM: lossless_webp, quality, embed_workflow, + save_with_metadata, add_counter_to_filename, ) diff --git a/tests/nodes/test_save_image.py b/tests/nodes/test_save_image.py new file mode 100644 index 00000000..6b5a26db --- /dev/null +++ b/tests/nodes/test_save_image.py @@ -0,0 +1,120 @@ +import json + +import numpy as np +import piexif +from PIL import Image + +from py.nodes.save_image import SaveImageLM + + +class _DummyTensor: + def __init__(self, array): + self._array = array + self.shape = array.shape + + def cpu(self): + return self + + def numpy(self): + return self._array + + +def _make_image(): + return _DummyTensor( + np.array( + [ + [[0.0, 0.1, 0.2], [0.3, 0.4, 0.5]], + [[0.6, 0.7, 0.8], [0.9, 1.0, 0.0]], + ], + dtype="float32", + ) + ) + + +def _configure_save_paths(monkeypatch, tmp_path): + monkeypatch.setattr("folder_paths.get_output_directory", lambda: str(tmp_path), raising=False) + monkeypatch.setattr( + "folder_paths.get_save_image_path", + lambda *_args, **_kwargs: (str(tmp_path), "sample", 1, "", "sample"), + raising=False, + ) + + +def _configure_metadata(monkeypatch, metadata_dict): + monkeypatch.setattr("py.nodes.save_image.get_metadata", lambda: {"raw": "metadata"}) + monkeypatch.setattr( + "py.nodes.save_image.MetadataProcessor.to_dict", + lambda raw_metadata, node_id: metadata_dict, + ) + + +def test_save_image_defaults_to_writing_png_metadata(monkeypatch, tmp_path): + _configure_save_paths(monkeypatch, tmp_path) + _configure_metadata(monkeypatch, {"prompt": "prompt text", "seed": 123}) + + node = SaveImageLM() + node.save_images([_make_image()], "ComfyUI", "png", id="node-1") + + image_path = tmp_path / "sample_00001_.png" + with Image.open(image_path) as img: + assert img.info["parameters"] == "prompt text\nSeed: 123" + + +def test_save_image_skips_png_parameters_when_metadata_disabled_and_keeps_workflow( + monkeypatch, tmp_path +): + _configure_save_paths(monkeypatch, tmp_path) + _configure_metadata(monkeypatch, {"prompt": "prompt text", "seed": 123}) + + node = SaveImageLM() + workflow = {"nodes": [{"id": 1}]} + node.save_images( + [_make_image()], + "ComfyUI", + "png", + id="node-1", + embed_workflow=True, + extra_pnginfo={"workflow": workflow}, + save_with_metadata=False, + ) + + image_path = tmp_path / "sample_00001_.png" + with Image.open(image_path) as img: + assert "parameters" not in img.info + assert img.info["workflow"] == json.dumps(workflow) + + +def test_save_image_skips_jpeg_metadata_when_disabled(monkeypatch, tmp_path): + _configure_save_paths(monkeypatch, tmp_path) + _configure_metadata(monkeypatch, {"prompt": "prompt text", "seed": 123}) + + node = SaveImageLM() + node.save_images( + [_make_image()], + "ComfyUI", + "jpeg", + id="node-1", + save_with_metadata=False, + ) + + image_path = tmp_path / "sample_00001_.jpg" + exif_dict = piexif.load(str(image_path)) + assert piexif.ExifIFD.UserComment not in exif_dict.get("Exif", {}) + + +def test_save_image_skips_webp_metadata_when_disabled(monkeypatch, tmp_path): + _configure_save_paths(monkeypatch, tmp_path) + _configure_metadata(monkeypatch, {"prompt": "prompt text", "seed": 123}) + + node = SaveImageLM() + node.save_images( + [_make_image()], + "ComfyUI", + "webp", + id="node-1", + save_with_metadata=False, + ) + + image_path = tmp_path / "sample_00001_.webp" + exif_dict = piexif.load(str(image_path)) + assert piexif.ExifIFD.UserComment not in exif_dict.get("Exif", {})