mirror of
https://github.com/calesthio/OpenMontage.git
synced 2026-08-24 09:00:20 +08:00
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
@@ -1,17 +1,19 @@
|
||||
---
|
||||
name: comfyui
|
||||
description: Use when working with ComfyUI workflows in OpenMontage, including comfyui_image/comfyui_video, custom workflow_json/workflow_path inputs, output_node selection, missing model setup, LoRAs, low-VRAM workflow choices, and community workflow imports.
|
||||
description: Use when working with ComfyUI workflows in OpenMontage, including comfyui_image/comfyui_video/comfyui_music, custom workflow_json/workflow_path inputs, output_node selection, missing model setup, LoRAs, low-VRAM workflow choices, and community workflow imports.
|
||||
---
|
||||
|
||||
# ComfyUI Workflows in OpenMontage
|
||||
|
||||
Use this skill before calling `comfyui_image` or `comfyui_video`, and when converting a community ComfyUI workflow into an OpenMontage tool call.
|
||||
Use this skill before calling `comfyui_image`, `comfyui_video`, or `comfyui_music`, and when converting a community ComfyUI workflow into an OpenMontage tool call.
|
||||
|
||||
## Server Contract
|
||||
|
||||
- ComfyUI must be running before the tool can generate. The default server is `http://localhost:8188`; override it with `COMFYUI_SERVER_URL`.
|
||||
- Running separate ComfyUI instances per capability (different GPU, different model set)? `COMFYUI_IMAGE_SERVER_URL` / `COMFYUI_VIDEO_SERVER_URL` / `COMFYUI_MUSIC_SERVER_URL` each override `COMFYUI_SERVER_URL` for that one tool only. Optional -- a single-server setup needs none of these.
|
||||
- Health and hardware status come from `GET /system_stats`.
|
||||
- Jobs are submitted to `POST /prompt`, completed outputs are read from `GET /history/{prompt_id}`, and artifact bytes are downloaded with `GET /view`.
|
||||
- Long waits (video, music) prefer ComfyUI's websocket feed for immediate completion/error detection and transparently fall back to REST polling if `websocket-client` isn't installed. Either way, a timeout is recoverable: pass the error's `prompt_id` back in as `resume_prompt_id` to resume waiting on the same job instead of resubmitting it.
|
||||
- Export workflows with ComfyUI's API-format JSON, not the UI layout format. If a downloaded workflow will not submit, re-export it from ComfyUI with API format enabled.
|
||||
|
||||
## Choosing a Workflow
|
||||
@@ -52,4 +54,13 @@ Use this skill before calling `comfyui_image` or `comfyui_video`, and when conve
|
||||
- If the server is unavailable, surface the structured setup offer. Starting ComfyUI or setting `COMFYUI_SERVER_URL` is the first fix.
|
||||
- If models are missing, read `data.missing_models[]`; each item should include the file name, role, destination hint, and download URL when OpenMontage knows it.
|
||||
- If custom nodes are missing, ask the user to install them through ComfyUI Manager or the workflow author's documented install path, then restart ComfyUI.
|
||||
- If a long render times out locally, check ComfyUI history before retrying from scratch; the server may still have completed the prompt.
|
||||
- If a long render times out locally, check ComfyUI history before retrying from scratch; the server may still have completed the prompt -- or just call again with `resume_prompt_id` set to the `prompt_id` from the timeout error.
|
||||
|
||||
## Music (`comfyui_music`)
|
||||
|
||||
- Bundled default is ACE-Step v1 (3.5B) text-to-audio, built from ComfyUI's *native* `TextEncodeAceStepAudio`/`EmptyAceStepLatentAudio` nodes (core, not a third-party pack) -- unlike ACE-Step 1.5 or other custom node packs, v1's interface is standardized enough to bundle safely.
|
||||
- `prompt` maps to the bundled workflow's `tags` field (style/genre/mood, e.g. `"upbeat electronic pop, female vocals"`), matching the same "prompt = music description" convention `suno_music` uses. `lyrics` is a separate optional field -- leave empty for instrumental, or use `[verse]`/`[chorus]`/`[bridge]` structure tags and `[zh]`/`[ja]`/`[ko]`-style language-code prefixes for non-English lines.
|
||||
- `duration_seconds`, `steps`, `cfg`, `lyrics_strength`, and `seed` are patchable on the bundled workflow. Missing `ace_step_v1_3.5b.safetensors` surfaces through the same `data.missing_models[]` contract as image/video.
|
||||
- Need ACE-Step 1.5, a different node pack, or a non-ACE-Step audio model? Fall back to `workflow_json`/`workflow_path` + `output_node`, exactly like a custom image/video workflow -- in that mode `prompt` becomes provenance/logging only again and must already be baked into the graph.
|
||||
- `output_node` (bundled or custom) should be the node that writes the final audio -- the bundled workflow's is `SaveAudioMP3`. The client reads artifacts from that node's `"audio"` output key (parallel to `"images"` for image/video savers).
|
||||
- For custom workflows, provide `workflow_name`/`workflow_model`/`workflow_model_stack` for provenance exactly as you would for a custom image/video workflow.
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -17,13 +17,14 @@ from tools.base_tool import (
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
from tools.audio.comfyui_music import ComfyUIMusic
|
||||
from tools.graphics.comfyui_image import ComfyUIImage
|
||||
from tools.graphics.image_selector import ImageSelector
|
||||
from tools.tool_registry import ToolRegistry
|
||||
from tools.video.video_selector import VideoSelector
|
||||
from tools.video.comfyui_video import ComfyUIVideo
|
||||
|
||||
TOOLS = [ComfyUIImage, ComfyUIVideo]
|
||||
TOOLS = [ComfyUIImage, ComfyUIVideo, ComfyUIMusic]
|
||||
WORKFLOW_DIR = Path(__file__).resolve().parent.parent.parent / "tools" / "_comfyui" / "workflows"
|
||||
PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent
|
||||
|
||||
@@ -136,6 +137,7 @@ EXPECTED_WORKFLOWS = [
|
||||
"flux2-txt2img.json",
|
||||
"wan22-i2v-4step.json",
|
||||
"wan22-t2v-4step.json",
|
||||
"ace-step-1-t2a.json",
|
||||
]
|
||||
|
||||
|
||||
@@ -282,6 +284,73 @@ class TestClientHelpers:
|
||||
"folder_type": "temp",
|
||||
}
|
||||
|
||||
def test_poll_timeout_carries_prompt_id_for_recovery(self, monkeypatch):
|
||||
from tools._comfyui.client import ComfyUIClient, ComfyUIError
|
||||
|
||||
client = ComfyUIClient("http://comfy.test")
|
||||
monkeypatch.setattr(
|
||||
"tools._comfyui.client.requests.get",
|
||||
lambda *a, **k: type("R", (), {
|
||||
"raise_for_status": lambda self: None,
|
||||
"json": lambda self: {},
|
||||
})(),
|
||||
)
|
||||
monkeypatch.setattr("tools._comfyui.client.time.sleep", lambda s: None)
|
||||
|
||||
with pytest.raises(ComfyUIError) as excinfo:
|
||||
client.poll("prompt-timeout-1", timeout=0, interval=0)
|
||||
|
||||
assert excinfo.value.prompt_id == "prompt-timeout-1"
|
||||
assert "prompt-timeout-1" in str(excinfo.value)
|
||||
|
||||
def test_generate_resume_prompt_id_skips_resubmit(self, monkeypatch, tmp_path):
|
||||
from tools._comfyui.client import ComfyUIClient
|
||||
import sys
|
||||
|
||||
client = ComfyUIClient("http://comfy.test")
|
||||
|
||||
def fail_submit(workflow):
|
||||
raise AssertionError("submit() should not be called when resuming")
|
||||
|
||||
monkeypatch.setattr(client, "submit", fail_submit)
|
||||
_install_fake_websocket(monkeypatch, frames=[])
|
||||
monkeypatch.setattr(
|
||||
sys.modules["websocket"],
|
||||
"create_connection",
|
||||
lambda *a, **k: (_ for _ in ()).throw(
|
||||
AssertionError("resumed jobs must use history polling")
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(client, "poll", lambda prompt_id, **kwargs: {
|
||||
"outputs": {"9": {"images": [{
|
||||
"filename": "resumed.png", "subfolder": "", "type": "output",
|
||||
}]}}
|
||||
})
|
||||
monkeypatch.setattr(client, "download", lambda filename, subfolder, dest, folder_type="output": Path(dest))
|
||||
|
||||
paths = client.generate(
|
||||
{"9": {"inputs": {}}}, "9", tmp_path / "out.png",
|
||||
resume_prompt_id="already-running-id",
|
||||
)
|
||||
assert paths == [tmp_path / "out.png"]
|
||||
|
||||
def test_generate_reads_audio_key_from_savaudio_node(self, monkeypatch, tmp_path):
|
||||
"""The native SaveAudio node writes outputs under "audio", not
|
||||
"images"/"gifs" -- comfyui_music depends on this being handled."""
|
||||
from tools._comfyui.client import ComfyUIClient
|
||||
|
||||
client = ComfyUIClient("http://comfy.test")
|
||||
monkeypatch.setattr(client, "submit", lambda workflow: "p1")
|
||||
monkeypatch.setattr(client, "poll", lambda prompt_id, **kwargs: {
|
||||
"outputs": {"9": {"audio": [{
|
||||
"filename": "track.flac", "subfolder": "", "type": "output",
|
||||
}]}}
|
||||
})
|
||||
monkeypatch.setattr(client, "download", lambda filename, subfolder, dest, folder_type="output": Path(dest))
|
||||
|
||||
paths = client.generate({"9": {"inputs": {}}}, "9", tmp_path / "out.flac")
|
||||
assert paths == [tmp_path / "out.flac"]
|
||||
|
||||
def test_is_default_url_when_env_not_set(self, monkeypatch):
|
||||
from tools._comfyui.client import ComfyUIClient
|
||||
monkeypatch.delenv("COMFYUI_SERVER_URL", raising=False)
|
||||
@@ -310,6 +379,290 @@ class TestClientHelpers:
|
||||
assert "myhost:9999" in msg
|
||||
assert "COMFYUI_SERVER_URL" not in msg
|
||||
|
||||
def test_submit_includes_client_id_for_websocket_targeting(self, monkeypatch):
|
||||
from tools._comfyui.client import ComfyUIClient
|
||||
|
||||
client = ComfyUIClient("http://comfy.test")
|
||||
seen = {}
|
||||
|
||||
def fake_post(url, json=None, timeout=None):
|
||||
seen.update(json)
|
||||
return type("R", (), {
|
||||
"raise_for_status": lambda self: None,
|
||||
"json": lambda self: {"prompt_id": "abc"},
|
||||
})()
|
||||
|
||||
monkeypatch.setattr("tools._comfyui.client.requests.post", fake_post)
|
||||
client.submit({"1": {"inputs": {}}})
|
||||
assert seen["client_id"] == client.client_id
|
||||
|
||||
|
||||
class TestMultiServer:
|
||||
|
||||
def test_capability_env_var_takes_priority_over_shared(self, monkeypatch):
|
||||
from tools._comfyui.client import ComfyUIClient
|
||||
|
||||
monkeypatch.setenv("COMFYUI_SERVER_URL", "http://shared:8188")
|
||||
monkeypatch.setenv("COMFYUI_VIDEO_SERVER_URL", "http://video-gpu:8188")
|
||||
client = ComfyUIClient(capability="video")
|
||||
assert client.server_url == "http://video-gpu:8188"
|
||||
|
||||
def test_falls_back_to_shared_when_capability_var_unset(self, monkeypatch):
|
||||
from tools._comfyui.client import ComfyUIClient
|
||||
|
||||
monkeypatch.setenv("COMFYUI_SERVER_URL", "http://shared:8188")
|
||||
monkeypatch.delenv("COMFYUI_IMAGE_SERVER_URL", raising=False)
|
||||
client = ComfyUIClient(capability="image")
|
||||
assert client.server_url == "http://shared:8188"
|
||||
|
||||
def test_other_capability_env_var_does_not_leak_across_tools(self, monkeypatch):
|
||||
from tools._comfyui.client import ComfyUIClient
|
||||
|
||||
monkeypatch.delenv("COMFYUI_SERVER_URL", raising=False)
|
||||
monkeypatch.setenv("COMFYUI_IMAGE_SERVER_URL", "http://image-gpu:8188")
|
||||
monkeypatch.delenv("COMFYUI_VIDEO_SERVER_URL", raising=False)
|
||||
video_client = ComfyUIClient(capability="video")
|
||||
assert video_client.server_url == "http://localhost:8188"
|
||||
|
||||
def test_explicit_server_url_wins_over_capability_env_var(self, monkeypatch):
|
||||
from tools._comfyui.client import ComfyUIClient
|
||||
|
||||
monkeypatch.setenv("COMFYUI_VIDEO_SERVER_URL", "http://video-gpu:8188")
|
||||
client = ComfyUIClient("http://explicit:1234", capability="video")
|
||||
assert client.server_url == "http://explicit:1234"
|
||||
|
||||
def test_no_capability_behaves_as_before(self, monkeypatch):
|
||||
from tools._comfyui.client import ComfyUIClient
|
||||
|
||||
monkeypatch.setenv("COMFYUI_SERVER_URL", "http://shared:8188")
|
||||
client = ComfyUIClient()
|
||||
assert client.server_url == "http://shared:8188"
|
||||
assert client.is_default_url is False
|
||||
|
||||
def test_is_default_url_true_only_when_both_vars_unset(self, monkeypatch):
|
||||
from tools._comfyui.client import ComfyUIClient
|
||||
|
||||
monkeypatch.delenv("COMFYUI_SERVER_URL", raising=False)
|
||||
monkeypatch.delenv("COMFYUI_IMAGE_SERVER_URL", raising=False)
|
||||
client = ComfyUIClient(capability="image")
|
||||
assert client.is_default_url is True
|
||||
|
||||
monkeypatch.setenv("COMFYUI_IMAGE_SERVER_URL", "http://image-gpu:8188")
|
||||
client2 = ComfyUIClient(capability="image")
|
||||
assert client2.is_default_url is False
|
||||
|
||||
def test_unavailable_reason_mentions_capability_and_shared_var(self, monkeypatch):
|
||||
from tools._comfyui.client import ComfyUIClient
|
||||
|
||||
monkeypatch.delenv("COMFYUI_SERVER_URL", raising=False)
|
||||
monkeypatch.delenv("COMFYUI_VIDEO_SERVER_URL", raising=False)
|
||||
client = ComfyUIClient(capability="video")
|
||||
msg = client.unavailable_reason()
|
||||
assert "COMFYUI_VIDEO_SERVER_URL" in msg
|
||||
assert "COMFYUI_SERVER_URL" in msg
|
||||
|
||||
def test_image_and_video_tools_use_independent_servers(self, monkeypatch):
|
||||
from tools.graphics.comfyui_image import ComfyUIImage
|
||||
from tools.video.comfyui_video import ComfyUIVideo
|
||||
|
||||
monkeypatch.delenv("COMFYUI_SERVER_URL", raising=False)
|
||||
monkeypatch.setenv("COMFYUI_IMAGE_SERVER_URL", "http://image-gpu:8188")
|
||||
monkeypatch.setenv("COMFYUI_VIDEO_SERVER_URL", "http://video-gpu:8188")
|
||||
|
||||
image_tool = ComfyUIImage()
|
||||
video_tool = ComfyUIVideo()
|
||||
|
||||
assert image_tool._client.server_url == "http://image-gpu:8188"
|
||||
assert video_tool._client.server_url == "http://video-gpu:8188"
|
||||
|
||||
|
||||
class _FakeWSTimeout(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class _FakeWSConn:
|
||||
def __init__(self, frames):
|
||||
self._frames = list(frames)
|
||||
|
||||
def settimeout(self, value):
|
||||
pass
|
||||
|
||||
def recv(self):
|
||||
if not self._frames:
|
||||
raise _FakeWSTimeout()
|
||||
return self._frames.pop(0)
|
||||
|
||||
def close(self):
|
||||
pass
|
||||
|
||||
|
||||
def _install_fake_websocket(monkeypatch, frames):
|
||||
"""Inject a fake `websocket` module so wait_ws() runs without the real
|
||||
optional websocket-client dependency installed."""
|
||||
import sys
|
||||
import types
|
||||
|
||||
fake_module = types.SimpleNamespace(
|
||||
WebSocketTimeoutException=_FakeWSTimeout,
|
||||
create_connection=lambda url, timeout=10: _FakeWSConn(frames),
|
||||
)
|
||||
monkeypatch.setitem(sys.modules, "websocket", fake_module)
|
||||
|
||||
|
||||
class TestWebsocketWait:
|
||||
|
||||
def test_wait_ws_returns_job_completed_before_connection(self, monkeypatch):
|
||||
from tools._comfyui.client import ComfyUIClient
|
||||
import sys
|
||||
|
||||
client = ComfyUIClient("http://comfy.test")
|
||||
_install_fake_websocket(monkeypatch, frames=[])
|
||||
monkeypatch.setattr(
|
||||
sys.modules["websocket"],
|
||||
"create_connection",
|
||||
lambda *a, **k: (_ for _ in ()).throw(
|
||||
AssertionError("completed history must avoid websocket connection")
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"tools._comfyui.client.requests.get",
|
||||
lambda *a, **k: type("R", (), {
|
||||
"raise_for_status": lambda self: None,
|
||||
"json": lambda self: {"done": {"outputs": {"9": {}}}},
|
||||
})(),
|
||||
)
|
||||
|
||||
assert client.wait_ws("done", timeout=5) == {"outputs": {"9": {}}}
|
||||
|
||||
def test_wait_ws_completes_on_executing_none_node(self, monkeypatch, tmp_path):
|
||||
from tools._comfyui.client import ComfyUIClient
|
||||
|
||||
client = ComfyUIClient("http://comfy.test")
|
||||
progress_events = []
|
||||
frames = [
|
||||
json.dumps({"type": "progress", "data": {
|
||||
"value": 2, "max": 20, "prompt_id": "p1",
|
||||
}}),
|
||||
json.dumps({"type": "executing", "data": {
|
||||
"node": None, "prompt_id": "p1",
|
||||
}}),
|
||||
]
|
||||
_install_fake_websocket(monkeypatch, frames)
|
||||
history_calls = iter(({}, {}, {"p1": {"outputs": {"9": {}}}}))
|
||||
monkeypatch.setattr(
|
||||
"tools._comfyui.client.requests.get",
|
||||
lambda *a, **k: type("R", (), {
|
||||
"raise_for_status": lambda self: None,
|
||||
"json": lambda self: next(history_calls),
|
||||
})(),
|
||||
)
|
||||
|
||||
entry = client.wait_ws("p1", timeout=5, on_progress=progress_events.append)
|
||||
|
||||
assert entry == {"outputs": {"9": {}}}
|
||||
assert progress_events == [{"value": 2, "max": 20, "prompt_id": "p1"}]
|
||||
|
||||
def test_wait_ws_execution_error_raises_with_prompt_id(self, monkeypatch):
|
||||
from tools._comfyui.client import ComfyUIClient, ComfyUIError
|
||||
|
||||
client = ComfyUIClient("http://comfy.test")
|
||||
frames = [
|
||||
json.dumps({"type": "execution_error", "data": {
|
||||
"prompt_id": "p2", "exception_message": "boom",
|
||||
}}),
|
||||
]
|
||||
_install_fake_websocket(monkeypatch, frames)
|
||||
|
||||
with pytest.raises(ComfyUIError) as excinfo:
|
||||
client.wait_ws("p2", timeout=5)
|
||||
|
||||
assert excinfo.value.prompt_id == "p2"
|
||||
|
||||
def test_wait_ws_ignores_other_prompts_on_shared_connection(self, monkeypatch):
|
||||
from tools._comfyui.client import ComfyUIClient
|
||||
|
||||
client = ComfyUIClient("http://comfy.test")
|
||||
frames = [
|
||||
# Another job's event on the same client_id -- must not trigger completion.
|
||||
json.dumps({"type": "executing", "data": {
|
||||
"node": None, "prompt_id": "someone-elses-job",
|
||||
}}),
|
||||
json.dumps({"type": "executing", "data": {
|
||||
"node": None, "prompt_id": "p3",
|
||||
}}),
|
||||
]
|
||||
_install_fake_websocket(monkeypatch, frames)
|
||||
monkeypatch.setattr(
|
||||
"tools._comfyui.client.requests.get",
|
||||
lambda *a, **k: type("R", (), {
|
||||
"raise_for_status": lambda self: None,
|
||||
"json": lambda self: {"p3": {"outputs": {}}},
|
||||
})(),
|
||||
)
|
||||
|
||||
entry = client.wait_ws("p3", timeout=5)
|
||||
assert entry == {"outputs": {}}
|
||||
|
||||
def test_wait_ws_timeout_raises_comfyuierror_with_prompt_id(self, monkeypatch):
|
||||
from tools._comfyui.client import ComfyUIClient, ComfyUIError
|
||||
|
||||
client = ComfyUIClient("http://comfy.test")
|
||||
_install_fake_websocket(monkeypatch, frames=[]) # recv() always times out
|
||||
|
||||
with pytest.raises(ComfyUIError) as excinfo:
|
||||
client.wait_ws("p4", timeout=0)
|
||||
|
||||
assert excinfo.value.prompt_id == "p4"
|
||||
|
||||
def test_wait_falls_back_to_poll_when_websocket_unavailable(self, monkeypatch):
|
||||
"""No websocket-client installed (or any transport failure) must
|
||||
silently fall back to REST polling, not blow up the whole call."""
|
||||
from tools._comfyui.client import ComfyUIClient
|
||||
import sys
|
||||
|
||||
client = ComfyUIClient("http://comfy.test")
|
||||
monkeypatch.delitem(sys.modules, "websocket", raising=False)
|
||||
monkeypatch.setattr(
|
||||
"builtins.__import__",
|
||||
_raise_on_websocket_import(__import__),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
client, "poll", lambda prompt_id, **kwargs: {"outputs": {"used": "poll"}}
|
||||
)
|
||||
|
||||
entry = client._wait("p5", timeout=5, interval=5)
|
||||
assert entry == {"outputs": {"used": "poll"}}
|
||||
|
||||
def test_wait_does_not_swallow_genuine_comfyuierror_from_websocket(self, monkeypatch):
|
||||
"""A real execution error detected over the websocket must propagate,
|
||||
not be masked by a fallback-to-poll retry."""
|
||||
from tools._comfyui.client import ComfyUIClient, ComfyUIError
|
||||
|
||||
client = ComfyUIClient("http://comfy.test")
|
||||
frames = [
|
||||
json.dumps({"type": "execution_error", "data": {
|
||||
"prompt_id": "p6", "exception_message": "bad node",
|
||||
}}),
|
||||
]
|
||||
_install_fake_websocket(monkeypatch, frames)
|
||||
|
||||
def fail_poll(prompt_id, **kwargs):
|
||||
raise AssertionError("poll() should not be called after a real ws error")
|
||||
|
||||
monkeypatch.setattr(client, "poll", fail_poll)
|
||||
|
||||
with pytest.raises(ComfyUIError) as excinfo:
|
||||
client._wait("p6", timeout=5, interval=5)
|
||||
assert excinfo.value.prompt_id == "p6"
|
||||
|
||||
|
||||
def _raise_on_websocket_import(real_import):
|
||||
def _import(name, *args, **kwargs):
|
||||
if name == "websocket":
|
||||
raise ImportError("no module named websocket")
|
||||
return real_import(name, *args, **kwargs)
|
||||
return _import
|
||||
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Model discovery (offline, no server needed)
|
||||
@@ -332,6 +685,11 @@ class TestModelRequirements:
|
||||
assert len(_REQUIRED_MODELS_T2V) > 0
|
||||
assert any("t2v" in m.lower() for m in _REQUIRED_MODELS_T2V)
|
||||
|
||||
def test_music_tool_has_required_models(self):
|
||||
from tools.audio.comfyui_music import _REQUIRED_MODELS
|
||||
assert len(_REQUIRED_MODELS) > 0
|
||||
assert any("ace_step" in m.lower() for m in _REQUIRED_MODELS)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Custom workflow contract and provenance
|
||||
@@ -418,6 +776,77 @@ class TestCustomWorkflowContract:
|
||||
assert provenance["model_stack"] == [{"role": "lora", "name": "style.safetensors"}]
|
||||
assert provenance["model_stack_source"] == "caller_supplied"
|
||||
|
||||
def test_video_timeout_surfaces_resumable_prompt_id(self, tmp_path):
|
||||
from tools._comfyui.client import ComfyUIError
|
||||
|
||||
tool = ComfyUIVideo()
|
||||
tool._client.is_available = lambda: True
|
||||
|
||||
def fake_generate(workflow, output_node, dest, **kwargs):
|
||||
raise ComfyUIError("Prompt timed-out-id did not complete within 5s", prompt_id="timed-out-id")
|
||||
|
||||
tool._client.generate = fake_generate
|
||||
|
||||
result = tool.execute({
|
||||
"prompt": "test",
|
||||
"workflow_json": json.dumps({"42": {"inputs": {}}}),
|
||||
"output_node": "42",
|
||||
"output_path": str(tmp_path / "video.mp4"),
|
||||
"timeout_seconds": 5,
|
||||
})
|
||||
|
||||
assert result.success is False
|
||||
assert result.data["prompt_id"] == "timed-out-id"
|
||||
assert "resume_prompt_id" in result.error
|
||||
assert "timed-out-id" in result.error
|
||||
|
||||
def test_video_passes_timeout_and_resume_prompt_id_through(self, tmp_path):
|
||||
tool = ComfyUIVideo()
|
||||
tool._client.is_available = lambda: True
|
||||
seen = {}
|
||||
|
||||
def fake_generate(workflow, output_node, dest, **kwargs):
|
||||
seen.update(kwargs)
|
||||
return [Path(dest)]
|
||||
|
||||
tool._client.generate = fake_generate
|
||||
|
||||
result = tool.execute({
|
||||
"prompt": "test",
|
||||
"workflow_json": json.dumps({"42": {"inputs": {}}}),
|
||||
"output_node": "42",
|
||||
"output_path": str(tmp_path / "video.mp4"),
|
||||
"timeout_seconds": 7200,
|
||||
"resume_prompt_id": "already-running-id",
|
||||
})
|
||||
|
||||
assert result.success is True
|
||||
assert seen["timeout"] == 7200
|
||||
assert seen["resume_prompt_id"] == "already-running-id"
|
||||
|
||||
def test_video_default_timeout_is_generous_not_900s(self, tmp_path):
|
||||
tool = ComfyUIVideo()
|
||||
tool._client.is_available = lambda: True
|
||||
seen = {}
|
||||
|
||||
def fake_generate(workflow, output_node, dest, **kwargs):
|
||||
seen.update(kwargs)
|
||||
return [Path(dest)]
|
||||
|
||||
tool._client.generate = fake_generate
|
||||
|
||||
tool.execute({
|
||||
"prompt": "test",
|
||||
"workflow_json": json.dumps({"42": {"inputs": {}}}),
|
||||
"output_node": "42",
|
||||
"output_path": str(tmp_path / "video.mp4"),
|
||||
})
|
||||
|
||||
# Regression guard: the old hardcoded 900s timeout false-failed real
|
||||
# renders on modest local GPUs (observed ~1360-1630s for non-accelerated
|
||||
# custom Wan 1.3B workflows at 832x480/81-97 frames).
|
||||
assert seen["timeout"] > 900
|
||||
|
||||
def test_image_missing_models_are_structured(self):
|
||||
tool = ComfyUIImage()
|
||||
tool._client.is_available = lambda: True
|
||||
@@ -466,6 +895,216 @@ class TestCustomWorkflowContract:
|
||||
assert any(item["role"] == "vae" for item in provenance["model_stack"])
|
||||
|
||||
|
||||
class TestComfyUIMusic:
|
||||
|
||||
def test_capability_and_provider(self):
|
||||
tool = ComfyUIMusic()
|
||||
assert tool.capability == "music_generation"
|
||||
assert tool.provider == "comfyui"
|
||||
|
||||
def test_bundled_path_requires_no_workflow_json_or_output_node(self, tmp_path):
|
||||
"""Without workflow_json/workflow_path it should attempt the bundled
|
||||
ACE-Step workflow, not demand a custom one."""
|
||||
tool = ComfyUIMusic()
|
||||
tool._client.is_available = lambda: True
|
||||
tool._client.check_models = lambda required: (list(required), [])
|
||||
tool._client.generate = lambda workflow, output_node, dest, **kwargs: [Path(dest)]
|
||||
|
||||
result = tool.execute({
|
||||
"prompt": "ambient pad",
|
||||
"output_path": str(tmp_path / "music.mp3"),
|
||||
})
|
||||
|
||||
assert result.success is True
|
||||
assert result.data["workflow_provenance"]["source"] == "bundled"
|
||||
|
||||
def test_custom_workflow_without_output_node_errors(self):
|
||||
tool = ComfyUIMusic()
|
||||
tool._client.is_available = lambda: True
|
||||
|
||||
result = tool.execute({
|
||||
"prompt": "ambient pad",
|
||||
"workflow_json": json.dumps({"9": {"inputs": {}}}),
|
||||
})
|
||||
|
||||
assert result.success is False
|
||||
assert "output_node" in result.error
|
||||
|
||||
def test_bundled_missing_models_returns_structured_payload(self):
|
||||
tool = ComfyUIMusic()
|
||||
tool._client.is_available = lambda: True
|
||||
tool._client.check_models = lambda required: ([], list(required))
|
||||
|
||||
result = tool.execute({"prompt": "ambient pad"})
|
||||
|
||||
assert result.success is False
|
||||
assert result.data["missing_models"][0]["name"] == "ace_step_v1_3.5b.safetensors"
|
||||
assert result.data["missing_models"][0]["download_url"]
|
||||
|
||||
def test_bundled_generation_patches_tags_lyrics_and_seed(self, tmp_path):
|
||||
tool = ComfyUIMusic()
|
||||
tool._client.is_available = lambda: True
|
||||
tool._client.check_models = lambda required: (list(required), [])
|
||||
seen = {}
|
||||
|
||||
def fake_generate(workflow, output_node, dest, **kwargs):
|
||||
seen["workflow"] = workflow
|
||||
seen["output_node"] = output_node
|
||||
return [Path(dest)]
|
||||
|
||||
tool._client.generate = fake_generate
|
||||
|
||||
result = tool.execute({
|
||||
"prompt": "lofi hip hop, chill, rain sounds",
|
||||
"lyrics": "[verse]\nquiet streets",
|
||||
"duration_seconds": 45,
|
||||
"seed": 777,
|
||||
"output_path": str(tmp_path / "music.mp3"),
|
||||
})
|
||||
|
||||
assert result.success is True
|
||||
assert seen["output_node"] == "10"
|
||||
assert seen["workflow"]["2"]["inputs"]["tags"] == "lofi hip hop, chill, rain sounds"
|
||||
assert seen["workflow"]["2"]["inputs"]["lyrics"] == "[verse]\nquiet streets"
|
||||
assert seen["workflow"]["4"]["inputs"]["seconds"] == 45
|
||||
assert seen["workflow"]["8"]["inputs"]["seed"] == 777
|
||||
assert result.data["model"] == "ace-step-v1-3.5b"
|
||||
|
||||
def test_bundled_generation_preserves_seed_zero(self, tmp_path):
|
||||
tool = ComfyUIMusic()
|
||||
tool._client.is_available = lambda: True
|
||||
tool._client.check_models = lambda required: (list(required), [])
|
||||
seen = {}
|
||||
|
||||
def fake_generate(workflow, output_node, dest, **kwargs):
|
||||
seen["seed"] = workflow["8"]["inputs"]["seed"]
|
||||
return [Path(dest)]
|
||||
|
||||
tool._client.generate = fake_generate
|
||||
result = tool.execute({
|
||||
"prompt": "deterministic test",
|
||||
"seed": 0,
|
||||
"output_path": str(tmp_path / "music.mp3"),
|
||||
})
|
||||
|
||||
assert result.success is True
|
||||
assert result.seed == 0
|
||||
assert seen["seed"] == 0
|
||||
|
||||
def test_get_status_degraded_when_model_missing(self):
|
||||
tool = ComfyUIMusic()
|
||||
tool._client.is_available = lambda: True
|
||||
tool._client.check_models = lambda required: ([], list(required))
|
||||
assert tool.get_status() == ToolStatus.DEGRADED
|
||||
|
||||
def test_unavailable_server_reports_unavailable_reason(self):
|
||||
tool = ComfyUIMusic()
|
||||
tool._client.is_available = lambda: False
|
||||
tool._client.unavailable_reason = lambda: "no server here"
|
||||
|
||||
result = tool.execute({
|
||||
"prompt": "ambient pad",
|
||||
"workflow_json": json.dumps({"9": {"inputs": {}}}),
|
||||
"output_node": "9",
|
||||
})
|
||||
|
||||
assert result.success is False
|
||||
assert result.error == "no server here"
|
||||
|
||||
def test_successful_generation_returns_provenance_and_duration(self, tmp_path, monkeypatch):
|
||||
tool = ComfyUIMusic()
|
||||
tool._client.is_available = lambda: True
|
||||
|
||||
dest_file = tmp_path / "music.mp3"
|
||||
|
||||
def fake_generate(workflow, output_node, dest, **kwargs):
|
||||
Path(dest).write_bytes(b"fake-audio-bytes")
|
||||
return [Path(dest)]
|
||||
|
||||
tool._client.generate = fake_generate
|
||||
monkeypatch.setattr("shutil.which", lambda name: None) # no ffprobe in test env
|
||||
|
||||
result = tool.execute({
|
||||
"prompt": "upbeat synthwave",
|
||||
"workflow_json": json.dumps({"9": {"inputs": {}}}),
|
||||
"output_node": "9",
|
||||
"output_path": str(dest_file),
|
||||
"workflow_name": "my-ace-step-graph",
|
||||
"workflow_model": "ace-step-v1-3.5b",
|
||||
})
|
||||
|
||||
assert result.success is True
|
||||
assert result.data["provider"] == "comfyui"
|
||||
assert result.data["model"] == "ace-step-v1-3.5b"
|
||||
assert result.data["output"] == str(dest_file)
|
||||
assert result.data["format"] == "mp3"
|
||||
assert result.data["duration_seconds"] is None # ffprobe unavailable
|
||||
provenance = result.data["workflow_provenance"]
|
||||
assert provenance["source"] == "user_supplied"
|
||||
assert provenance["output_node"] == "9"
|
||||
assert provenance["workflow_hash_sha256"]
|
||||
|
||||
def test_timeout_surfaces_resumable_prompt_id(self, tmp_path):
|
||||
from tools._comfyui.client import ComfyUIError
|
||||
|
||||
tool = ComfyUIMusic()
|
||||
tool._client.is_available = lambda: True
|
||||
|
||||
def fake_generate(workflow, output_node, dest, **kwargs):
|
||||
raise ComfyUIError("timed out", prompt_id="music-prompt-id")
|
||||
|
||||
tool._client.generate = fake_generate
|
||||
|
||||
result = tool.execute({
|
||||
"prompt": "ambient pad",
|
||||
"workflow_json": json.dumps({"9": {"inputs": {}}}),
|
||||
"output_node": "9",
|
||||
"output_path": str(tmp_path / "music.mp3"),
|
||||
})
|
||||
|
||||
assert result.success is False
|
||||
assert result.data["prompt_id"] == "music-prompt-id"
|
||||
assert "resume_prompt_id" in result.error
|
||||
|
||||
def test_passes_timeout_and_resume_prompt_id_through(self, tmp_path):
|
||||
tool = ComfyUIMusic()
|
||||
tool._client.is_available = lambda: True
|
||||
seen = {}
|
||||
|
||||
def fake_generate(workflow, output_node, dest, **kwargs):
|
||||
seen.update(kwargs)
|
||||
return [Path(dest)]
|
||||
|
||||
tool._client.generate = fake_generate
|
||||
|
||||
tool.execute({
|
||||
"prompt": "ambient pad",
|
||||
"workflow_json": json.dumps({"9": {"inputs": {}}}),
|
||||
"output_node": "9",
|
||||
"output_path": str(tmp_path / "music.mp3"),
|
||||
"timeout_seconds": 3600,
|
||||
"resume_prompt_id": "already-running-id",
|
||||
})
|
||||
|
||||
assert seen["timeout"] == 3600
|
||||
assert seen["resume_prompt_id"] == "already-running-id"
|
||||
|
||||
def test_registry_discovers_comfyui_music_under_music_generation(self):
|
||||
registry = ToolRegistry()
|
||||
tool = ComfyUIMusic()
|
||||
registry.register(tool)
|
||||
registry._discovered_packages.add("tools")
|
||||
|
||||
by_capability = registry.get_by_capability("music_generation")
|
||||
assert any(t.name == "comfyui_music" for t in by_capability)
|
||||
|
||||
def test_uses_music_capability_env_var_for_multi_server(self, monkeypatch):
|
||||
monkeypatch.delenv("COMFYUI_SERVER_URL", raising=False)
|
||||
monkeypatch.setenv("COMFYUI_MUSIC_SERVER_URL", "http://music-gpu:8188")
|
||||
tool = ComfyUIMusic()
|
||||
assert tool._client.server_url == "http://music-gpu:8188"
|
||||
|
||||
|
||||
class TestComfyUISetupOffer:
|
||||
|
||||
def test_provider_menu_summary_includes_structured_setup_offer(self):
|
||||
|
||||
@@ -11,14 +11,25 @@ import json
|
||||
import os
|
||||
import random
|
||||
import time
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from typing import Any, Callable
|
||||
|
||||
import requests
|
||||
|
||||
|
||||
class ComfyUIError(Exception):
|
||||
"""Raised when ComfyUI returns an error or times out."""
|
||||
"""Raised when ComfyUI returns an error or times out.
|
||||
|
||||
``prompt_id`` is set when the error follows a successful ``submit()``,
|
||||
so callers can recover a timed-out-but-still-running job instead of
|
||||
losing track of it: poll ``GET /history/{prompt_id}`` directly, or
|
||||
pass ``resume_prompt_id`` back into ``ComfyUIVideo.execute()``.
|
||||
"""
|
||||
|
||||
def __init__(self, message: str, prompt_id: str | None = None) -> None:
|
||||
super().__init__(message)
|
||||
self.prompt_id = prompt_id
|
||||
|
||||
|
||||
class ComfyUIClient:
|
||||
@@ -31,11 +42,32 @@ class ComfyUIClient:
|
||||
4. POST /upload/image → stage a local image for I2V workflows
|
||||
"""
|
||||
|
||||
def __init__(self, server_url: str | None = None) -> None:
|
||||
self.server_url = (
|
||||
server_url
|
||||
or os.environ.get("COMFYUI_SERVER_URL", "http://localhost:8188")
|
||||
).rstrip("/")
|
||||
def __init__(
|
||||
self, server_url: str | None = None, capability: str | None = None
|
||||
) -> None:
|
||||
"""*capability*, if given (e.g. ``"image"``, ``"video"``), lets a
|
||||
per-capability env var (``COMFYUI_{CAPABILITY}_SERVER_URL``) point
|
||||
this client at its own ComfyUI instance -- useful when image and
|
||||
video generation are split across separate servers/GPUs. Falls back
|
||||
to the shared ``COMFYUI_SERVER_URL`` when the capability-specific
|
||||
var isn't set, so single-server setups need no extra configuration.
|
||||
"""
|
||||
self.capability = capability
|
||||
self._capability_env_var = (
|
||||
f"COMFYUI_{capability.upper()}_SERVER_URL" if capability else None
|
||||
)
|
||||
resolved = server_url or self._capability_url() or os.environ.get(
|
||||
"COMFYUI_SERVER_URL"
|
||||
)
|
||||
self.server_url = (resolved or "http://localhost:8188").rstrip("/")
|
||||
# Scopes websocket execution events to this client (see wait_ws) and
|
||||
# is echoed back on /prompt so the server targets messages to us.
|
||||
self.client_id = str(uuid.uuid4())
|
||||
|
||||
def _capability_url(self) -> str | None:
|
||||
if self._capability_env_var:
|
||||
return os.environ.get(self._capability_env_var)
|
||||
return None
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Health
|
||||
@@ -43,8 +75,25 @@ class ComfyUIClient:
|
||||
|
||||
@property
|
||||
def is_default_url(self) -> bool:
|
||||
"""True if using the fallback URL (user didn't set COMFYUI_SERVER_URL)."""
|
||||
return not os.environ.get("COMFYUI_SERVER_URL")
|
||||
"""True if neither the capability-specific nor shared env var is set."""
|
||||
return not (self._capability_url() or os.environ.get("COMFYUI_SERVER_URL"))
|
||||
|
||||
def unavailable_reason(self) -> str:
|
||||
"""Human-readable explanation of why the server can't be reached."""
|
||||
env_var_hint = self._capability_env_var or "COMFYUI_SERVER_URL"
|
||||
if self._capability_env_var:
|
||||
env_var_hint += " (or the shared COMFYUI_SERVER_URL)"
|
||||
if self.is_default_url:
|
||||
return (
|
||||
f"No ComfyUI server found at {self.server_url} "
|
||||
f"(default — no server URL configured).\n"
|
||||
f"Set {env_var_hint} in your .env file to the address of "
|
||||
f"your ComfyUI server (e.g. http://localhost:8188)."
|
||||
)
|
||||
return (
|
||||
f"ComfyUI server not reachable at {self.server_url}.\n"
|
||||
f"Check that ComfyUI is running and the URL is correct."
|
||||
)
|
||||
|
||||
def is_available(self) -> bool:
|
||||
"""Return True if the ComfyUI server is reachable."""
|
||||
@@ -56,20 +105,6 @@ class ComfyUIClient:
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def unavailable_reason(self) -> str:
|
||||
"""Human-readable explanation of why the server can't be reached."""
|
||||
if self.is_default_url:
|
||||
return (
|
||||
f"No ComfyUI server found at {self.server_url} "
|
||||
f"(default — no COMFYUI_SERVER_URL configured).\n"
|
||||
f"Set COMFYUI_SERVER_URL in your .env file to the address of "
|
||||
f"your ComfyUI server (e.g. http://localhost:8188)."
|
||||
)
|
||||
return (
|
||||
f"ComfyUI server not reachable at {self.server_url}.\n"
|
||||
f"Check that ComfyUI is running and the URL is correct."
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Model discovery
|
||||
# ------------------------------------------------------------------
|
||||
@@ -137,7 +172,7 @@ class ComfyUIClient:
|
||||
"""Queue a workflow for execution. Returns the ``prompt_id``."""
|
||||
resp = requests.post(
|
||||
f"{self.server_url}/prompt",
|
||||
json={"prompt": workflow},
|
||||
json={"prompt": workflow, "client_id": self.client_id},
|
||||
timeout=30,
|
||||
)
|
||||
try:
|
||||
@@ -164,23 +199,170 @@ class ComfyUIClient:
|
||||
"""Block until *prompt_id* finishes. Returns the history entry."""
|
||||
deadline = time.time() + timeout
|
||||
while time.time() < deadline:
|
||||
resp = requests.get(
|
||||
f"{self.server_url}/history/{prompt_id}", timeout=10
|
||||
)
|
||||
resp.raise_for_status()
|
||||
history = resp.json()
|
||||
if prompt_id in history:
|
||||
entry = history[prompt_id]
|
||||
status = entry.get("status", {})
|
||||
if status.get("status_str") == "error":
|
||||
msgs = status.get("messages", [])
|
||||
raise ComfyUIError(f"Execution error: {msgs}")
|
||||
entry = self._history_entry(prompt_id)
|
||||
if entry is not None:
|
||||
return entry
|
||||
time.sleep(interval)
|
||||
raise ComfyUIError(
|
||||
f"Prompt {prompt_id} did not complete within {timeout}s"
|
||||
f"Prompt {prompt_id} did not complete within {timeout}s. "
|
||||
f"The job is very likely still running on the ComfyUI server "
|
||||
f"(local/custom workflows on modest GPUs routinely exceed the "
|
||||
f"client wait) — it was not cancelled. Poll "
|
||||
f"GET {{server_url}}/history/{prompt_id} directly, or call "
|
||||
f"generate()/execute() again with a longer timeout and this "
|
||||
f"prompt_id to resume waiting without resubmitting.",
|
||||
prompt_id=prompt_id,
|
||||
)
|
||||
|
||||
def _history_entry(self, prompt_id: str) -> dict | None:
|
||||
"""Return a completed history entry, or ``None`` while it is absent."""
|
||||
resp = requests.get(f"{self.server_url}/history/{prompt_id}", timeout=10)
|
||||
resp.raise_for_status()
|
||||
entry = resp.json().get(prompt_id)
|
||||
if entry is None:
|
||||
return None
|
||||
status = entry.get("status", {})
|
||||
if status.get("status_str") == "error":
|
||||
msgs = status.get("messages", [])
|
||||
raise ComfyUIError(f"Execution error: {msgs}", prompt_id=prompt_id)
|
||||
return entry
|
||||
|
||||
def _history_entry_if_reachable(self, prompt_id: str) -> dict | None:
|
||||
"""Best-effort history probe while the websocket remains usable."""
|
||||
try:
|
||||
return self._history_entry(prompt_id)
|
||||
except ComfyUIError:
|
||||
raise
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
def wait_ws(
|
||||
self,
|
||||
prompt_id: str,
|
||||
*,
|
||||
timeout: int = 600,
|
||||
interval: int = 5,
|
||||
on_progress: Callable[[dict], None] | None = None,
|
||||
) -> dict:
|
||||
"""Block until *prompt_id* finishes, watching ComfyUI's websocket feed.
|
||||
|
||||
Reacts to server-pushed ``executing``/``progress``/``execution_error``
|
||||
events instead of sleeping between REST polls, so completion and
|
||||
errors are detected immediately rather than up to *interval* seconds
|
||||
late. *on_progress*, if given, is called with each ``progress``
|
||||
message's ``data`` dict (``value``, ``max``, ``node``, ``prompt_id``).
|
||||
|
||||
Requires the optional ``websocket-client`` package. Any transport
|
||||
failure (missing dependency, connection refused, dropped socket,
|
||||
malformed frame) propagates as a plain exception — callers should
|
||||
catch it and fall back to :meth:`poll`, which is what :meth:`generate`
|
||||
does. A genuine ComfyUI-side execution error or an unmet deadline is
|
||||
raised as :class:`ComfyUIError` with ``prompt_id`` set, exactly like
|
||||
:meth:`poll`, so ``resume_prompt_id`` recovery works the same way
|
||||
regardless of which wait strategy was used.
|
||||
"""
|
||||
import websocket # websocket-client; optional, see docstring
|
||||
|
||||
# History is authoritative and websocket events are not replayed. The
|
||||
# job may already have finished between submit() and this wait call.
|
||||
entry = self._history_entry_if_reachable(prompt_id)
|
||||
if entry is not None:
|
||||
return entry
|
||||
|
||||
ws_url = self.server_url.replace("http://", "ws://", 1).replace(
|
||||
"https://", "wss://", 1
|
||||
)
|
||||
conn = websocket.create_connection(
|
||||
f"{ws_url}/ws?clientId={self.client_id}", timeout=10
|
||||
)
|
||||
try:
|
||||
conn.settimeout(interval)
|
||||
deadline = time.time() + timeout
|
||||
finished = False
|
||||
# Close the remaining race between the first history probe and
|
||||
# websocket connection establishment. Events after this point are
|
||||
# queued on the open socket; earlier completion is in history.
|
||||
entry = self._history_entry_if_reachable(prompt_id)
|
||||
if entry is not None:
|
||||
return entry
|
||||
while time.time() < deadline:
|
||||
try:
|
||||
raw = conn.recv()
|
||||
except websocket.WebSocketTimeoutException:
|
||||
entry = self._history_entry_if_reachable(prompt_id)
|
||||
if entry is not None:
|
||||
return entry
|
||||
continue
|
||||
if not isinstance(raw, str):
|
||||
continue # binary preview-image frame, not a status message
|
||||
try:
|
||||
message = json.loads(raw)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
data = message.get("data", {})
|
||||
if data.get("prompt_id") not in (None, prompt_id):
|
||||
continue # another job sharing this connection
|
||||
msg_type = message.get("type")
|
||||
if msg_type == "progress":
|
||||
if on_progress:
|
||||
on_progress(data)
|
||||
elif msg_type == "execution_error":
|
||||
raise ComfyUIError(
|
||||
f"Execution error: {data}", prompt_id=prompt_id
|
||||
)
|
||||
elif msg_type == "executing" and data.get("node") is None:
|
||||
finished = True
|
||||
break
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
if not finished:
|
||||
entry = self._history_entry_if_reachable(prompt_id)
|
||||
if entry is not None:
|
||||
return entry
|
||||
raise ComfyUIError(
|
||||
f"Prompt {prompt_id} did not complete within {timeout}s "
|
||||
f"(websocket wait). The job was not cancelled — resume with "
|
||||
f"resume_prompt_id={prompt_id!r} and a longer timeout.",
|
||||
prompt_id=prompt_id,
|
||||
)
|
||||
|
||||
entry = self._history_entry(prompt_id)
|
||||
if entry is None:
|
||||
raise ComfyUIError(
|
||||
f"No history entry for {prompt_id} after completion",
|
||||
prompt_id=prompt_id,
|
||||
)
|
||||
return entry
|
||||
|
||||
def _wait(
|
||||
self,
|
||||
prompt_id: str,
|
||||
*,
|
||||
timeout: int,
|
||||
interval: int,
|
||||
on_progress: Callable[[dict], None] | None = None,
|
||||
) -> dict:
|
||||
"""Wait for *prompt_id*, preferring the websocket feed over polling.
|
||||
|
||||
Falls back to :meth:`poll` when ``websocket-client`` isn't installed
|
||||
or the websocket can't be established/maintained. A genuine
|
||||
:class:`ComfyUIError` (execution error or deadline reached) is never
|
||||
swallowed by the fallback — only transport-level failures are. The
|
||||
fallback gets whatever's left of *timeout*, not a fresh budget, so a
|
||||
mid-wait websocket drop can't double the caller's worst-case wait.
|
||||
"""
|
||||
started = time.time()
|
||||
try:
|
||||
return self.wait_ws(
|
||||
prompt_id, timeout=timeout, interval=interval, on_progress=on_progress
|
||||
)
|
||||
except ComfyUIError:
|
||||
raise
|
||||
except Exception:
|
||||
remaining = max(timeout - (time.time() - started), 0)
|
||||
return self.poll(prompt_id, timeout=remaining, interval=interval)
|
||||
|
||||
def download(
|
||||
self,
|
||||
filename: str,
|
||||
@@ -229,16 +411,41 @@ class ComfyUIClient:
|
||||
*,
|
||||
timeout: int = 600,
|
||||
interval: int = 5,
|
||||
resume_prompt_id: str | None = None,
|
||||
on_progress: Callable[[dict], None] | None = None,
|
||||
) -> list[Path]:
|
||||
"""Submit → poll → download. Returns list of artifact paths."""
|
||||
prompt_id = self.submit(workflow)
|
||||
entry = self.poll(prompt_id, timeout=timeout, interval=interval)
|
||||
"""Submit → wait → download. Returns list of artifact paths.
|
||||
|
||||
Pass ``resume_prompt_id`` (from a previous ``ComfyUIError.prompt_id``)
|
||||
to skip re-submitting an already-queued/running job and just resume
|
||||
waiting on it — the common recovery path after a timeout.
|
||||
|
||||
Waiting prefers ComfyUI's websocket feed (immediate completion/error
|
||||
detection, optional live ``on_progress`` callback) and transparently
|
||||
falls back to REST polling if ``websocket-client`` isn't installed or
|
||||
the connection can't be used. See :meth:`_wait`.
|
||||
"""
|
||||
prompt_id = resume_prompt_id or self.submit(workflow)
|
||||
if resume_prompt_id:
|
||||
# A prompt resumed by a new client instance was submitted with the
|
||||
# original instance's client_id, so its websocket events are not
|
||||
# guaranteed to reach this socket. Poll authoritative history.
|
||||
entry = self.poll(prompt_id, timeout=timeout, interval=interval)
|
||||
else:
|
||||
entry = self._wait(
|
||||
prompt_id, timeout=timeout, interval=interval, on_progress=on_progress
|
||||
)
|
||||
|
||||
outputs = entry.get("outputs", {})
|
||||
node_output = outputs.get(output_node, {})
|
||||
|
||||
# ComfyUI stores images and videos under the "images" key
|
||||
items = node_output.get("images", []) or node_output.get("gifs", [])
|
||||
# ComfyUI stores images/video frames under "images", legacy GIFs
|
||||
# under "gifs", and the native SaveAudio node's output under "audio".
|
||||
items = (
|
||||
node_output.get("images", [])
|
||||
or node_output.get("gifs", [])
|
||||
or node_output.get("audio", [])
|
||||
)
|
||||
if not items:
|
||||
raise ComfyUIError(
|
||||
f"No output artifacts on node {output_node}. "
|
||||
|
||||
@@ -18,6 +18,14 @@ COMFYUI_SETUP_OFFER: dict[str, Any] = {
|
||||
"free local video generation through ComfyUI workflows",
|
||||
"community workflow_json/workflow_path execution",
|
||||
],
|
||||
# Optional: point image/video generation at separate ComfyUI instances
|
||||
# (e.g. different GPUs). Each overrides COMFYUI_SERVER_URL for its own
|
||||
# tool only; single-server setups can ignore this entirely.
|
||||
"per_capability_env_var_overrides": {
|
||||
"comfyui_image": "COMFYUI_IMAGE_SERVER_URL",
|
||||
"comfyui_video": "COMFYUI_VIDEO_SERVER_URL",
|
||||
"comfyui_music": "COMFYUI_MUSIC_SERVER_URL",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@@ -176,6 +184,17 @@ BUNDLED_MODEL_STACKS: dict[str, list[dict[str, Any]]] = {
|
||||
),
|
||||
},
|
||||
],
|
||||
"ace-step-1-t2a": [
|
||||
{
|
||||
"role": "checkpoint",
|
||||
"name": "ace_step_v1_3.5b.safetensors",
|
||||
"destination_hint": "ComfyUI/models/checkpoints/",
|
||||
"download_url": (
|
||||
"https://huggingface.co/Comfy-Org/ACE-Step_ComfyUI_repackaged/"
|
||||
"blob/main/all_in_one/ace_step_v1_3.5b.safetensors"
|
||||
),
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
|
||||
80
tools/_comfyui/workflows/ace-step-1-t2a.json
Normal file
80
tools/_comfyui/workflows/ace-step-1-t2a.json
Normal file
@@ -0,0 +1,80 @@
|
||||
{
|
||||
"1": {
|
||||
"class_type": "CheckpointLoaderSimple",
|
||||
"inputs": {
|
||||
"ckpt_name": "ace_step_v1_3.5b.safetensors"
|
||||
}
|
||||
},
|
||||
"2": {
|
||||
"class_type": "TextEncodeAceStepAudio",
|
||||
"inputs": {
|
||||
"clip": ["1", 1],
|
||||
"tags": "",
|
||||
"lyrics": "",
|
||||
"lyrics_strength": 0.99
|
||||
}
|
||||
},
|
||||
"3": {
|
||||
"class_type": "ConditioningZeroOut",
|
||||
"inputs": {
|
||||
"conditioning": ["2", 0]
|
||||
}
|
||||
},
|
||||
"4": {
|
||||
"class_type": "EmptyAceStepLatentAudio",
|
||||
"inputs": {
|
||||
"seconds": 120,
|
||||
"batch_size": 1
|
||||
}
|
||||
},
|
||||
"5": {
|
||||
"class_type": "ModelSamplingSD3",
|
||||
"inputs": {
|
||||
"model": ["1", 0],
|
||||
"shift": 5.0
|
||||
}
|
||||
},
|
||||
"6": {
|
||||
"class_type": "LatentOperationTonemapReinhard",
|
||||
"inputs": {
|
||||
"multiplier": 1.0
|
||||
}
|
||||
},
|
||||
"7": {
|
||||
"class_type": "LatentApplyOperationCFG",
|
||||
"inputs": {
|
||||
"model": ["5", 0],
|
||||
"operation": ["6", 0]
|
||||
}
|
||||
},
|
||||
"8": {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {
|
||||
"model": ["7", 0],
|
||||
"positive": ["2", 0],
|
||||
"negative": ["3", 0],
|
||||
"latent_image": ["4", 0],
|
||||
"seed": 0,
|
||||
"steps": 50,
|
||||
"cfg": 5.0,
|
||||
"sampler_name": "euler",
|
||||
"scheduler": "simple",
|
||||
"denoise": 1.0
|
||||
}
|
||||
},
|
||||
"9": {
|
||||
"class_type": "VAEDecodeAudio",
|
||||
"inputs": {
|
||||
"samples": ["8", 0],
|
||||
"vae": ["1", 2]
|
||||
}
|
||||
},
|
||||
"10": {
|
||||
"class_type": "SaveAudioMP3",
|
||||
"inputs": {
|
||||
"audio": ["9", 0],
|
||||
"filename_prefix": "openmontage",
|
||||
"quality": "V0"
|
||||
}
|
||||
}
|
||||
}
|
||||
367
tools/audio/comfyui_music.py
Normal file
367
tools/audio/comfyui_music.py
Normal file
@@ -0,0 +1,367 @@
|
||||
"""ComfyUI music generation via a local or remote ComfyUI server.
|
||||
|
||||
Default workflow: ACE-Step v1 (3.5B) text-to-audio using ComfyUI's native
|
||||
``TextEncodeAceStepAudio``/``EmptyAceStepLatentAudio`` nodes (built into
|
||||
ComfyUI core, not a third-party pack). Custom workflows are still accepted
|
||||
via ``workflow_json``/``workflow_path`` for other ACE-Step node packs, other
|
||||
versions (e.g. ACE-Step 1.5), or entirely different audio models -- the same
|
||||
override contract ``comfyui_image``/``comfyui_video`` offer.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import shutil
|
||||
import subprocess
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
RetryPolicy,
|
||||
ToolResult,
|
||||
ToolRuntime,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
from tools._comfyui.client import ComfyUIClient, ComfyUIError
|
||||
from tools._comfyui.metadata import (
|
||||
BUNDLED_MODEL_STACKS,
|
||||
COMFYUI_SETUP_OFFER,
|
||||
missing_models_payload,
|
||||
model_stack,
|
||||
workflow_hash,
|
||||
)
|
||||
|
||||
_WORKFLOWS = Path(__file__).resolve().parent.parent / "_comfyui" / "workflows"
|
||||
|
||||
# Model required by the bundled ACE-Step v1 workflow
|
||||
_REQUIRED_MODELS = ["ace_step_v1_3.5b.safetensors"]
|
||||
|
||||
|
||||
class ComfyUIMusic(BaseTool):
|
||||
name = "comfyui_music"
|
||||
version = "0.2.0"
|
||||
tier = ToolTier.GENERATE
|
||||
capability = "music_generation"
|
||||
provider = "comfyui"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.SEEDED
|
||||
runtime = ToolRuntime.LOCAL_GPU
|
||||
|
||||
dependencies = [] # checked at runtime via server health
|
||||
setup_offer = COMFYUI_SETUP_OFFER
|
||||
install_instructions = (
|
||||
"Start a ComfyUI server and set COMFYUI_SERVER_URL "
|
||||
"(default http://localhost:8188).\n"
|
||||
"Requires ace_step_v1_3.5b.safetensors in ComfyUI's checkpoints "
|
||||
"directory for the bundled workflow.\n"
|
||||
"Running a separate ComfyUI instance for music? Set "
|
||||
"COMFYUI_MUSIC_SERVER_URL instead -- it takes priority over "
|
||||
"COMFYUI_SERVER_URL for this tool only."
|
||||
)
|
||||
agent_skills = ["comfyui"]
|
||||
|
||||
capabilities = ["generate_background_music", "generate_song", "generate_instrumental"]
|
||||
supports = {
|
||||
"seed": True,
|
||||
"lyrics": True,
|
||||
"custom_workflow": True,
|
||||
"custom_output_node": True,
|
||||
"offline": True,
|
||||
}
|
||||
best_for = [
|
||||
"local GPU music generation without API costs",
|
||||
"instrumentals and songs with lyrics via the bundled ACE-Step v1 workflow",
|
||||
"full control over sampling or other ACE-Step versions/node packs via custom ComfyUI workflows",
|
||||
]
|
||||
not_good_for = [
|
||||
"setups without a running ComfyUI server",
|
||||
"CPU-only machines",
|
||||
]
|
||||
fallback_tools = ["suno_music", "music_gen"]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["prompt"],
|
||||
"properties": {
|
||||
"prompt": {
|
||||
"type": "string",
|
||||
"description": (
|
||||
"Style/mood/genre description (ACE-Step 'tags'), e.g. "
|
||||
"'upbeat electronic pop, female vocals, driving bassline'. "
|
||||
"Comma-separated tags work best. Not injected for custom workflows."
|
||||
),
|
||||
},
|
||||
"lyrics": {
|
||||
"type": "string",
|
||||
"default": "",
|
||||
"description": (
|
||||
"Optional lyrics. Leave empty for instrumental. Supports structure "
|
||||
"tags like [verse]/[chorus]/[bridge] and language-code prefixes "
|
||||
"(e.g. [zh], [ja]) for non-English lines."
|
||||
),
|
||||
},
|
||||
"duration_seconds": {"type": "number", "default": 120.0},
|
||||
"steps": {"type": "integer", "default": 50},
|
||||
"cfg": {"type": "number", "default": 5.0},
|
||||
"lyrics_strength": {"type": "number", "default": 0.99},
|
||||
"seed": {"type": "integer", "description": "Random if omitted"},
|
||||
"output_path": {"type": "string", "description": "Where to save the audio"},
|
||||
"workflow_json": {
|
||||
"type": "string",
|
||||
"description": "Optional full ComfyUI workflow JSON. Requires output_node.",
|
||||
},
|
||||
"workflow_path": {
|
||||
"type": "string",
|
||||
"description": "Optional path to a ComfyUI workflow JSON file. Requires output_node.",
|
||||
},
|
||||
"output_node": {
|
||||
"type": "string",
|
||||
"description": "ComfyUI output node ID for custom workflow_json/workflow_path.",
|
||||
},
|
||||
"workflow_name": {
|
||||
"type": "string",
|
||||
"description": "Optional human-readable provenance label for a custom workflow.",
|
||||
},
|
||||
"workflow_model": {
|
||||
"type": "string",
|
||||
"description": "Optional model/provenance label for a custom workflow.",
|
||||
},
|
||||
"workflow_model_stack": {
|
||||
"type": "array",
|
||||
"description": (
|
||||
"Optional provenance metadata for custom workflow dependencies. "
|
||||
"Items should include name, role, and node-pack origin when known."
|
||||
),
|
||||
"items": {"type": "object"},
|
||||
},
|
||||
"timeout_seconds": {
|
||||
"type": "integer",
|
||||
"description": "How long to wait for the ComfyUI job before giving up. Default 1800s (30min).",
|
||||
},
|
||||
"resume_prompt_id": {
|
||||
"type": "string",
|
||||
"description": "A prompt_id from a previous timed-out call. Skips resubmission and resumes waiting/downloading.",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(
|
||||
cpu_cores=2, ram_mb=8000, vram_mb=8000, disk_mb=500, network_required=False,
|
||||
)
|
||||
retry_policy = RetryPolicy(max_retries=1, retryable_errors=["timeout"])
|
||||
idempotency_key_fields = ["prompt", "lyrics", "duration_seconds", "seed"]
|
||||
side_effects = ["writes audio file to output_path"]
|
||||
user_visible_verification = ["Listen to generated audio for mood, genre accuracy, and quality"]
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._client = ComfyUIClient(capability="music")
|
||||
self._last_progress_log = 0.0
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
if not self._client.is_available():
|
||||
return ToolStatus.UNAVAILABLE
|
||||
_, missing = self._client.check_models(_REQUIRED_MODELS)
|
||||
if missing:
|
||||
return ToolStatus.DEGRADED
|
||||
return ToolStatus.AVAILABLE
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
return 0.0
|
||||
|
||||
def estimate_runtime(self, inputs: dict[str, Any]) -> float:
|
||||
return float(inputs.get("steps", 50)) * 2.0
|
||||
|
||||
def get_info(self) -> dict[str, Any]:
|
||||
info = super().get_info()
|
||||
info["setup_offer"] = self.setup_offer
|
||||
info["bundled_model_stack"] = BUNDLED_MODEL_STACKS["ace-step-1-t2a"]
|
||||
return info
|
||||
|
||||
def _log_progress(self, data: dict) -> None:
|
||||
"""Throttled progress line (see comfyui_video for rationale)."""
|
||||
now = time.monotonic()
|
||||
if now - self._last_progress_log < 10:
|
||||
return
|
||||
self._last_progress_log = now
|
||||
value, max_value = data.get("value"), data.get("max")
|
||||
if value is not None and max_value:
|
||||
print(f"[comfyui_music] step {value}/{max_value}")
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
custom_workflow = bool(inputs.get("workflow_json") or inputs.get("workflow_path"))
|
||||
if custom_workflow and not inputs.get("output_node"):
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=(
|
||||
"Custom ComfyUI workflows require output_node so OpenMontage "
|
||||
"knows which ComfyUI node to download artifacts from."
|
||||
),
|
||||
)
|
||||
|
||||
if not self._client.is_available():
|
||||
return ToolResult(success=False, error=self._client.unavailable_reason())
|
||||
|
||||
if not custom_workflow:
|
||||
_, missing = self._client.check_models(_REQUIRED_MODELS)
|
||||
if missing:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
data=missing_models_payload(
|
||||
missing,
|
||||
workflow_key="ace-step-1-t2a",
|
||||
workflow_name="ace-step-1-t2a.json",
|
||||
),
|
||||
error=(
|
||||
f"ComfyUI server is running but missing required models: "
|
||||
f"{', '.join(missing)}.\n"
|
||||
f"See data.missing_models for destination hints and download URLs."
|
||||
),
|
||||
)
|
||||
|
||||
start = time.time()
|
||||
seed = inputs.get("seed")
|
||||
if seed is None:
|
||||
seed = ComfyUIClient.random_seed()
|
||||
output_path = Path(inputs.get("output_path", f"comfyui_music_{seed}.mp3"))
|
||||
|
||||
try:
|
||||
if custom_workflow:
|
||||
workflow = self._load_custom_workflow(inputs)
|
||||
output_node = str(inputs["output_node"])
|
||||
else:
|
||||
workflow = ComfyUIClient.load_workflow(_WORKFLOWS / "ace-step-1-t2a.json")
|
||||
workflow = ComfyUIClient.patch_workflow(workflow, {
|
||||
"2": {
|
||||
"tags": inputs["prompt"],
|
||||
"lyrics": inputs.get("lyrics", ""),
|
||||
"lyrics_strength": inputs.get("lyrics_strength", 0.99),
|
||||
},
|
||||
"4": {"seconds": inputs.get("duration_seconds", 120.0)},
|
||||
"8": {
|
||||
"seed": seed,
|
||||
"steps": inputs.get("steps", 50),
|
||||
"cfg": inputs.get("cfg", 5.0),
|
||||
},
|
||||
"10": {"filename_prefix": output_path.stem},
|
||||
})
|
||||
output_node = "10"
|
||||
|
||||
provenance = self._workflow_provenance(inputs, custom_workflow, output_node, workflow)
|
||||
paths = self._client.generate(
|
||||
workflow,
|
||||
output_node=output_node,
|
||||
dest=output_path,
|
||||
timeout=inputs.get("timeout_seconds", 1800),
|
||||
interval=10,
|
||||
resume_prompt_id=inputs.get("resume_prompt_id"),
|
||||
on_progress=self._log_progress,
|
||||
)
|
||||
|
||||
except ComfyUIError as exc:
|
||||
data = {"prompt_id": exc.prompt_id} if exc.prompt_id else {}
|
||||
if exc.prompt_id:
|
||||
error_msg = (
|
||||
f"{exc}\n\nThis job was NOT cancelled and is very likely still "
|
||||
f"running server-side. To recover it without resubmitting, call "
|
||||
f"execute() again with resume_prompt_id={exc.prompt_id!r} "
|
||||
f"(and a longer timeout_seconds if it needs more time), or poll "
|
||||
f"GET {{COMFYUI_SERVER_URL}}/history/{exc.prompt_id} directly."
|
||||
)
|
||||
else:
|
||||
error_msg = str(exc)
|
||||
return ToolResult(success=False, error=error_msg, data=data)
|
||||
except Exception as exc:
|
||||
return ToolResult(success=False, error=f"ComfyUI music generation failed: {exc}")
|
||||
|
||||
duration = self._probe_duration(paths[0])
|
||||
model_name = self._model_name(inputs, custom_workflow)
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"provider": "comfyui",
|
||||
"model": model_name,
|
||||
"prompt": inputs["prompt"],
|
||||
"lyrics": inputs.get("lyrics", ""),
|
||||
"duration_seconds": duration,
|
||||
"output": str(paths[0]),
|
||||
"format": paths[0].suffix.lstrip("."),
|
||||
"workflow_provenance": provenance,
|
||||
},
|
||||
artifacts=[str(p) for p in paths],
|
||||
cost_usd=0.0,
|
||||
duration_seconds=round(time.time() - start, 2),
|
||||
seed=seed,
|
||||
model=model_name,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _load_custom_workflow(inputs: dict[str, Any]) -> dict:
|
||||
if inputs.get("workflow_json"):
|
||||
return json.loads(inputs["workflow_json"])
|
||||
return ComfyUIClient.load_workflow(Path(inputs["workflow_path"]))
|
||||
|
||||
@staticmethod
|
||||
def _model_name(inputs: dict[str, Any], custom_workflow: bool) -> str:
|
||||
if not custom_workflow:
|
||||
return "ace-step-v1-3.5b"
|
||||
return (
|
||||
inputs.get("workflow_model")
|
||||
or inputs.get("model")
|
||||
or inputs.get("workflow_name")
|
||||
or "custom-comfyui-workflow"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _workflow_provenance(
|
||||
inputs: dict[str, Any],
|
||||
custom_workflow: bool,
|
||||
output_node: str,
|
||||
workflow: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
if not custom_workflow:
|
||||
return {
|
||||
"source": "bundled",
|
||||
"workflow": "ace-step-1-t2a.json",
|
||||
"workflow_hash_sha256": workflow_hash(workflow),
|
||||
"model_stack": model_stack("ace-step-1-t2a", inputs),
|
||||
"output_node": output_node,
|
||||
}
|
||||
stack = inputs.get("workflow_model_stack")
|
||||
return {
|
||||
"source": "user_supplied",
|
||||
"workflow_name": inputs.get("workflow_name"),
|
||||
"workflow_path": inputs.get("workflow_path"),
|
||||
"model": inputs.get("workflow_model") or inputs.get("model"),
|
||||
"workflow_hash_sha256": workflow_hash(workflow),
|
||||
"model_stack": stack if isinstance(stack, list) else [],
|
||||
"model_stack_source": "caller_supplied" if stack else "unknown_custom_workflow",
|
||||
"output_node": output_node,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _probe_duration(path: Path) -> float | None:
|
||||
"""Best-effort track duration via ffprobe; None if unavailable."""
|
||||
if shutil.which("ffprobe") is None:
|
||||
return None
|
||||
try:
|
||||
out = subprocess.run(
|
||||
[
|
||||
"ffprobe", "-v", "error",
|
||||
"-show_entries", "format=duration",
|
||||
"-of", "default=noprint_wrappers=1:nokey=1",
|
||||
str(path),
|
||||
],
|
||||
capture_output=True, text=True, timeout=15, check=True,
|
||||
)
|
||||
value = out.stdout.strip()
|
||||
return round(float(value), 2) if value else None
|
||||
except (subprocess.SubprocessError, ValueError):
|
||||
return None
|
||||
@@ -58,7 +58,9 @@ class ComfyUIImage(BaseTool):
|
||||
install_instructions = (
|
||||
"Start a ComfyUI server and set COMFYUI_SERVER_URL "
|
||||
"(default http://localhost:8188).\n"
|
||||
"See https://github.com/comfyanonymous/ComfyUI for setup."
|
||||
"See https://github.com/comfyanonymous/ComfyUI for setup.\n"
|
||||
"Running a separate ComfyUI instance for images? Set COMFYUI_IMAGE_SERVER_URL "
|
||||
"instead -- it takes priority over COMFYUI_SERVER_URL for this tool only."
|
||||
)
|
||||
agent_skills = ["comfyui", "flux-best-practices"]
|
||||
|
||||
@@ -133,7 +135,7 @@ class ComfyUIImage(BaseTool):
|
||||
user_visible_verification = ["Inspect generated image for quality and prompt adherence"]
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._client = ComfyUIClient()
|
||||
self._client = ComfyUIClient(capability="image")
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
if not self._client.is_available():
|
||||
|
||||
@@ -107,7 +107,9 @@ class ComfyUIVideo(BaseTool):
|
||||
install_instructions = (
|
||||
"Start a ComfyUI server and set COMFYUI_SERVER_URL "
|
||||
"(default http://localhost:8188).\n"
|
||||
"Requires WAN 2.2 models and LightX2V LoRAs in ComfyUI's model directory."
|
||||
"Requires WAN 2.2 models and LightX2V LoRAs in ComfyUI's model directory.\n"
|
||||
"Running a separate ComfyUI instance for video? Set COMFYUI_VIDEO_SERVER_URL "
|
||||
"instead -- it takes priority over COMFYUI_SERVER_URL for this tool only."
|
||||
)
|
||||
agent_skills = ["comfyui", "ai-video-gen", "ltx2"]
|
||||
|
||||
@@ -186,6 +188,24 @@ class ComfyUIVideo(BaseTool):
|
||||
),
|
||||
"items": {"type": "object"},
|
||||
},
|
||||
"timeout_seconds": {
|
||||
"type": "integer",
|
||||
"description": (
|
||||
"How long to wait for the ComfyUI job to finish before giving up. "
|
||||
"Default 3600s (1hr) covers slow/local GPUs and non-accelerated "
|
||||
"custom workflows; raise it further for large frame counts or "
|
||||
"high resolutions. On timeout the job is NOT cancelled server-side "
|
||||
"and the error's data.prompt_id can be passed back via "
|
||||
"resume_prompt_id to keep waiting without resubmitting."
|
||||
),
|
||||
},
|
||||
"resume_prompt_id": {
|
||||
"type": "string",
|
||||
"description": (
|
||||
"A prompt_id from a previous timed-out call (see error data on "
|
||||
"timeout). Skips resubmission and just resumes waiting/downloading."
|
||||
),
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
@@ -198,7 +218,24 @@ class ComfyUIVideo(BaseTool):
|
||||
user_visible_verification = ["Watch generated clip for motion coherence and artifacts"]
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._client = ComfyUIClient()
|
||||
self._client = ComfyUIClient(capability="video")
|
||||
self._last_progress_log = 0.0
|
||||
|
||||
def _log_progress(self, data: dict) -> None:
|
||||
"""Print a throttled progress line for long video renders.
|
||||
|
||||
Video jobs can run for tens of minutes; without this the process
|
||||
looks hung. Throttled to once per 10s since ComfyUI pushes a
|
||||
``progress`` event per sampling step, which would otherwise flood
|
||||
stdout on fast GPUs.
|
||||
"""
|
||||
now = time.monotonic()
|
||||
if now - self._last_progress_log < 10:
|
||||
return
|
||||
self._last_progress_log = now
|
||||
value, max_value = data.get("value"), data.get("max")
|
||||
if value is not None and max_value:
|
||||
print(f"[comfyui_video] step {value}/{max_value}")
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
if not self._client.is_available():
|
||||
@@ -320,12 +357,25 @@ class ComfyUIVideo(BaseTool):
|
||||
workflow,
|
||||
output_node=output_node,
|
||||
dest=output_path,
|
||||
timeout=900,
|
||||
timeout=inputs.get("timeout_seconds", 3600),
|
||||
interval=10,
|
||||
resume_prompt_id=inputs.get("resume_prompt_id"),
|
||||
on_progress=self._log_progress,
|
||||
)
|
||||
|
||||
except ComfyUIError as exc:
|
||||
return ToolResult(success=False, error=str(exc))
|
||||
data = {"prompt_id": exc.prompt_id} if exc.prompt_id else {}
|
||||
if exc.prompt_id:
|
||||
error_msg = (
|
||||
f"{exc}\n\nThis job was NOT cancelled and is very likely still "
|
||||
f"running server-side. To recover it without resubmitting, call "
|
||||
f"execute() again with resume_prompt_id={exc.prompt_id!r} "
|
||||
f"(and a longer timeout_seconds if it needs more time), or poll "
|
||||
f"GET {{COMFYUI_SERVER_URL}}/history/{exc.prompt_id} directly."
|
||||
)
|
||||
else:
|
||||
error_msg = str(exc)
|
||||
return ToolResult(success=False, error=error_msg, data=data)
|
||||
except Exception as exc:
|
||||
return ToolResult(success=False, error=f"ComfyUI video generation failed: {exc}")
|
||||
|
||||
|
||||
Reference in New Issue
Block a user