From 5e162eee45eab9bc030a492b95f58005c654b7a7 Mon Sep 17 00:00:00 2001
From: Chris <little-god@web.de>
Date: Sun, 12 Jul 2026 11:32:43 +0200
Subject: [PATCH] feat: PNG metadata builder (parameters + optional
 workflow/prompt)

---
 nodes/civitai_metadata/png_info.py | 19 +++++++++++++++++++
 tests/test_png_info.py             | 24 ++++++++++++++++++++++++
 2 files changed, 43 insertions(+)
 create mode 100644 nodes/civitai_metadata/png_info.py
 create mode 100644 tests/test_png_info.py

diff --git a/nodes/civitai_metadata/png_info.py b/nodes/civitai_metadata/png_info.py
new file mode 100644
index 0000000..4c49100
--- /dev/null
+++ b/nodes/civitai_metadata/png_info.py
@@ -0,0 +1,19 @@
+# PNG text-chunk builder — part of ComfyUI-AI2Go-Utils. GPL-3.0.
+#
+# Always writes the A1111 `parameters` chunk (what Civitai reads). Optionally also writes ComfyUI's
+# native `prompt` + `workflow` chunks (when the node's save_workflow toggle is on) so the PNG stays
+# drag-and-drop reproducible in ComfyUI.
+import json
+
+from PIL.PngImagePlugin import PngInfo
+
+
+def build_pnginfo(parameters, *, extra_pnginfo=None, prompt=None):
+    info = PngInfo()
+    info.add_text("parameters", parameters)
+    if prompt is not None:
+        info.add_text("prompt", json.dumps(prompt))
+    if extra_pnginfo:
+        for key, value in extra_pnginfo.items():
+            info.add_text(key, json.dumps(value))
+    return info
diff --git a/tests/test_png_info.py b/tests/test_png_info.py
new file mode 100644
index 0000000..7d5bac6
--- /dev/null
+++ b/tests/test_png_info.py
@@ -0,0 +1,24 @@
+# Tests for the PNG metadata builder — part of ComfyUI-AI2Go-Utils. GPL-3.0.
+from PIL import Image
+from nodes.civitai_metadata.png_info import build_pnginfo
+
+
+def _roundtrip(info, tmp_path):
+    p = tmp_path / "img.png"
+    Image.new("RGB", (2, 2)).save(str(p), pnginfo=info)
+    return Image.open(str(p)).text
+
+
+def test_parameters_only(tmp_path):
+    text = _roundtrip(build_pnginfo("hello params"), tmp_path)
+    assert text["parameters"] == "hello params"
+    assert "workflow" not in text
+    assert "prompt" not in text
+
+
+def test_with_workflow_and_prompt(tmp_path):
+    info = build_pnginfo("p", extra_pnginfo={"workflow": {"nodes": []}}, prompt={"1": {"class_type": "X"}})
+    text = _roundtrip(info, tmp_path)
+    assert text["parameters"] == "p"
+    assert '"nodes"' in text["workflow"]
+    assert '"class_type"' in text["prompt"]