fix(nodes): add save_with_metadata toggle to save image

This commit is contained in:
Will Miao
2026-03-28 11:17:36 +08:00
parent ad3bdddb72
commit 16e30ea689
2 changed files with 133 additions and 3 deletions

View File

@@ -72,6 +72,13 @@ class SaveImageLM:
"tooltip": "Embeds the complete workflow data into the image metadata. Only works with PNG and WebP formats.", "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": ( "add_counter_to_filename": (
"BOOLEAN", "BOOLEAN",
{ {
@@ -350,6 +357,7 @@ class SaveImageLM:
lossless_webp=True, lossless_webp=True,
quality=100, quality=100,
embed_workflow=False, embed_workflow=False,
save_with_metadata=True,
add_counter_to_filename=True, add_counter_to_filename=True,
): ):
"""Save images with metadata""" """Save images with metadata"""
@@ -421,7 +429,7 @@ class SaveImageLM:
try: try:
if file_format == "png": if file_format == "png":
assert pnginfo is not None assert pnginfo is not None
if metadata: if save_with_metadata and metadata:
pnginfo.add_text("parameters", metadata) pnginfo.add_text("parameters", metadata)
if embed_workflow and extra_pnginfo is not None: if embed_workflow and extra_pnginfo is not None:
workflow_json = json.dumps(extra_pnginfo["workflow"]) workflow_json = json.dumps(extra_pnginfo["workflow"])
@@ -430,7 +438,7 @@ class SaveImageLM:
img.save(file_path, format="PNG", **save_kwargs) img.save(file_path, format="PNG", **save_kwargs)
elif file_format == "jpeg": elif file_format == "jpeg":
# For JPEG, use piexif # For JPEG, use piexif
if metadata: if save_with_metadata and metadata:
try: try:
exif_dict = { exif_dict = {
"Exif": { "Exif": {
@@ -448,7 +456,7 @@ class SaveImageLM:
# For WebP, use piexif for metadata # For WebP, use piexif for metadata
exif_dict = {} exif_dict = {}
if metadata: if save_with_metadata and metadata:
exif_dict["Exif"] = { exif_dict["Exif"] = {
piexif.ExifIFD.UserComment: b"UNICODE\0" piexif.ExifIFD.UserComment: b"UNICODE\0"
+ metadata.encode("utf-16be") + metadata.encode("utf-16be")
@@ -489,6 +497,7 @@ class SaveImageLM:
lossless_webp=True, lossless_webp=True,
quality=100, quality=100,
embed_workflow=False, embed_workflow=False,
save_with_metadata=True,
add_counter_to_filename=True, add_counter_to_filename=True,
): ):
"""Process and save image with metadata""" """Process and save image with metadata"""
@@ -516,6 +525,7 @@ class SaveImageLM:
lossless_webp, lossless_webp,
quality, quality,
embed_workflow, embed_workflow,
save_with_metadata,
add_counter_to_filename, add_counter_to_filename,
) )

View File

@@ -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", {})