diff --git a/comfy_extras/nodes_images.py b/comfy_extras/nodes_images.py
index a47de69f9..a80c77587 100644
--- a/comfy_extras/nodes_images.py
+++ b/comfy_extras/nodes_images.py
@@ -947,9 +947,9 @@ def inject_png_metadata(png_bytes: bytes, prompt: dict | None, extra_pnginfo: di
def build_avif_xmp(prompt: dict | None, extra_pnginfo: dict | None) -> bytes | None:
- """Build the standard XMP item used to import ComfyUI AVIF workflows."""
+ """Build the XMP item used to preserve ComfyUI AVIF metadata."""
workflow = extra_pnginfo.get("workflow") if extra_pnginfo else None
- if prompt is None and workflow is None:
+ if prompt is None and not extra_pnginfo:
return None
prompt_attribute = ""
@@ -960,6 +960,14 @@ def build_avif_xmp(prompt: dict | None, extra_pnginfo: dict | None) -> bytes | N
if workflow is not None:
workflow_element = f" {escape(json.dumps(workflow))}"
+ extra_pnginfo_element = ""
+ if extra_pnginfo:
+ extra_pnginfo_element = (
+ " "
+ f"{escape(json.dumps(extra_pnginfo))}"
+ ""
+ )
+
packet = [
'',
'',
@@ -968,6 +976,8 @@ def build_avif_xmp(prompt: dict | None, extra_pnginfo: dict | None) -> bytes | N
]
if workflow_element:
packet.append(workflow_element)
+ if extra_pnginfo_element:
+ packet.append(extra_pnginfo_element)
packet.extend([
" ",
" ",
diff --git a/tests-unit/comfy_extras_test/nodes_images_save_test.py b/tests-unit/comfy_extras_test/nodes_images_save_test.py
index eadc4aeb8..3df02d9c7 100644
--- a/tests-unit/comfy_extras_test/nodes_images_save_test.py
+++ b/tests-unit/comfy_extras_test/nodes_images_save_test.py
@@ -75,26 +75,43 @@ def test_save_image_advanced_schema_exposes_avif_quality():
assert serialized_options["avif"]["required"]["quality"][:1] == ("INT",)
-def test_build_avif_xmp_escapes_and_preserves_prompt_and_workflow():
+def test_build_avif_xmp_escapes_and_preserves_comfyui_metadata():
prompt = {"1": {"class_type": "KSampler", "inputs": {"text": " & value"}}}
workflow = {"nodes": [{"id": 1, "title": 'Save "AVIF"'}], "version": 1}
+ extra_pnginfo = {"workflow": workflow, "custom metadata": {"value": " & two"}}
- xmp = nodes_images.build_avif_xmp(
- prompt, {"workflow": workflow, "ignored": {"value": 1}}
- )
+ xmp = nodes_images.build_avif_xmp(prompt, extra_pnginfo)
assert xmp is not None
root = ElementTree.fromstring(xmp)
description = root.find(f".//{{{RDF_NAMESPACE}}}Description")
workflow_element = root.find(f".//{{{COMFYUI_XMP_NAMESPACE}}}workflow")
+ extra_pnginfo_element = root.find(
+ f".//{{{COMFYUI_XMP_NAMESPACE}}}extra_pnginfo"
+ )
assert description is not None
assert workflow_element is not None
+ assert extra_pnginfo_element is not None
assert description.attrib[f"{{{COMFYUI_XMP_NAMESPACE}}}prompt"] == json.dumps(
prompt
)
assert workflow_element.text == json.dumps(workflow)
- assert b"ignored" not in xmp
+ assert json.loads(extra_pnginfo_element.text) == extra_pnginfo
+
+
+def test_build_avif_xmp_preserves_extra_pnginfo_without_workflow():
+ extra_pnginfo = {"custom": {"value": 1}}
+
+ xmp = nodes_images.build_avif_xmp(None, extra_pnginfo)
+
+ assert xmp is not None
+ root = ElementTree.fromstring(xmp)
+ extra_pnginfo_element = root.find(
+ f".//{{{COMFYUI_XMP_NAMESPACE}}}extra_pnginfo"
+ )
+ assert extra_pnginfo_element is not None
+ assert json.loads(extra_pnginfo_element.text) == extra_pnginfo
def test_encode_avif_image_passes_fixed_encoding_options(monkeypatch):
@@ -170,6 +187,7 @@ def test_save_image_advanced_writes_avif_batch_and_respects_metadata_setting(
):
prompt = {"1": {"class_type": "KSampler"}}
workflow = {"nodes": [], "version": 1}
+ extra_pnginfo = {"workflow": workflow, "custom": {"value": 1}}
images = torch.from_numpy(
np.random.default_rng(7).random((2, 12, 16, 3), dtype=np.float32)
)
@@ -186,7 +204,7 @@ def test_save_image_advanced_writes_avif_batch_and_respects_metadata_setting(
monkeypatch.setattr(
nodes_images.SaveImageAdvanced,
"hidden",
- SimpleNamespace(prompt=prompt, extra_pnginfo={"workflow": workflow}),
+ SimpleNamespace(prompt=prompt, extra_pnginfo=extra_pnginfo),
)
output = nodes_images.SaveImageAdvanced.execute(
@@ -206,5 +224,5 @@ def test_save_image_advanced_writes_avif_batch_and_respects_metadata_setting(
assert "xmp" not in saved_image.info
else:
assert saved_image.info["xmp"] == nodes_images.build_avif_xmp(
- prompt, {"workflow": workflow}
+ prompt, extra_pnginfo
)