From dc2cc6140731f57e9afb7afe8a4c07b5c858a52e Mon Sep 17 00:00:00 2001
From: Tim Pietrusky <timpietrusky@gmail.com>
Date: Tue, 2 Jun 2026 13:35:09 +0200
Subject: [PATCH] feat: embed workflow in generated PNG metadata

Drop --disable-metadata from the comfyui launch and forward
extra_pnginfo.workflow alongside the prompt in the /prompt payload.
This is what ComfyUI's Save Image node looks at to embed the
"workflow" key in PNG metadata. Without both, generated PNGs cannot
be dragged back into the editor to recover the source workflow.

Refs #139
---
 .changeset/enable-png-workflow-metadata.md | 14 ++++++++++++++
 handler.py                                 | 14 +++++++++++---
 src/start.sh                               |  4 ++--
 3 files changed, 27 insertions(+), 5 deletions(-)
 create mode 100644 .changeset/enable-png-workflow-metadata.md

diff --git a/.changeset/enable-png-workflow-metadata.md b/.changeset/enable-png-workflow-metadata.md
new file mode 100644
index 000000000..fc6ca2557
--- /dev/null
+++ b/.changeset/enable-png-workflow-metadata.md
@@ -0,0 +1,14 @@
+---
+"worker-comfyui": minor
+---
+
+feat: embed the workflow in generated PNG metadata so output images can be loaded back into ComfyUI
+
+Before this change, ComfyUI was launched with `--disable-metadata` and the handler did not forward `extra_pnginfo` to the `/prompt` endpoint. As a result, generated PNGs did not contain the workflow JSON, and users couldn't drag a result back into the editor to recover the workflow that produced it (#139).
+
+Two changes are required for round-tripping to work:
+
+1. `src/start.sh`: drop `--disable-metadata` from the two `python /comfyui/main.py` invocations.
+2. `handler.py`: pass `extra_data: {"extra_pnginfo": {"workflow": workflow}}` in the `/prompt` payload alongside the existing API-format `prompt`. The `extra_pnginfo.workflow` is what the `Save Image` node embeds as the `workflow` key in the PNG.
+
+Side effect: each generated PNG is now ~10-50 KB larger than before (depends on workflow complexity). For users who prefer minimal images, this can be reversed per-job by submitting a payload that uses a workflow without `Save Image` (e.g. only `PreviewImage`), or by post-processing strip on the client.
diff --git a/handler.py b/handler.py
index 817cad8a9..9d5d6b557 100644
--- a/handler.py
+++ b/handler.py
@@ -417,8 +417,16 @@ def queue_workflow(workflow, client_id, comfy_org_api_key=None):
     Raises:
         ValueError: If the workflow validation fails with detailed error information
     """
-    # Include client_id in the prompt payload
-    payload = {"prompt": workflow, "client_id": client_id}
+    # Include client_id in the prompt payload.
+    # extra_pnginfo.workflow is what ComfyUI embeds into generated PNGs so users
+    # can drag a result back into the editor and recover the workflow. Without
+    # it, even with metadata enabled, the PNG only carries the API-format
+    # `prompt` and is not directly loadable in the UI.
+    payload = {
+        "prompt": workflow,
+        "client_id": client_id,
+        "extra_data": {"extra_pnginfo": {"workflow": workflow}},
+    }
 
     # Optionally inject Comfy.org API key for API Nodes.
     # Precedence: per-request key (argument) overrides environment variable.
@@ -427,7 +435,7 @@ def queue_workflow(workflow, client_id, comfy_org_api_key=None):
     key_from_env = os.environ.get("COMFY_ORG_API_KEY")
     effective_key = comfy_org_api_key if comfy_org_api_key else key_from_env
     if effective_key:
-        payload["extra_data"] = {"api_key_comfy_org": effective_key}
+        payload["extra_data"]["api_key_comfy_org"] = effective_key
     data = json.dumps(payload).encode("utf-8")
 
     # Use requests for consistency and timeout
diff --git a/src/start.sh b/src/start.sh
index c490b0db1..f995a5cbf 100644
--- a/src/start.sh
+++ b/src/start.sh
@@ -60,13 +60,13 @@ COMFY_PID_FILE="/tmp/comfyui.pid"
 
 # Serve the API and don't shutdown the container
 if [ "$SERVE_API_LOCALLY" == "true" ]; then
-    python -u /comfyui/main.py --disable-auto-launch --disable-metadata --listen --verbose "${COMFY_LOG_LEVEL}" --log-stdout &
+    python -u /comfyui/main.py --disable-auto-launch --listen --verbose "${COMFY_LOG_LEVEL}" --log-stdout &
     echo $! > "$COMFY_PID_FILE"
 
     echo "worker-comfyui: Starting RunPod Handler"
     python -u /handler.py --rp_serve_api --rp_api_host=0.0.0.0
 else
-    python -u /comfyui/main.py --disable-auto-launch --disable-metadata --verbose "${COMFY_LOG_LEVEL}" --log-stdout &
+    python -u /comfyui/main.py --disable-auto-launch --verbose "${COMFY_LOG_LEVEL}" --log-stdout &
     echo $! > "$COMFY_PID_FILE"
 
     echo "worker-comfyui: Starting RunPod Handler"