Skip to content

feat: embed workflow in generated PNG metadata#225

Open
TimPietruskyRunPod wants to merge 1 commit into
mainfrom
fix/enable-png-workflow-metadata
Open

feat: embed workflow in generated PNG metadata#225
TimPietruskyRunPod wants to merge 1 commit into
mainfrom
fix/enable-png-workflow-metadata

Conversation

@TimPietruskyRunPod
Copy link
Copy Markdown
Contributor

@TimPietruskyRunPod TimPietruskyRunPod commented Jun 2, 2026

Summary

Closes #139.

Generated PNGs currently cannot be dragged back into ComfyUI to recover the source workflow — because the worker was suppressing PNG metadata in two places:

  1. `src/start.sh` launched ComfyUI with `--disable-metadata`
  2. `handler.py` did not forward `extra_pnginfo.workflow` in the `/prompt` payload

`extra_pnginfo.workflow` is what the `Save Image` node reads to populate the `workflow` key embedded in PNG metadata (separate from the API-format `prompt`). Without both, the round-trip from output PNG → editor workflow doesn't work.

Bonus fix: crystools NoneType crash

This PR also fixes a separate bug discovered while investigating #167. Without `extra_pnginfo` in the payload, ComfyUI-Crystools' `Save image with extra metadata` node crashes with:

```
AttributeError: 'NoneType' object has no attribute 'copy'
```

at `nodes/image.py:392` (`extra_pnginfo_new = extra_pnginfo.copy()`) when `with_workflow=True`. Reproduced on a fresh pod today; PR fixes it.

Test plan

  • Reproduced the crystools NoneType crash on `main` (5.8.5 image)
  • After applying this PR's patch on a live pod, crystools `Save image with extra metadata` runs successfully with `with_workflow=True` and the handler returns a normal-sized image response
  • Validate end-to-end on a serverless endpoint built from this PR (via ci: build base image on every pull request #223 once merged) that the resulting PNG actually contains both `prompt` and `workflow` keys in its metadata

Trade-off

Each generated PNG is now ~10–50 KB larger (workflow JSON embedded). For users who want smaller responses, they can route through `PreviewImage` instead of `Save Image`, or strip metadata client-side. Worth documenting in a follow-up.

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
@TimPietruskyRunPod
Copy link
Copy Markdown
Contributor Author

Validation status update

Test Status Notes
Diff logic in isolation Verified the new `extra_data` initialization + `api_key_comfy_org` merge logic correctly produces both `extra_pnginfo.workflow` and `api_key_comfy_org` keys when applicable, and only `extra_pnginfo` when no key
--disable-metadata removal from `start.sh` Manually inspected, two occurrences (SERVE_API_LOCALLY=true branch and default branch) both removed
crystools NoneType crash fix Reproduced the crash on main's handler (5.8.5-flux1-dev-fp8 pod), applied this PR's diff to /handler.py in the running pod, restarted, re-submitted the workflow — completed successfully (585 KB response, no crash)
End-to-end PNG metadata embedding on a serverless endpoint ⚠️ pending Built image runpod/worker-comfyui:fix-enable-png-workflow-metadata-base published, deployed to a fresh serverless endpoint with a LoadImageSaveImage workflow. Worker came up healthy, picked up the job, then went unhealthy mid-execution. Could not retrieve worker logs (RunPod serverless doesn't expose them via API). The same image pattern worked when I tested via the pod manually, so likely an infra-level issue on that specific assigned host (same class as #186's original CUDA mismatch). Re-test once #223 merges and a real release goes out

The code change is correct and the crystools side-effect is verified. The final PNG-metadata roundtrip test (decode the returned PNG with PIL and confirm workflow chunk is present) hasn't been validated on a serverless endpoint yet — the test pod test wasn't a useful substitute because the pod image had --disable-metadata baked into its start.sh (only the handler.py was hot-patched).

I'd recommend merging this anyway and validating on the next release — the surface area is small, the diff logic checks out, and the worst case is "PNG metadata not embedded" which is the current state of main.

@TimPietruskyRunPod
Copy link
Copy Markdown
Contributor Author

Verified both changes are in the built image:

```
$ docker run --rm --platform linux/amd64 --entrypoint cat \
runpod/worker-comfyui:fix-enable-png-workflow-metadata-base /start.sh | grep main.py
python -u /comfyui/main.py --disable-auto-launch --listen --verbose "${COMFY_LOG_LEVEL}" --log-stdout &
python -u /comfyui/main.py --disable-auto-launch --verbose "${COMFY_LOG_LEVEL}" --log-stdout &
```

Both `python /comfyui/main.py` invocations are missing `--disable-metadata` ✅. And:

```
$ docker run --rm --platform linux/amd64 --entrypoint bash \
runpod/worker-comfyui:fix-enable-png-workflow-metadata-base \
-c "grep 'extra_pnginfo' /handler.py | head -2"
# extra_pnginfo.workflow is what ComfyUI embeds into generated PNGs so users
"extra_data": {"extra_pnginfo": {"workflow": workflow}},
```

The `extra_pnginfo` payload change is in handler.py ✅. So both halves of the fix are correctly baked. End-to-end PNG embedding test (decode returned PNG and confirm `workflow` chunk) still pending a live deploy that doesn't hit the host-side infra issue I ran into earlier today.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG]: Metadata not included in output image

2 participants