From ceee7c7d5654a250c086927e4e9ac59a3e653efd Mon Sep 17 00:00:00 2001 From: Ntsako Date: Thu, 6 Aug 2026 10:55:56 +0200 Subject: [PATCH 1/6] fix: raise ComfyUI video timeout default and add job resume support Non-accelerated local-GPU workflows (e.g. Wan 1.3B at 832x480/81-97 frames) routinely took ~1360-1630s, so the old 900s default false-failed real renders that were still completing server-side. Timeout is now a configurable timeout_seconds input (default 3600s), and ComfyUIError carries the prompt_id on error/timeout so a timed-out-but-still-running job can be resumed via resume_prompt_id instead of resubmitted. Co-Authored-By: Claude Sonnet 5 --- docs/comfyui-adapter-plan.md | 19 ++++- tests/contracts/test_comfyui_tools.py | 112 ++++++++++++++++++++++++++ tools/_comfyui/client.py | 33 ++++++-- tools/video/comfyui_video.py | 34 +++++++- 4 files changed, 189 insertions(+), 9 deletions(-) diff --git a/docs/comfyui-adapter-plan.md b/docs/comfyui-adapter-plan.md index d60fb89d1..c0a941e60 100644 --- a/docs/comfyui-adapter-plan.md +++ b/docs/comfyui-adapter-plan.md @@ -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 diff --git a/tests/contracts/test_comfyui_tools.py b/tests/contracts/test_comfyui_tools.py index a8ee8304d..ff3e279b7 100644 --- a/tests/contracts/test_comfyui_tools.py +++ b/tests/contracts/test_comfyui_tools.py @@ -282,6 +282,47 @@ def fake_download(filename, subfolder, dest, folder_type="output"): "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 + + client = ComfyUIClient("http://comfy.test") + + def fail_submit(workflow): + raise AssertionError("submit() should not be called when resuming") + + monkeypatch.setattr(client, "submit", fail_submit) + 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_is_default_url_when_env_not_set(self, monkeypatch): from tools._comfyui.client import ComfyUIClient monkeypatch.delenv("COMFYUI_SERVER_URL", raising=False) @@ -418,6 +459,77 @@ def test_custom_workflow_accepts_model_stack_provenance(self, tmp_path): 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 diff --git a/tools/_comfyui/client.py b/tools/_comfyui/client.py index fa80e8434..41fee64be 100644 --- a/tools/_comfyui/client.py +++ b/tools/_comfyui/client.py @@ -18,7 +18,17 @@ 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: @@ -174,11 +184,18 @@ def poll( status = entry.get("status", {}) if status.get("status_str") == "error": msgs = status.get("messages", []) - raise ComfyUIError(f"Execution error: {msgs}") + raise ComfyUIError(f"Execution error: {msgs}", prompt_id=prompt_id) 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 download( @@ -229,9 +246,15 @@ def generate( *, timeout: int = 600, interval: int = 5, + resume_prompt_id: str | None = None, ) -> list[Path]: - """Submit → poll → download. Returns list of artifact paths.""" - prompt_id = self.submit(workflow) + """Submit → poll → 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. + """ + prompt_id = resume_prompt_id or self.submit(workflow) entry = self.poll(prompt_id, timeout=timeout, interval=interval) outputs = entry.get("outputs", {}) diff --git a/tools/video/comfyui_video.py b/tools/video/comfyui_video.py index 409264f5a..5bab999d5 100644 --- a/tools/video/comfyui_video.py +++ b/tools/video/comfyui_video.py @@ -186,6 +186,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." + ), + }, }, } @@ -320,12 +338,24 @@ def execute(self, inputs: dict[str, Any]) -> ToolResult: 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"), ) 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}") From ca203e49b713ab9eb45877e04be0cff100b98eb3 Mon Sep 17 00:00:00 2001 From: Ntsako Date: Thu, 6 Aug 2026 12:31:24 +0200 Subject: [PATCH 2/6] feat: wait on ComfyUI websocket feed instead of polling for completion Resolves the "async generation" open question from the adapter plan. generate() now watches ComfyUI's websocket events (executing/progress/ execution_error) and reacts immediately instead of sleeping between REST polls, with an optional on_progress callback that comfyui_video uses to print step progress on long renders. websocket-client is an optional import; _wait() falls back to the original poll() loop (with the remaining time budget, not a fresh one) when it's unavailable or the connection drops, so resume_prompt_id recovery is unaffected either way. Co-Authored-By: Claude Sonnet 5 --- docs/comfyui-adapter-plan.md | 13 +- tests/contracts/test_comfyui_tools.py | 181 ++++++++++++++++++++++++++ tools/_comfyui/client.py | 131 ++++++++++++++++++- tools/video/comfyui_video.py | 18 +++ 4 files changed, 336 insertions(+), 7 deletions(-) diff --git a/docs/comfyui-adapter-plan.md b/docs/comfyui-adapter-plan.md index c0a941e60..c8b7f79e0 100644 --- a/docs/comfyui-adapter-plan.md +++ b/docs/comfyui-adapter-plan.md @@ -474,9 +474,16 @@ 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? diff --git a/tests/contracts/test_comfyui_tools.py b/tests/contracts/test_comfyui_tools.py index ff3e279b7..21dc5166f 100644 --- a/tests/contracts/test_comfyui_tools.py +++ b/tests/contracts/test_comfyui_tools.py @@ -351,6 +351,187 @@ def test_unavailable_reason_custom_url(self, monkeypatch): 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 _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_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) + monkeypatch.setattr( + "tools._comfyui.client.requests.get", + lambda *a, **k: type("R", (), { + "raise_for_status": lambda self: None, + "json": lambda self: {"p1": {"outputs": {"9": {}}}}, + })(), + ) + + 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) diff --git a/tools/_comfyui/client.py b/tools/_comfyui/client.py index 41fee64be..85b694211 100644 --- a/tools/_comfyui/client.py +++ b/tools/_comfyui/client.py @@ -11,8 +11,9 @@ import os import random import time +import uuid from pathlib import Path -from typing import Any +from typing import Any, Callable import requests @@ -46,6 +47,9 @@ def __init__(self, server_url: str | None = None) -> None: server_url or os.environ.get("COMFYUI_SERVER_URL", "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()) # ------------------------------------------------------------------ # Health @@ -147,7 +151,7 @@ def submit(self, workflow: dict) -> str: """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: @@ -198,6 +202,117 @@ def poll( prompt_id=prompt_id, ) + 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 + + 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 + while time.time() < deadline: + try: + raw = conn.recv() + except websocket.WebSocketTimeoutException: + 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: + 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, + ) + + 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: + 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, @@ -247,15 +362,23 @@ def generate( 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. + """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) - entry = self.poll(prompt_id, timeout=timeout, interval=interval) + entry = self._wait( + prompt_id, timeout=timeout, interval=interval, on_progress=on_progress + ) outputs = entry.get("outputs", {}) node_output = outputs.get(output_node, {}) diff --git a/tools/video/comfyui_video.py b/tools/video/comfyui_video.py index 5bab999d5..84270c27e 100644 --- a/tools/video/comfyui_video.py +++ b/tools/video/comfyui_video.py @@ -217,6 +217,23 @@ class ComfyUIVideo(BaseTool): def __init__(self) -> None: self._client = ComfyUIClient() + 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(): @@ -341,6 +358,7 @@ def execute(self, inputs: dict[str, Any]) -> ToolResult: 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: From 2f114682e8897e3956a03a9926e47eed17955cfb Mon Sep 17 00:00:00 2001 From: Ntsako Date: Thu, 6 Aug 2026 12:49:57 +0200 Subject: [PATCH 3/6] feat: support per-capability ComfyUI server URLs for image/video MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Resolves the "multi-server" open question from the adapter plan. ComfyUIClient(capability="image"|"video") now resolves its server URL from COMFYUI_IMAGE_SERVER_URL / COMFYUI_VIDEO_SERVER_URL first, falling back to the shared COMFYUI_SERVER_URL and then the localhost default — so comfyui_image and comfyui_video can point at separate ComfyUI instances (different GPUs, different model sets) with zero extra config for single-server setups. is_default_url/unavailable_reason() and the setup_offer metadata account for the override. Co-Authored-By: Claude Sonnet 5 --- docs/comfyui-adapter-plan.md | 24 +++++++- tests/contracts/test_comfyui_tools.py | 79 +++++++++++++++++++++++++++ tools/_comfyui/client.py | 59 +++++++++++++------- tools/_comfyui/metadata.py | 7 +++ tools/graphics/comfyui_image.py | 6 +- tools/video/comfyui_video.py | 6 +- 6 files changed, 156 insertions(+), 25 deletions(-) diff --git a/docs/comfyui-adapter-plan.md b/docs/comfyui-adapter-plan.md index c8b7f79e0..c8eef2e36 100644 --- a/docs/comfyui-adapter-plan.md +++ b/docs/comfyui-adapter-plan.md @@ -373,6 +373,17 @@ COMFYUI_POLL_TIMEOUT=600 # max wait for image gen COMFYUI_VIDEO_TIMEOUT=900 # max wait for video gen ``` +**Multi-server (optional):** point `comfyui_image` and `comfyui_video` at +separate ComfyUI instances -- e.g. one GPU running FLUX 2, another running +WAN 2.2 -- by setting a per-capability override. Each takes priority over +`COMFYUI_SERVER_URL` for its own tool only; leave both unset and everything +still talks to the single shared server. + +```bash +COMFYUI_IMAGE_SERVER_URL=http://gpu-a:8188 +COMFYUI_VIDEO_SERVER_URL=http://gpu-b:8188 +``` + **For Docker Compose setups** (ComfyUI in a container): ```bash @@ -485,8 +496,17 @@ pipeline definition, or any schema. `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")` + resolves its server URL from a per-capability env var first + (`COMFYUI_IMAGE_SERVER_URL` / `COMFYUI_VIDEO_SERVER_URL`), then the shared + `COMFYUI_SERVER_URL`, then the `http://localhost:8188` default. `comfyui_image` + and `comfyui_video` pass their capability at construction, so image and video + generation can point at different ComfyUI instances (different GPUs, different + model sets) with zero code changes -- single-server setups need no extra + configuration since both 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`. diff --git a/tests/contracts/test_comfyui_tools.py b/tests/contracts/test_comfyui_tools.py index 21dc5166f..44c451c0f 100644 --- a/tests/contracts/test_comfyui_tools.py +++ b/tests/contracts/test_comfyui_tools.py @@ -369,6 +369,85 @@ def fake_post(url, json=None, timeout=None): 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 diff --git a/tools/_comfyui/client.py b/tools/_comfyui/client.py index 85b694211..6632bfa6a 100644 --- a/tools/_comfyui/client.py +++ b/tools/_comfyui/client.py @@ -42,41 +42,52 @@ 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 # ------------------------------------------------------------------ @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") - - def is_available(self) -> bool: - """Return True if the ComfyUI server is reachable.""" - try: - resp = requests.get( - f"{self.server_url}/system_stats", timeout=5 - ) - return resp.status_code == 200 - except Exception: - return False + """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 COMFYUI_SERVER_URL configured).\n" - f"Set COMFYUI_SERVER_URL in your .env file to the address of " + 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 ( @@ -84,6 +95,16 @@ def unavailable_reason(self) -> str: f"Check that ComfyUI is running and the URL is correct." ) + def is_available(self) -> bool: + """Return True if the ComfyUI server is reachable.""" + try: + resp = requests.get( + f"{self.server_url}/system_stats", timeout=5 + ) + return resp.status_code == 200 + except Exception: + return False + # ------------------------------------------------------------------ # Model discovery # ------------------------------------------------------------------ diff --git a/tools/_comfyui/metadata.py b/tools/_comfyui/metadata.py index dcca514b7..757078680 100644 --- a/tools/_comfyui/metadata.py +++ b/tools/_comfyui/metadata.py @@ -18,6 +18,13 @@ "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", + }, } diff --git a/tools/graphics/comfyui_image.py b/tools/graphics/comfyui_image.py index 4c55d94b9..e91e71c33 100644 --- a/tools/graphics/comfyui_image.py +++ b/tools/graphics/comfyui_image.py @@ -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(): diff --git a/tools/video/comfyui_video.py b/tools/video/comfyui_video.py index 84270c27e..85ed3717e 100644 --- a/tools/video/comfyui_video.py +++ b/tools/video/comfyui_video.py @@ -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"] @@ -216,7 +218,7 @@ 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: From 172acca6ea4ee86cb0f48d63cd863a81a4e5117b Mon Sep 17 00:00:00 2001 From: Ntsako Date: Thu, 6 Aug 2026 14:19:05 +0200 Subject: [PATCH 4/6] feat: ship comfyui_music as a custom-workflow-only ACE-Step tool Resolves the "music generation" open question from the adapter plan. Unlike comfyui_image/comfyui_video there is no bundled workflow: ACE-Step's ComfyUI node interface isn't standardized across custom node packs (AceStepModelLoader vs native TextEncodeAceStepAudio, etc.), so instead of picking one pack and breaking for everyone else, comfyui_music always requires a caller-supplied workflow_json/workflow_path + output_node -- the same override contract image/video offer as an alternative, just mandatory here. prompt is provenance-only, never injected into the graph. Routed through the existing registry.get_by_capability("music_generation") path alongside suno_music/music_gen -- no dedicated selector needed. ComfyUIClient.generate() now also reads the "audio" output key (what ComfyUI's native SaveAudio node writes), and gets timeout/resume/websocket- wait/multi-server support for free via the shared client. Duration is a best-effort ffprobe probe of the downloaded file since a custom workflow gives no other way to know it ahead of time. Co-Authored-By: Claude Sonnet 5 --- .agents/skills/comfyui/SKILL.md | 15 +- docs/comfyui-adapter-plan.md | 90 +++++--- tests/contracts/test_comfyui_tools.py | 156 +++++++++++++- tools/_comfyui/client.py | 9 +- tools/_comfyui/metadata.py | 1 + tools/audio/comfyui_music.py | 300 ++++++++++++++++++++++++++ 6 files changed, 535 insertions(+), 36 deletions(-) create mode 100644 tools/audio/comfyui_music.py diff --git a/.agents/skills/comfyui/SKILL.md b/.agents/skills/comfyui/SKILL.md index 64950870f..ea89cc3e7 100644 --- a/.agents/skills/comfyui/SKILL.md +++ b/.agents/skills/comfyui/SKILL.md @@ -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,11 @@ 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`) + +- Unlike `comfyui_image`/`comfyui_video`, there is **no bundled workflow**. ACE-Step's ComfyUI node interface isn't standardized across custom node packs (`AceStepModelLoader` vs native `TextEncodeAceStepAudio`, etc.), so `workflow_json`/`workflow_path` + `output_node` are always required, not optional. +- `prompt` is provenance/logging only -- it is never injected into the workflow. Bake the actual tags/lyrics into the workflow JSON yourself before calling, the same way you would patch a custom image/video workflow. +- `output_node` should be the node that writes the final audio, typically ComfyUI's native `SaveAudio`. The client reads artifacts from that node's `"audio"` output key (parallel to `"images"` for image/video savers). +- Provide `workflow_name`/`workflow_model`/`workflow_model_stack` for provenance exactly as you would for a custom image/video workflow -- there's no bundled model stack to fall back on here. diff --git a/docs/comfyui-adapter-plan.md b/docs/comfyui-adapter-plan.md index c8eef2e36..c3d3884fc 100644 --- a/docs/comfyui-adapter-plan.md +++ b/docs/comfyui-adapter-plan.md @@ -296,19 +296,44 @@ not promote ComfyUI for an operation whose bundled models are missing. --- -### `comfyui_music` -- Music Generation (not shipped) - -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. - -**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. +### `comfyui_music` -- Music Generation (shipped, custom-workflow-only) + +`tools/audio/comfyui_music.py`. `capability="music_generation"`, `provider="comfyui"`. +Ships with **no bundled workflow** -- the ACE-Step node-pack fragmentation +described below is real and unsolved, so instead of picking one pack and +breaking for everyone else, the tool always requires a caller-supplied +`workflow_json`/`workflow_path` + `output_node`, exactly like the image/video +tools' *optional* override path, just mandatory here. `prompt` is accepted +for provenance/logging only and is never injected into the workflow -- +tags/lyrics must already be baked into the graph before calling, same +convention as image/video custom workflows. + +Originally not shipped because: 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. The custom-workflow-only design sidesteps this +entirely: whichever node pack is installed, the caller exports it themselves. + +**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 `duration_seconds` -- a best-effort `ffprobe` probe of the downloaded +file (`None` if `ffprobe` isn't on PATH), since a custom workflow gives no +other reliable way to know actual output duration ahead of time. + +**Workflow/output-node contract:** identical to image/video -- `output_node` +must be the ID of the node that writes the final artifact (typically ComfyUI's +native `SaveAudio` node). `ComfyUIClient.generate()`'s artifact extraction now +also checks the `"audio"` output key (previously only `"images"`/`"gifs"`), +which is what `SaveAudio` writes to in ComfyUI's `/history` response -- +this is the one part of the contract that *is* standardized regardless of +which ACE-Step loader pack sits upstream of it. --- @@ -373,15 +398,17 @@ COMFYUI_POLL_TIMEOUT=600 # max wait for image gen COMFYUI_VIDEO_TIMEOUT=900 # max wait for video gen ``` -**Multi-server (optional):** point `comfyui_image` and `comfyui_video` at -separate ComfyUI instances -- e.g. one GPU running FLUX 2, another running -WAN 2.2 -- by setting a per-capability override. Each takes priority over -`COMFYUI_SERVER_URL` for its own tool only; leave both unset and everything -still talks to the single shared server. +**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): @@ -496,20 +523,23 @@ pipeline definition, or any schema. `poll()` REST loop when it isn't installed or the connection fails — `resume_prompt_id` recovery behaves identically either way. -3. ~~**Multi-server:**~~ **Resolved.** `ComfyUIClient(capability="image"|"video")` +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`), then the shared - `COMFYUI_SERVER_URL`, then the `http://localhost:8188` default. `comfyui_image` - and `comfyui_video` pass their capability at construction, so image and video - generation can point at different ComfyUI instances (different GPUs, different - model sets) with zero code changes -- single-server setups need no extra - configuration since both env vars are optional. `client.capability`/ + (`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 as custom-workflow-only.** + `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`. It has no bundled workflow -- the node-pack + fragmentation that originally blocked this is real, so the tool always + requires caller-supplied `workflow_json`/`workflow_path` + `output_node` + rather than betting on one pack. See the `comfyui_music` section above for + the artifact schema and workflow/output-node contract. diff --git a/tests/contracts/test_comfyui_tools.py b/tests/contracts/test_comfyui_tools.py index 44c451c0f..27e32f7c7 100644 --- a/tests/contracts/test_comfyui_tools.py +++ b/tests/contracts/test_comfyui_tools.py @@ -17,13 +17,14 @@ 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 @@ -323,6 +324,23 @@ def fail_submit(workflow): ) 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) @@ -838,6 +856,142 @@ def test_bundled_workflow_provenance_records_hash_and_stack(self, tmp_path): 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_requires_workflow_json_or_path(self): + tool = ComfyUIMusic() + tool._client.is_available = lambda: True + + result = tool.execute({"prompt": "ambient pad", "output_node": "9"}) + + assert result.success is False + assert "workflow_json" in result.error or "workflow_path" in result.error + + def test_requires_output_node(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_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): diff --git a/tools/_comfyui/client.py b/tools/_comfyui/client.py index 6632bfa6a..e280bea6d 100644 --- a/tools/_comfyui/client.py +++ b/tools/_comfyui/client.py @@ -404,8 +404,13 @@ def generate( 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}. " diff --git a/tools/_comfyui/metadata.py b/tools/_comfyui/metadata.py index 757078680..3c0ee13bb 100644 --- a/tools/_comfyui/metadata.py +++ b/tools/_comfyui/metadata.py @@ -24,6 +24,7 @@ "per_capability_env_var_overrides": { "comfyui_image": "COMFYUI_IMAGE_SERVER_URL", "comfyui_video": "COMFYUI_VIDEO_SERVER_URL", + "comfyui_music": "COMFYUI_MUSIC_SERVER_URL", }, } diff --git a/tools/audio/comfyui_music.py b/tools/audio/comfyui_music.py new file mode 100644 index 000000000..1d006f77b --- /dev/null +++ b/tools/audio/comfyui_music.py @@ -0,0 +1,300 @@ +"""ComfyUI music generation via a local or remote ComfyUI server. + +No bundled workflow: ACE-Step's ComfyUI node interface is not standardized +across custom node packs (``AceStepModelLoader`` vs native +``TextEncodeAceStepAudio``, etc.), so a hardcoded template would break for +most installs. This tool always runs a caller-supplied ``workflow_json`` or +``workflow_path`` -- the same override contract ``comfyui_image``/ +``comfyui_video`` offer as an alternative to their bundled workflow, just +mandatory here instead of optional. See the ``comfyui`` skill for how to +convert a community ACE-Step workflow into a call. +""" + +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 COMFYUI_SETUP_OFFER, workflow_hash + + +class ComfyUIMusic(BaseTool): + name = "comfyui_music" + version = "0.1.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 with ACE-Step installed (any node pack) and " + "set COMFYUI_SERVER_URL (default http://localhost:8188).\n" + "There is no bundled workflow for this tool -- export your ACE-Step " + "graph in API format and pass it as workflow_json/workflow_path.\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, + "custom_workflow": True, + "custom_output_node": True, + "offline": True, + } + best_for = [ + "local GPU music generation without API costs, using whatever ACE-Step node pack is installed", + "full control over sampling via custom ComfyUI workflows", + ] + not_good_for = [ + "setups without a running ComfyUI server", + "quick generation without first exporting/adapting an ACE-Step workflow", + "CPU-only machines", + ] + fallback_tools = ["suno_music", "music_gen"] + + input_schema = { + "type": "object", + "required": ["prompt", "output_node"], + "properties": { + "prompt": { + "type": "string", + "description": ( + "Description of the desired music, for provenance/logging only. " + "Not injected into the workflow -- bake the actual tags/lyrics " + "into workflow_json/workflow_path before calling." + ), + }, + "seed": {"type": "integer", "description": "Random if omitted"}, + "output_path": {"type": "string", "description": "Where to save the audio"}, + "workflow_json": { + "type": "string", + "description": "Full ComfyUI ACE-Step workflow JSON (API format). Required if workflow_path is omitted.", + }, + "workflow_path": { + "type": "string", + "description": "Path to a ComfyUI ACE-Step workflow JSON file. Required if workflow_json is omitted.", + }, + "output_node": { + "type": "string", + "description": "ComfyUI output node ID (e.g. the SaveAudio node) to download the artifact from.", + }, + "workflow_name": { + "type": "string", + "description": "Optional human-readable provenance label for the workflow.", + }, + "workflow_model": { + "type": "string", + "description": "Optional model/provenance label (e.g. 'ace-step-v1-3.5b').", + }, + "workflow_model_stack": { + "type": "array", + "description": ( + "Optional provenance metadata for 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", "seed", "workflow_json", "workflow_path", "output_node"] + 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 + return ToolStatus.AVAILABLE + + def estimate_cost(self, inputs: dict[str, Any]) -> float: + return 0.0 + + def estimate_runtime(self, inputs: dict[str, Any]) -> float: + # Actual runtime depends entirely on the caller's custom workflow + # (steps, duration, sampler); this is a conservative flat estimate. + return 180.0 + + def get_info(self) -> dict[str, Any]: + info = super().get_info() + info["setup_offer"] = self.setup_offer + info["bundled_workflow"] = None + info["custom_workflow_required"] = True + 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: + if not (inputs.get("workflow_json") or inputs.get("workflow_path")): + return ToolResult( + success=False, + error=( + "comfyui_music requires workflow_json or workflow_path -- there " + "is no bundled default. ACE-Step's ComfyUI node interface isn't " + "standardized across custom node packs, so a hardcoded template " + "would break for most installs. Export the ACE-Step workflow " + "you actually have installed (API format) and pass it in." + ), + ) + + if not inputs.get("output_node"): + return ToolResult( + success=False, + error="output_node is required so OpenMontage knows which ComfyUI node to download the audio from.", + ) + + if not self._client.is_available(): + return ToolResult(success=False, error=self._client.unavailable_reason()) + + start = time.time() + seed = inputs.get("seed") or ComfyUIClient.random_seed() + output_path = Path(inputs.get("output_path", f"comfyui_music_{seed}.mp3")) + output_node = str(inputs["output_node"]) + + try: + workflow = self._load_custom_workflow(inputs) + provenance = self._workflow_provenance(inputs, 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) + return ToolResult( + success=True, + data={ + "provider": "comfyui", + "model": model_name, + "prompt": inputs["prompt"], + "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]) -> str: + 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], output_node: str, workflow: dict[str, Any] + ) -> dict[str, Any]: + 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 From cf1a13a722dffe26fab1050eac0b0cd79ebb7055 Mon Sep 17 00:00:00 2001 From: Ntsako Date: Thu, 6 Aug 2026 15:10:48 +0200 Subject: [PATCH 5/6] feat: bundle a native-node ACE-Step v1 workflow for comfyui_music ACE-Step v1's node-pack fragmentation turns out to be moot: 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 repo has an official ACE-Step-v1 template built from those plus long-stable core nodes. tools/_comfyui/workflows/ace-step-1-t2a.json was built by cross-checking every node's class_type and input names against ComfyUI's own source (nodes_ace.py, nodes_audio.py, nodes_latent.py, nodes.py) rather than trusting the UI-format export directly. comfyui_music now defaults to this bundled workflow: prompt maps to ACE-Step's tags field (matching suno_music's "prompt = music description" convention), lyrics/duration_seconds/steps/cfg/lyrics_strength/seed are all patchable, and missing ace_step_v1_3.5b.safetensors surfaces through the same missing_models contract as image/video. workflow_json/workflow_path + output_node remains available for ACE-Step 1.5, other node packs, or different audio models entirely. Co-Authored-By: Claude Sonnet 5 --- .agents/skills/comfyui/SKILL.md | 10 +- docs/comfyui-adapter-plan.md | 71 ++++---- tests/contracts/test_comfyui_tools.py | 69 +++++++- tools/_comfyui/metadata.py | 11 ++ tools/_comfyui/workflows/ace-step-1-t2a.json | 80 +++++++++ tools/audio/comfyui_music.py | 167 +++++++++++++------ 6 files changed, 317 insertions(+), 91 deletions(-) create mode 100644 tools/_comfyui/workflows/ace-step-1-t2a.json diff --git a/.agents/skills/comfyui/SKILL.md b/.agents/skills/comfyui/SKILL.md index ea89cc3e7..05bec8c29 100644 --- a/.agents/skills/comfyui/SKILL.md +++ b/.agents/skills/comfyui/SKILL.md @@ -58,7 +58,9 @@ Use this skill before calling `comfyui_image`, `comfyui_video`, or `comfyui_musi ## Music (`comfyui_music`) -- Unlike `comfyui_image`/`comfyui_video`, there is **no bundled workflow**. ACE-Step's ComfyUI node interface isn't standardized across custom node packs (`AceStepModelLoader` vs native `TextEncodeAceStepAudio`, etc.), so `workflow_json`/`workflow_path` + `output_node` are always required, not optional. -- `prompt` is provenance/logging only -- it is never injected into the workflow. Bake the actual tags/lyrics into the workflow JSON yourself before calling, the same way you would patch a custom image/video workflow. -- `output_node` should be the node that writes the final audio, typically ComfyUI's native `SaveAudio`. The client reads artifacts from that node's `"audio"` output key (parallel to `"images"` for image/video savers). -- Provide `workflow_name`/`workflow_model`/`workflow_model_stack` for provenance exactly as you would for a custom image/video workflow -- there's no bundled model stack to fall back on here. +- 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. diff --git a/docs/comfyui-adapter-plan.md b/docs/comfyui-adapter-plan.md index c3d3884fc..0e29d6019 100644 --- a/docs/comfyui-adapter-plan.md +++ b/docs/comfyui-adapter-plan.md @@ -296,24 +296,33 @@ not promote ComfyUI for an operation whose bundled models are missing. --- -### `comfyui_music` -- Music Generation (shipped, custom-workflow-only) +### `comfyui_music` -- Music Generation (shipped, with a native-node bundled workflow) `tools/audio/comfyui_music.py`. `capability="music_generation"`, `provider="comfyui"`. -Ships with **no bundled workflow** -- the ACE-Step node-pack fragmentation -described below is real and unsolved, so instead of picking one pack and -breaking for everyone else, the tool always requires a caller-supplied -`workflow_json`/`workflow_path` + `output_node`, exactly like the image/video -tools' *optional* override path, just mandatory here. `prompt` is accepted -for provenance/logging only and is never injected into the workflow -- -tags/lyrics must already be baked into the graph before calling, same -convention as image/video custom workflows. - -Originally not shipped because: 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. The custom-workflow-only design sidesteps this -entirely: whichever node pack is installed, the caller exports it themselves. + +**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 @@ -323,17 +332,16 @@ 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 `duration_seconds` -- a best-effort `ffprobe` probe of the downloaded -file (`None` if `ffprobe` isn't on PATH), since a custom workflow gives no -other reliable way to know actual output duration ahead of time. +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 (typically ComfyUI's -native `SaveAudio` node). `ComfyUIClient.generate()`'s artifact extraction now -also checks the `"audio"` output key (previously only `"images"`/`"gifs"`), -which is what `SaveAudio` writes to in ComfyUI's `/history` response -- -this is the one part of the contract that *is* standardized regardless of -which ACE-Step loader pack sits upstream of it. +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. --- @@ -535,11 +543,12 @@ pipeline definition, or any schema. override, and `COMFYUI_SETUP_OFFER.per_capability_env_var_overrides` documents it for the setup-offer surfacing in `provider_menu()`. -4. ~~**Music generation:**~~ **Resolved -- shipped as custom-workflow-only.** +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`. It has no bundled workflow -- the node-pack - fragmentation that originally blocked this is real, so the tool always - requires caller-supplied `workflow_json`/`workflow_path` + `output_node` - rather than betting on one pack. See the `comfyui_music` section above for - the artifact schema and workflow/output-node contract. + 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. diff --git a/tests/contracts/test_comfyui_tools.py b/tests/contracts/test_comfyui_tools.py index 27e32f7c7..ac533f30b 100644 --- a/tests/contracts/test_comfyui_tools.py +++ b/tests/contracts/test_comfyui_tools.py @@ -137,6 +137,7 @@ def test_custom_workflow_requires_output_node(self, cls): "flux2-txt2img.json", "wan22-i2v-4step.json", "wan22-t2v-4step.json", + "ace-step-1-t2a.json", ] @@ -651,6 +652,11 @@ def test_video_tool_has_required_models_t2v(self): 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 @@ -863,16 +869,23 @@ def test_capability_and_provider(self): assert tool.capability == "music_generation" assert tool.provider == "comfyui" - def test_requires_workflow_json_or_path(self): + 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_node": "9"}) + result = tool.execute({ + "prompt": "ambient pad", + "output_path": str(tmp_path / "music.mp3"), + }) - assert result.success is False - assert "workflow_json" in result.error or "workflow_path" in result.error + assert result.success is True + assert result.data["workflow_provenance"]["source"] == "bundled" - def test_requires_output_node(self): + def test_custom_workflow_without_output_node_errors(self): tool = ComfyUIMusic() tool._client.is_available = lambda: True @@ -884,6 +897,52 @@ def test_requires_output_node(self): 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_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 diff --git a/tools/_comfyui/metadata.py b/tools/_comfyui/metadata.py index 3c0ee13bb..084fbeb3e 100644 --- a/tools/_comfyui/metadata.py +++ b/tools/_comfyui/metadata.py @@ -184,6 +184,17 @@ ), }, ], + "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" + ), + }, + ], } diff --git a/tools/_comfyui/workflows/ace-step-1-t2a.json b/tools/_comfyui/workflows/ace-step-1-t2a.json new file mode 100644 index 000000000..fdeba4053 --- /dev/null +++ b/tools/_comfyui/workflows/ace-step-1-t2a.json @@ -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" + } + } +} diff --git a/tools/audio/comfyui_music.py b/tools/audio/comfyui_music.py index 1d006f77b..9751ad469 100644 --- a/tools/audio/comfyui_music.py +++ b/tools/audio/comfyui_music.py @@ -1,13 +1,11 @@ """ComfyUI music generation via a local or remote ComfyUI server. -No bundled workflow: ACE-Step's ComfyUI node interface is not standardized -across custom node packs (``AceStepModelLoader`` vs native -``TextEncodeAceStepAudio``, etc.), so a hardcoded template would break for -most installs. This tool always runs a caller-supplied ``workflow_json`` or -``workflow_path`` -- the same override contract ``comfyui_image``/ -``comfyui_video`` offer as an alternative to their bundled workflow, just -mandatory here instead of optional. See the ``comfyui`` skill for how to -convert a community ACE-Step workflow into a call. +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 @@ -32,12 +30,23 @@ ToolTier, ) from tools._comfyui.client import ComfyUIClient, ComfyUIError -from tools._comfyui.metadata import COMFYUI_SETUP_OFFER, workflow_hash +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.1.0" + version = "0.2.0" tier = ToolTier.GENERATE capability = "music_generation" provider = "comfyui" @@ -49,10 +58,10 @@ class ComfyUIMusic(BaseTool): dependencies = [] # checked at runtime via server health setup_offer = COMFYUI_SETUP_OFFER install_instructions = ( - "Start a ComfyUI server with ACE-Step installed (any node pack) and " - "set COMFYUI_SERVER_URL (default http://localhost:8188).\n" - "There is no bundled workflow for this tool -- export your ACE-Step " - "graph in API format and pass it as workflow_json/workflow_path.\n" + "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." @@ -62,59 +71,73 @@ class ComfyUIMusic(BaseTool): 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, using whatever ACE-Step node pack is installed", - "full control over sampling via custom ComfyUI workflows", + "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", - "quick generation without first exporting/adapting an ACE-Step workflow", "CPU-only machines", ] fallback_tools = ["suno_music", "music_gen"] input_schema = { "type": "object", - "required": ["prompt", "output_node"], + "required": ["prompt"], "properties": { "prompt": { "type": "string", "description": ( - "Description of the desired music, for provenance/logging only. " - "Not injected into the workflow -- bake the actual tags/lyrics " - "into workflow_json/workflow_path before calling." + "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": "Full ComfyUI ACE-Step workflow JSON (API format). Required if workflow_path is omitted.", + "description": "Optional full ComfyUI workflow JSON. Requires output_node.", }, "workflow_path": { "type": "string", - "description": "Path to a ComfyUI ACE-Step workflow JSON file. Required if workflow_json is omitted.", + "description": "Optional path to a ComfyUI workflow JSON file. Requires output_node.", }, "output_node": { "type": "string", - "description": "ComfyUI output node ID (e.g. the SaveAudio node) to download the artifact from.", + "description": "ComfyUI output node ID for custom workflow_json/workflow_path.", }, "workflow_name": { "type": "string", - "description": "Optional human-readable provenance label for the workflow.", + "description": "Optional human-readable provenance label for a custom workflow.", }, "workflow_model": { "type": "string", - "description": "Optional model/provenance label (e.g. 'ace-step-v1-3.5b').", + "description": "Optional model/provenance label for a custom workflow.", }, "workflow_model_stack": { "type": "array", "description": ( - "Optional provenance metadata for workflow dependencies. " + "Optional provenance metadata for custom workflow dependencies. " "Items should include name, role, and node-pack origin when known." ), "items": {"type": "object"}, @@ -134,7 +157,7 @@ class ComfyUIMusic(BaseTool): 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", "seed", "workflow_json", "workflow_path", "output_node"] + 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"] @@ -145,21 +168,21 @@ def __init__(self) -> None: 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: - # Actual runtime depends entirely on the caller's custom workflow - # (steps, duration, sampler); this is a conservative flat estimate. - return 180.0 + 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_workflow"] = None - info["custom_workflow_required"] = True + info["bundled_model_stack"] = BUNDLED_MODEL_STACKS["ace-step-1-t2a"] return info def _log_progress(self, data: dict) -> None: @@ -173,35 +196,63 @@ def _log_progress(self, data: dict) -> None: print(f"[comfyui_music] step {value}/{max_value}") def execute(self, inputs: dict[str, Any]) -> ToolResult: - if not (inputs.get("workflow_json") or inputs.get("workflow_path")): + 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=( - "comfyui_music requires workflow_json or workflow_path -- there " - "is no bundled default. ACE-Step's ComfyUI node interface isn't " - "standardized across custom node packs, so a hardcoded template " - "would break for most installs. Export the ACE-Step workflow " - "you actually have installed (API format) and pass it in." + "Custom ComfyUI workflows require output_node so OpenMontage " + "knows which ComfyUI node to download artifacts from." ), ) - if not inputs.get("output_node"): - return ToolResult( - success=False, - error="output_node is required so OpenMontage knows which ComfyUI node to download the audio 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") or ComfyUIClient.random_seed() output_path = Path(inputs.get("output_path", f"comfyui_music_{seed}.mp3")) - output_node = str(inputs["output_node"]) try: - workflow = self._load_custom_workflow(inputs) - provenance = self._workflow_provenance(inputs, output_node, workflow) + 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, @@ -229,13 +280,14 @@ def execute(self, inputs: dict[str, Any]) -> ToolResult: return ToolResult(success=False, error=f"ComfyUI music generation failed: {exc}") duration = self._probe_duration(paths[0]) - model_name = self._model_name(inputs) + 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("."), @@ -255,7 +307,9 @@ def _load_custom_workflow(inputs: dict[str, Any]) -> dict: return ComfyUIClient.load_workflow(Path(inputs["workflow_path"])) @staticmethod - def _model_name(inputs: dict[str, Any]) -> str: + 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") @@ -265,8 +319,19 @@ def _model_name(inputs: dict[str, Any]) -> str: @staticmethod def _workflow_provenance( - inputs: dict[str, Any], output_node: str, workflow: dict[str, Any] + 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", From 41793226b5d0ee5932fc39b12efac2d269bafc2d Mon Sep 17 00:00:00 2001 From: calesthio Date: Thu, 13 Aug 2026 09:37:02 -0700 Subject: [PATCH 6/6] fix: make ComfyUI history authoritative --- tests/contracts/test_comfyui_tools.py | 56 +++++++++++++++++++++- tools/_comfyui/client.py | 69 ++++++++++++++++++++------- tools/audio/comfyui_music.py | 4 +- 3 files changed, 110 insertions(+), 19 deletions(-) diff --git a/tests/contracts/test_comfyui_tools.py b/tests/contracts/test_comfyui_tools.py index ac533f30b..cf678a85c 100644 --- a/tests/contracts/test_comfyui_tools.py +++ b/tests/contracts/test_comfyui_tools.py @@ -305,6 +305,7 @@ def test_poll_timeout_carries_prompt_id_for_recovery(self, monkeypatch): 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") @@ -312,6 +313,14 @@ 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", @@ -502,6 +511,29 @@ def _install_fake_websocket(monkeypatch, frames): 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 @@ -516,11 +548,12 @@ def test_wait_ws_completes_on_executing_none_node(self, monkeypatch, tmp_path): }}), ] _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: {"p1": {"outputs": {"9": {}}}}, + "json": lambda self: next(history_calls), })(), ) @@ -937,6 +970,27 @@ def fake_generate(workflow, output_node, dest, **kwargs): 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 diff --git a/tools/_comfyui/client.py b/tools/_comfyui/client.py index e280bea6d..e37bc2b43 100644 --- a/tools/_comfyui/client.py +++ b/tools/_comfyui/client.py @@ -199,17 +199,8 @@ def poll( """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}", prompt_id=prompt_id) + entry = self._history_entry(prompt_id) + if entry is not None: return entry time.sleep(interval) raise ComfyUIError( @@ -223,6 +214,28 @@ def poll( 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, @@ -250,6 +263,12 @@ def wait_ws( """ 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 ) @@ -260,10 +279,19 @@ def wait_ws( 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 @@ -289,6 +317,9 @@ def wait_ws( 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 " @@ -296,9 +327,7 @@ def wait_ws( prompt_id=prompt_id, ) - resp = requests.get(f"{self.server_url}/history/{prompt_id}", timeout=10) - resp.raise_for_status() - entry = resp.json().get(prompt_id) + entry = self._history_entry(prompt_id) if entry is None: raise ComfyUIError( f"No history entry for {prompt_id} after completion", @@ -397,9 +426,15 @@ def generate( the connection can't be used. See :meth:`_wait`. """ prompt_id = resume_prompt_id or self.submit(workflow) - entry = self._wait( - prompt_id, timeout=timeout, interval=interval, on_progress=on_progress - ) + 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, {}) diff --git a/tools/audio/comfyui_music.py b/tools/audio/comfyui_music.py index 9751ad469..05dfddf0d 100644 --- a/tools/audio/comfyui_music.py +++ b/tools/audio/comfyui_music.py @@ -227,7 +227,9 @@ def execute(self, inputs: dict[str, Any]) -> ToolResult: ) start = time.time() - seed = inputs.get("seed") or ComfyUIClient.random_seed() + 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: