feat: embed workflow in generated PNG metadata

.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.

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

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"