Merge pull request #475 from nbsumbana-pixel/fix/comfyui-video-timeout-resume

fix: ComfyUI timeout/resume + websocket wait + multi-server + music tool
This commit is contained in:
Calesthio
2026-08-13 10:13:11 -07:00
committed by GitHub
9 changed files with 1530 additions and 74 deletions

View File

@@ -257,21 +257,36 @@ output_node: string # required for custom workflows
workflow_name: string # optional custom workflow provenance label
workflow_model: string # optional custom model/provenance label
workflow_model_stack: [] # optional custom dependency provenance
timeout_seconds: integer # optional, default 3600 (see below)
resume_prompt_id: string # optional, resume a timed-out job without resubmitting
```
**execute() flow (i2v):**
1. Upload reference image via `client.upload_image()`
2. Deep-copy i2v workflow template
3. Inject prompt, uploaded image name, seed, dimensions
4. `client.generate(workflow, output_node="108", dest=output_path, timeout=900)`
4. `client.generate(workflow, output_node="108", dest=output_path, timeout=inputs.get("timeout_seconds", 3600), resume_prompt_id=inputs.get("resume_prompt_id"))`
5. Return `ToolResult`
**execute() flow (t2v):**
1. Deep-copy t2v workflow template
2. Inject prompt, seed, dimensions
3. `client.generate(workflow, output_node="16", dest=output_path, timeout=900)`
3. `client.generate(workflow, output_node="16", dest=output_path, timeout=inputs.get("timeout_seconds", 3600), resume_prompt_id=inputs.get("resume_prompt_id"))`
4. Return `ToolResult`
**Timeout and resume (added after real-world local-GPU testing):** the
default client wait was raised from 900s to 3600s — non-accelerated custom
Wan 1.3B workflows on modest local GPUs were observed taking ~1360-1630s at
832x480/81-97 frames, and the old 900s default false-failed those jobs even
though ComfyUI kept rendering server-side. `ComfyUIError` now carries a
`prompt_id` on both execution errors and timeouts (`ComfyUIError.prompt_id`),
and `ComfyUIVideo`'s `ToolResult.error`/`.data` surface it on timeout so the
caller isn't left guessing whether the job is dead. Callers recover a
timed-out-but-still-running job by calling `execute()` again with
`resume_prompt_id` set to that `prompt_id` (and a longer `timeout_seconds` if
needed) — `client.generate()` then skips `submit()` entirely and just resumes
polling/downloading the existing job instead of queuing a duplicate.
`comfyui_video` publishes `operation_statuses` in `get_info()` and implements
`is_operation_available(operation)` for selector routing. This keeps partial
ComfyUI installs useful for the installed mode without advertising unavailable
@@ -281,19 +296,52 @@ not promote ComfyUI for an operation whose bundled models are missing.
---
### `comfyui_music` -- Music Generation (not shipped)
### `comfyui_music` -- Music Generation (shipped, with a native-node bundled workflow)
We explored adding a `comfyui_music` tool using the ACE-Step 3.5B model.
The model runs well in ComfyUI, but the ComfyUI node interface for
ACE-Step is not standardized -- there are multiple custom node packs with
different class names (`AceStepModelLoader` vs native `TextEncodeAceStepAudio`,
etc.). Shipping a workflow that only works with one specific custom node
pack would break for most users.
`tools/audio/comfyui_music.py`. `capability="music_generation"`, `provider="comfyui"`.
**Future path:** ACE-Step support should be revisited once OpenMontage decides
the music-generation routing shape and a portable ComfyUI audio workflow
contract. Current image/video workflow overrides are intentionally scoped to
image and video artifacts, not arbitrary audio workflows.
**Bundled default:** ACE-Step v1 (3.5B) text-to-audio, via `tools/_comfyui/workflows/ace-step-1-t2a.json`.
The node-pack fragmentation that originally blocked this tool (`AceStepModelLoader`
vs native `TextEncodeAceStepAudio`, etc.) turned out to be moot for ACE-Step v1:
ComfyUI ships `TextEncodeAceStepAudio`/`EmptyAceStepLatentAudio` as **native core
nodes** (`comfy_extras/nodes_ace.py`), not a third-party pack, and Comfy-Org's own
[`workflow_templates`](https://github.com/Comfy-Org/workflow_templates) repo bundles
an official ACE-Step-v1 template built entirely from those native nodes plus
long-stable core nodes (`CheckpointLoaderSimple`, `KSampler`, `ModelSamplingSD3`,
`VAEDecodeAudio`, `SaveAudioMP3`). Every node's `class_type` and input names in
`ace-step-1-t2a.json` were cross-checked against ComfyUI's own source
(`comfy_extras/nodes_ace.py`, `nodes_audio.py`, `nodes_latent.py`, `nodes.py`) --
not guessed from the UI export -- since the UI-format template Comfy-Org ships
isn't directly usable as the API-format JSON this client submits.
`prompt` maps to ACE-Step's `tags` field (style/genre/mood description, matching
the "prompt = description of desired music" convention `suno_music` already uses).
`lyrics` is a separate optional field (empty for instrumental). `duration_seconds`,
`steps`, `cfg`, `lyrics_strength`, and `seed` are all patchable; `shift` and the
tonemap `multiplier` stay at the official template's defaults.
Newer/different setups aren't locked out: `workflow_json`/`workflow_path` +
`output_node` still works exactly like the image/video tools' override path --
for ACE-Step 1.5, a different node pack, or a non-ACE-Step audio model entirely.
**Selector integration:** no dedicated `music_selector` exists in OpenMontage
(unlike `tts_selector`/`image_selector`/`video_selector`) -- music tools are
already routed directly via `registry.get_by_capability("music_generation")`,
and `comfyui_music` participates in that the same way `suno_music`/`music_gen`
do. `fallback_tools = ["suno_music", "music_gen"]`.
**Audio artifact schema:** `ToolResult.data` follows the same shape as the
image/video tools (`provider`, `model`, `output`, `format`, `workflow_provenance`),
plus `lyrics` and `duration_seconds` -- the latter a best-effort `ffprobe` probe
of the downloaded file (`None` if `ffprobe` isn't on PATH), since even the bundled
workflow doesn't report actual rendered duration back through `/history`.
**Workflow/output-node contract:** identical to image/video -- `output_node`
must be the ID of the node that writes the final artifact (the bundled workflow's
is `SaveAudioMP3`, ComfyUI's native audio saver). `ComfyUIClient.generate()`'s
artifact extraction now also checks the `"audio"` output key (previously only
`"images"`/`"gifs"`), which is what `SaveAudioMP3`/`SaveAudio` write to in
ComfyUI's `/history` response.
---
@@ -358,6 +406,19 @@ COMFYUI_POLL_TIMEOUT=600 # max wait for image gen
COMFYUI_VIDEO_TIMEOUT=900 # max wait for video gen
```
**Multi-server (optional):** point `comfyui_image`, `comfyui_video`, and
`comfyui_music` at separate ComfyUI instances -- e.g. one GPU running FLUX 2,
another running WAN 2.2, another running ACE-Step -- by setting a
per-capability override. Each takes priority over `COMFYUI_SERVER_URL` for
its own tool only; leave all three unset and everything talks to the single
shared server.
```bash
COMFYUI_IMAGE_SERVER_URL=http://gpu-a:8188
COMFYUI_VIDEO_SERVER_URL=http://gpu-b:8188
COMFYUI_MUSIC_SERVER_URL=http://gpu-c:8188
```
**For Docker Compose setups** (ComfyUI in a container):
```bash
@@ -459,15 +520,35 @@ pipeline definition, or any schema.
user-provided via a config directory? Bundling gives reproducibility;
external gives flexibility.
2. **Async generation:** ComfyUI supports websocket connections for real-time
progress. Worth implementing for long video generations, or is polling
sufficient?
2. ~~**Async generation:**~~ **Resolved.** `ComfyUIClient.generate()` now
waits via ComfyUI's websocket feed (`wait_ws()`) by default, reacting to
`executing`/`execution_error` events immediately instead of sleeping
between REST polls — completion and errors are caught without the
`interval`-seconds lag, and an optional `on_progress` callback gets live
`progress` events (`comfyui_video` uses this to print step progress on
long renders). No new hard dependency: `websocket-client` is an optional
import, and `_wait()` transparently falls back to the original
`poll()` REST loop when it isn't installed or the connection fails —
`resume_prompt_id` recovery behaves identically either way.
3. **Multi-server:** Should the adapter support multiple ComfyUI instances
(e.g., one for images, one for video) via per-capability URLs?
3. ~~**Multi-server:**~~ **Resolved.** `ComfyUIClient(capability="image"|"video"|"music")`
resolves its server URL from a per-capability env var first
(`COMFYUI_IMAGE_SERVER_URL` / `COMFYUI_VIDEO_SERVER_URL` / `COMFYUI_MUSIC_SERVER_URL`),
then the shared `COMFYUI_SERVER_URL`, then the `http://localhost:8188` default.
All three tools pass their capability at construction, so image, video, and
music generation can each point at different ComfyUI instances (different GPUs,
different model sets) with zero code changes -- single-server setups need no extra
configuration since all three env vars are optional. `client.capability`/
`client.is_default_url`/`client.unavailable_reason()` all account for the
override, and `COMFYUI_SETUP_OFFER.per_capability_env_var_overrides` documents
it for the setup-offer surfacing in `provider_menu()`.
4. **Music generation:** ACE-Step works in ComfyUI but OpenMontage needs a
dedicated music-generation routing contract before adding `comfyui_music`.
The follow-up should decide selector integration, audio artifact schemas, and
a portable workflow/output-node contract rather than treating music as a
hidden image/video workflow override.
4. ~~**Music generation:**~~ **Resolved -- shipped with a bundled ACE-Step v1 workflow.**
`comfyui_music` is a real tool now (not a hidden image/video override), routed
through the existing `registry.get_by_capability("music_generation")` path
like `suno_music`/`music_gen`. The node-pack fragmentation that originally
blocked this turned out not to apply to ACE-Step v1: its ComfyUI nodes are
native core nodes, not a third-party pack, so `ace-step-1-t2a.json` ships as
the default, verified node-by-node against ComfyUI's own source. Custom
`workflow_json`/`workflow_path` + `output_node` remains available for other
versions/packs. See the `comfyui_music` section above for the full contract.