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"