Preserve extra PNG metadata in AVIF XMP

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"      <comfy:workflow>{escape(json.dumps(workflow))}</comfy:workflow>"
 
+    extra_pnginfo_element = ""
+    if extra_pnginfo:
+        extra_pnginfo_element = (
+            "      <comfy:extra_pnginfo>"
+            f"{escape(json.dumps(extra_pnginfo))}"
+            "</comfy:extra_pnginfo>"
+        )
+
     packet = [
         '<?xpacket begin="\ufeff" id="W5M0MpCehiHzreSzNTczkc9d"?>',
         '<x:xmpmeta xmlns:x="adobe:ns:meta/">',
@@ -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([
         "    </rdf:Description>",
         "  </rdf:RDF>",

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": "<tag> & value"}}}
     workflow = {"nodes": [{"id": 1, "title": 'Save "AVIF"'}], "version": 1}
+    extra_pnginfo = {"workflow": workflow, "custom metadata": {"value": "<one> & 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
                 )