From 6ec2bbb0909ee9ab6b5492f9e58390bbed483680 Mon Sep 17 00:00:00 2001 From: martimramos Date: Thu, 16 Apr 2026 23:59:33 +0100 Subject: [PATCH] comfyui: add native ComfyUI provider for image, video, and music generation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds three new BaseTool providers that delegate GPU work to a running ComfyUI server via its REST API. This avoids the need to install PyTorch/diffusers directly, which is critical on hardware where the ecosystem hasn't caught up (e.g. NVIDIA Blackwell / DGX Spark, aarch64 + CUDA 13.0). New files: - tools/_comfyui/client.py — shared REST client (submit/poll/download) - tools/_comfyui/workflows/ — 4 bundled workflow templates - tools/graphics/comfyui_image.py — FLUX 2 Dev NVFP4 text-to-image - tools/video/comfyui_video.py — WAN 2.2 14B t2v + i2v (4-step LightX2V) - tools/audio/comfyui_music.py — ACE-Step 3.5B music generation - tests/contracts/test_comfyui_tools.py — 41 contract tests - docs/comfyui-adapter-plan.md — design document Zero changes to existing tools, selectors, registry, or pipelines. Tools are auto-discovered and selectors pick them up via capability match. Co-Authored-By: Claude Opus 4.6 (1M context) --- docs/comfyui-adapter-plan.md | 385 ++++++++++++++++++ tests/contracts/test_comfyui_tools.py | 173 ++++++++ tools/_comfyui/__init__.py | 1 + tools/_comfyui/client.py | 207 ++++++++++ tools/_comfyui/workflows/ace-step-music.json | 27 ++ tools/_comfyui/workflows/flux2-txt2img.json | 96 +++++ tools/_comfyui/workflows/wan22-i2v-4step.json | 154 +++++++ tools/_comfyui/workflows/wan22-t2v-4step.json | 143 +++++++ tools/audio/comfyui_music.py | 192 +++++++++ tools/graphics/comfyui_image.py | 165 ++++++++ tools/video/comfyui_video.py | 250 ++++++++++++ 11 files changed, 1793 insertions(+) create mode 100644 docs/comfyui-adapter-plan.md create mode 100644 tests/contracts/test_comfyui_tools.py create mode 100644 tools/_comfyui/__init__.py create mode 100644 tools/_comfyui/client.py create mode 100644 tools/_comfyui/workflows/ace-step-music.json create mode 100644 tools/_comfyui/workflows/flux2-txt2img.json create mode 100644 tools/_comfyui/workflows/wan22-i2v-4step.json create mode 100644 tools/_comfyui/workflows/wan22-t2v-4step.json create mode 100644 tools/audio/comfyui_music.py create mode 100644 tools/graphics/comfyui_image.py create mode 100644 tools/video/comfyui_video.py diff --git a/docs/comfyui-adapter-plan.md b/docs/comfyui-adapter-plan.md new file mode 100644 index 00000000..126a9308 --- /dev/null +++ b/docs/comfyui-adapter-plan.md @@ -0,0 +1,385 @@ +# ComfyUI Provider Adapter for OpenMontage + +**RFC: Native ComfyUI backend for image, video, and music generation** + +--- + +## Motivation + +OpenMontage's local GPU tools (`wan_video`, `hunyuan_video`, `cogvideo_video`, +`local_diffusion`) use HuggingFace `diffusers` directly. This works on x86 + +consumer GPUs but breaks on newer hardware where the PyTorch ecosystem hasn't +caught up: + +| Issue | Detail | +|-------|--------| +| **NVIDIA Blackwell (sm_121)** | No stable PyTorch wheels for aarch64 + CUDA 13.0. Requires NGC containers or nightly builds. | +| **Flash Attention** | Does not support sm_121. Must be replaced with SageAttention v3 or native SDPA. | +| **Unified Memory (GB10/DGX Spark)** | `nvidia-smi` cannot report VRAM. Diffusers' memory estimation breaks. | +| **Model format mismatch** | Diffusers expects HF repos. Production deployments use `.safetensors` checkpoints with quantized variants (NVFP4, FP8) that diffusers doesn't natively load. | + +ComfyUI already solves all of these. NVIDIA ships official ComfyUI containers +for DGX Spark. The community has optimized workflows for Blackwell (SageAttention, +NVFP4 quantization, LightX2V 4-step LoRAs). Models like WAN 2.2, FLUX 2, +and ACE-Step run reliably through ComfyUI on hardware where diffusers cannot. + +A ComfyUI adapter gives OpenMontage access to any model ComfyUI supports, +on any hardware ComfyUI runs on, without shipping or maintaining PyTorch builds. + +--- + +## Design + +### Architecture + +``` +OpenMontage Agent + | + v +video_selector / image_selector / music_selector + | + v +comfyui_video comfyui_image comfyui_music (new tools) + | | | + v v v +ComfyUI REST API (POST /prompt, GET /history, GET /view) + | + v +GPU (any hardware ComfyUI supports) +``` + +### Integration model + +Three new `BaseTool` subclasses plus one shared client library: + +``` +tools/ + _comfyui/ + __init__.py + client.py # Shared ComfyUI REST client + workflows/ # Bundled workflow templates + flux2-txt2img.json + wan22-t2v-4step.json + wan22-i2v-4step.json + ace-step-music.json + graphics/ + comfyui_image.py # capability="image_generation", provider="comfyui" + video/ + comfyui_video.py # capability="video_generation", provider="comfyui" + audio/ + comfyui_music.py # capability="music_generation", provider="comfyui" +``` + +### Zero changes to selectors or registry + +The tools declare `capability` and `provider` as class attributes. +`tool_registry.discover()` picks them up automatically via `pkgutil.walk_packages`. +`video_selector`, `image_selector`, and `music_selector` find them via +`registry.get_by_capability()` -- no hardcoded references needed. + +--- + +## Shared Client: `tools/_comfyui/client.py` + +Encapsulates the ComfyUI REST API pattern proven in production (used by the +Bard project's Airflow DAGs for thousands of generations): + +```python +class ComfyUIClient: + """Thin client for the ComfyUI REST API.""" + + def __init__(self, server_url: str | None = None): + self.server_url = server_url or os.environ.get( + "COMFYUI_SERVER_URL", "http://localhost:8188" + ) + + def is_available(self) -> bool: + """Health check -- can we reach the server?""" + + def submit(self, workflow: dict) -> str: + """POST /prompt. Returns prompt_id. Raises on node_errors.""" + + def poll(self, prompt_id: str, timeout: int = 600, interval: int = 5) -> dict: + """GET /history/{prompt_id} until complete. Returns outputs dict.""" + + def download(self, filename: str, subfolder: str, dest: Path) -> Path: + """GET /view?filename=...&type=output. Writes bytes to dest.""" + + def upload_image(self, local_path: Path, name: str) -> str: + """POST /upload/image. Returns server-side filename for LoadImage nodes.""" + + def generate(self, workflow: dict, output_node: str, dest: Path, + timeout: int = 600) -> Path: + """Full cycle: submit -> poll -> download. Returns artifact path.""" +``` + +**Why a shared client?** The submit/poll/download cycle is identical across +image, video, and music generation. The only differences are: which workflow +template, which nodes to customize, and which output node to read from. + +--- + +## Tool Specifications + +### `comfyui_image` -- Image Generation + +| Field | Value | +|-------|-------| +| capability | `image_generation` | +| provider | `comfyui` | +| runtime | `LOCAL_GPU` | +| tier | `GENERATE` | +| stability | `EXPERIMENTAL` | +| capabilities | `text_to_image`, `image_to_image` | +| dependencies | (runtime: ComfyUI server reachable) | +| fallback_tools | `flux_image`, `local_diffusion`, `openai_image` | +| cost | `$0.00` (local compute) | + +**Bundled workflow:** `flux2-txt2img.json` + +Loads FLUX 2 Dev (NVFP4) with Mistral text encoder. Templated nodes: + +| Node | Class | Templated field | +|------|-------|-----------------| +| 4 | CLIPTextEncode | `text` (prompt) | +| 6 | EmptyFlux2LatentImage | `width`, `height` | +| 7 | RandomNoise | `noise_seed` | +| 10 | Flux2Scheduler | `steps` | +| 13 | SaveImage | `filename_prefix` | + +**Input schema:** + +```yaml +prompt: string # required +width: integer # default 1024 +height: integer # default 1024 +steps: integer # default 20 +seed: integer # optional (random if omitted) +guidance: number # default 3.5 +output_path: string # where to save the image +workflow_json: string # optional override (full custom workflow) +``` + +**get_status():** Pings ComfyUI server. Returns `AVAILABLE` if reachable, `UNAVAILABLE` otherwise. + +**execute() flow:** +1. Deep-copy workflow template +2. Inject prompt, seed, dimensions, steps into templated nodes +3. `client.generate(workflow, output_node="13", dest=output_path)` +4. Return `ToolResult` with artifact path, seed, model info + +--- + +### `comfyui_video` -- Video Generation + +| Field | Value | +|-------|-------| +| capability | `video_generation` | +| provider | `comfyui` | +| runtime | `LOCAL_GPU` | +| tier | `GENERATE` | +| stability | `EXPERIMENTAL` | +| capabilities | `text_to_video`, `image_to_video` | +| dependencies | (runtime: ComfyUI server reachable) | +| fallback_tools | `wan_video`, `hunyuan_video`, `ltx_video_local` | +| cost | `$0.00` (local compute) | + +**Bundled workflows:** + +1. **`wan22-i2v-4step.json`** -- Image-to-video (WAN 2.2 14B, fp8, 4-step LightX2V LoRA) +2. **`wan22-t2v-4step.json`** -- Text-to-video (WAN 2.2 14B, fp8, 4-step LightX2V LoRA) + +**I2V workflow -- templated nodes:** + +| Node | Class | Templated field | +|------|-------|-----------------| +| 93 | CLIPTextEncode | `text` (positive prompt) | +| 97 | LoadImage | `image` (server filename from upload) | +| 98 | WanImageToVideo | `width`, `height`, `length` | +| 86 | KSamplerAdvanced | `noise_seed` | +| 108 | SaveVideo | `filename_prefix` | + +**Input schema:** + +```yaml +prompt: string # required +operation: string # "text_to_video" | "image_to_video" (default: t2v) +reference_image_path: string # local path (for i2v) +reference_image_url: string # URL (for i2v, downloaded first) +width: integer # default 640 +height: integer # default 640 +num_frames: integer # default 81 (5s at 16fps) +seed: integer # optional +output_path: string # where to save the video +workflow_json: string # optional override +``` + +**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)` +5. Return `ToolResult` + +**execute() flow (t2v):** +1. Deep-copy t2v workflow template +2. Inject prompt, seed, dimensions +3. `client.generate(workflow, output_node="108", dest=output_path, timeout=900)` +4. Return `ToolResult` + +--- + +### `comfyui_music` -- Music Generation + +| Field | Value | +|-------|-------| +| capability | `music_generation` | +| provider | `comfyui` | +| runtime | `LOCAL_GPU` | +| tier | `GENERATE` | +| stability | `EXPERIMENTAL` | +| capabilities | `text_to_music` | +| dependencies | (runtime: ComfyUI server reachable + ACE-Step model) | +| fallback_tools | `suno_music`, `elevenlabs_music` | +| cost | `$0.00` (local compute) | + +**Bundled workflow:** `ace-step-music.json` + +Uses ACE-Step v1 3.5B for text-to-music generation. Workflow to be authored +based on the ComfyUI ACE-Step custom node. + +**Input schema:** + +```yaml +prompt: string # required (music description) +duration: number # seconds (default 30) +seed: integer # optional +output_path: string # where to save the audio +``` + +--- + +## Workflow Override Mechanism + +Every tool accepts an optional `workflow_json` input. When provided, it +replaces the bundled template entirely. This enables: + +- Using newer model checkpoints without code changes +- Custom sampling strategies (different schedulers, step counts, LoRAs) +- Community workflows dropped in as-is +- A/B testing different generation approaches + +The agent can also read workflow files from `tools/_comfyui/workflows/` and +modify them programmatically before passing to `execute()`. + +--- + +## Configuration + +**Environment variables:** + +```bash +# .env +COMFYUI_SERVER_URL=http://localhost:8188 # ComfyUI API endpoint +COMFYUI_POLL_INTERVAL=5 # seconds between status checks +COMFYUI_POLL_TIMEOUT=600 # max wait for image gen +COMFYUI_VIDEO_TIMEOUT=900 # max wait for video gen +``` + +**For Docker Compose setups** (ComfyUI in a container): + +```bash +COMFYUI_SERVER_URL=http://host.docker.internal:8188 +# or +COMFYUI_SERVER_URL=http://comfyui:8188 # if on same docker network +``` + +--- + +## Provider Selection Behavior + +When the adapter is available, selectors will rank it alongside other providers +using OpenMontage's 7-dimension scoring: + +| Dimension | ComfyUI score | Rationale | +|-----------|---------------|-----------| +| Task fit | High | Supports t2i, i2v, t2v, music | +| Quality | High | Latest models (FLUX 2, WAN 2.2 14B) | +| Control | Highest | Full workflow customization | +| Reliability | High | Proven in production | +| Cost | $0 | Local compute | +| Latency | Medium | GPU-bound, no network round-trip | +| Continuity | High | Deterministic with seeds | + +When ComfyUI is unavailable (server down), the selector falls through to +`fallback_tools` automatically -- API providers like FLUX via fal.ai or +HeyGen take over transparently. + +--- + +## What This Unlocks + +### Immediate (with existing models) + +- **FLUX 2 Dev NVFP4** image generation -- Blackwell-optimized, ~60s per image +- **WAN 2.2 14B** i2v with 4-step acceleration -- ~3.5 min per 5s clip +- **WAN 2.2 14B** t2v (models downloaded, workflow needed) +- **ACE-Step 3.5B** local music generation (model downloaded, workflow needed) + +### Future (add models to ComfyUI, no code changes to OpenMontage) + +- Newer checkpoints (WAN 3.x, FLUX 3, etc.) -- just update workflow JSON +- ControlNet, IP-Adapter, AnimateDiff -- supported via ComfyUI custom nodes +- Upscaling, inpainting, outpainting -- ComfyUI nodes exist +- Any model the ComfyUI ecosystem supports + +### Hardware portability + +The same adapter works on: +- NVIDIA DGX Spark (GB10, aarch64, CUDA 13.0) +- Consumer GPUs (RTX 3090/4090, x86) +- Cloud instances (A100, H100) +- Multi-GPU setups (ComfyUI handles device placement) + +No PyTorch version pinning, no architecture-specific wheels, no CUDA +compatibility matrices. ComfyUI is the abstraction layer. + +--- + +## Implementation Scope + +| Component | Files | Estimated size | +|-----------|-------|----------------| +| Shared client | `tools/_comfyui/client.py` | ~120 lines | +| Image tool | `tools/graphics/comfyui_image.py` | ~130 lines | +| Video tool | `tools/video/comfyui_video.py` | ~160 lines | +| Music tool | `tools/audio/comfyui_music.py` | ~100 lines | +| Workflow templates | `tools/_comfyui/workflows/*.json` | 4 files | +| T2V workflow | `tools/_comfyui/workflows/wan22-t2v-4step.json` | 1 file (to author) | +| Music workflow | `tools/_comfyui/workflows/ace-step-music.json` | 1 file (to author) | +| Tests | `tests/contracts/test_comfyui_*.py` | ~80 lines | +| Docs | `skills/creative/comfyui-workflows.md` | Agent skill file | + +**Total:** ~600 lines of Python + 4-6 workflow JSONs. + +No changes to: `base_tool.py`, `tool_registry.py`, any selector, any +existing tool, any pipeline definition, or any schema. + +--- + +## Open Questions + +1. **Workflow versioning:** Should workflow JSONs live in the repo or be + user-provided via a config directory? Bundling gives reproducibility; + external gives flexibility. + +2. **Model discovery:** ComfyUI has a `/object_info` endpoint that lists + available nodes and models. Should `get_status()` also report which + models are loaded, so the selector can make informed routing decisions? + +3. **Async generation:** ComfyUI supports websocket connections for real-time + progress. Worth implementing for long video generations, or is polling + sufficient? + +4. **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 new file mode 100644 index 00000000..3c68fdc1 --- /dev/null +++ b/tests/contracts/test_comfyui_tools.py @@ -0,0 +1,173 @@ +"""Contract tests for ComfyUI provider tools. + +These tests verify that the tools satisfy the BaseTool contract without +requiring a running ComfyUI server. They check class attributes, +schemas, status reporting, and cost estimates. +""" + +import json +from pathlib import Path + +import pytest + +from tools.base_tool import ( + BaseTool, + ToolRuntime, + ToolStability, + ToolStatus, + ToolTier, +) +from tools.graphics.comfyui_image import ComfyUIImage +from tools.video.comfyui_video import ComfyUIVideo +from tools.audio.comfyui_music import ComfyUIMusic + +TOOLS = [ComfyUIImage, ComfyUIVideo, ComfyUIMusic] +WORKFLOW_DIR = Path(__file__).resolve().parent.parent.parent / "tools" / "_comfyui" / "workflows" + + +# ------------------------------------------------------------------ +# Contract compliance +# ------------------------------------------------------------------ + +@pytest.mark.parametrize("cls", TOOLS, ids=lambda c: c.name) +class TestContract: + + def test_inherits_base_tool(self, cls): + assert issubclass(cls, BaseTool) + + def test_has_required_identity(self, cls): + tool = cls() + assert tool.name + assert tool.version + assert tool.capability + assert tool.provider == "comfyui" + assert tool.tier == ToolTier.GENERATE + assert tool.stability == ToolStability.EXPERIMENTAL + assert tool.runtime == ToolRuntime.LOCAL_GPU + + def test_has_input_schema(self, cls): + tool = cls() + schema = tool.input_schema + assert schema.get("type") == "object" + assert "prompt" in schema.get("properties", {}) + assert "prompt" in schema.get("required", []) + + def test_has_capabilities(self, cls): + tool = cls() + assert len(tool.capabilities) > 0 + + def test_has_fallbacks(self, cls): + tool = cls() + assert tool.fallback or tool.fallback_tools + + def test_cost_is_zero(self, cls): + tool = cls() + assert tool.estimate_cost({"prompt": "test"}) == 0.0 + + def test_runtime_estimate_positive(self, cls): + tool = cls() + assert tool.estimate_runtime({"prompt": "test"}) > 0 + + def test_get_info_returns_dict(self, cls): + tool = cls() + info = tool.get_info() + assert isinstance(info, dict) + assert info["name"] == tool.name + assert info["provider"] == "comfyui" + assert info["runtime"] == "local_gpu" + + def test_status_unavailable_without_server(self, cls): + """Without a running server, status should be UNAVAILABLE.""" + tool = cls() + # Point to a port that's almost certainly not running ComfyUI + tool._client.server_url = "http://127.0.0.1:19999" + assert tool.get_status() == ToolStatus.UNAVAILABLE + + def test_idempotency_key_fields(self, cls): + tool = cls() + assert len(tool.idempotency_key_fields) > 0 + assert "prompt" in tool.idempotency_key_fields + + +# ------------------------------------------------------------------ +# Workflow files +# ------------------------------------------------------------------ + +EXPECTED_WORKFLOWS = [ + "flux2-txt2img.json", + "wan22-i2v-4step.json", + "wan22-t2v-4step.json", + "ace-step-music.json", +] + + +@pytest.mark.parametrize("filename", EXPECTED_WORKFLOWS) +def test_workflow_exists_and_valid_json(filename): + path = WORKFLOW_DIR / filename + assert path.exists(), f"Missing workflow: {path}" + with open(path) as f: + data = json.load(f) + assert isinstance(data, dict) + assert len(data) > 0 + + +def test_flux2_workflow_has_templated_nodes(): + with open(WORKFLOW_DIR / "flux2-txt2img.json") as f: + w = json.load(f) + assert "4" in w # CLIPTextEncode (prompt) + assert "7" in w # RandomNoise (seed) + assert "13" in w # SaveImage (output) + + +def test_i2v_workflow_has_templated_nodes(): + with open(WORKFLOW_DIR / "wan22-i2v-4step.json") as f: + w = json.load(f) + assert "93" in w # CLIPTextEncode (prompt) + assert "97" in w # LoadImage (reference) + assert "86" in w # KSamplerAdvanced (seed) + assert "108" in w # SaveVideo (output) + + +def test_t2v_workflow_has_templated_nodes(): + with open(WORKFLOW_DIR / "wan22-t2v-4step.json") as f: + w = json.load(f) + assert "2" in w # CLIPTextEncode (prompt) + assert "12" in w # KSamplerAdvanced (seed) + assert "16" in w # SaveVideo (output) + + +# ------------------------------------------------------------------ +# Client unit tests +# ------------------------------------------------------------------ + +class TestClientHelpers: + + def test_load_workflow(self): + from tools._comfyui.client import ComfyUIClient + w = ComfyUIClient.load_workflow(WORKFLOW_DIR / "flux2-txt2img.json") + assert isinstance(w, dict) + assert "1" in w + + def test_patch_workflow(self): + from tools._comfyui.client import ComfyUIClient + w = ComfyUIClient.load_workflow(WORKFLOW_DIR / "flux2-txt2img.json") + patched = ComfyUIClient.patch_workflow(w, { + "4": {"text": "hello world"}, + "7": {"noise_seed": 123}, + }) + assert patched["4"]["inputs"]["text"] == "hello world" + assert patched["7"]["inputs"]["noise_seed"] == 123 + # Original unchanged + assert w["4"]["inputs"]["text"] == "" + + def test_patch_workflow_bad_node(self): + from tools._comfyui.client import ComfyUIClient, ComfyUIError + w = {"1": {"inputs": {"x": 1}}} + with pytest.raises(ComfyUIError, match="not found"): + ComfyUIClient.patch_workflow(w, {"99": {"x": 2}}) + + def test_random_seed_range(self): + from tools._comfyui.client import ComfyUIClient + for _ in range(100): + s = ComfyUIClient.random_seed() + assert 0 <= s < 2**32 diff --git a/tools/_comfyui/__init__.py b/tools/_comfyui/__init__.py new file mode 100644 index 00000000..fa7beab6 --- /dev/null +++ b/tools/_comfyui/__init__.py @@ -0,0 +1 @@ +"""ComfyUI integration — shared client and bundled workflow templates.""" diff --git a/tools/_comfyui/client.py b/tools/_comfyui/client.py new file mode 100644 index 00000000..62f7ca0d --- /dev/null +++ b/tools/_comfyui/client.py @@ -0,0 +1,207 @@ +"""Thin REST client for a running ComfyUI server. + +Handles the full generation cycle: submit workflow, poll for completion, +download artifacts. Used by comfyui_image, comfyui_video, and comfyui_music. +""" + +from __future__ import annotations + +import copy +import json +import os +import random +import time +from pathlib import Path +from typing import Any + +import requests + + +class ComfyUIError(Exception): + """Raised when ComfyUI returns an error or times out.""" + + +class ComfyUIClient: + """Client for the ComfyUI REST API. + + The protocol is simple and battle-tested: + 1. POST /prompt → queue a workflow, get a prompt_id + 2. GET /history/{id} → poll until outputs appear + 3. GET /view?filename=… → download the generated artifact + 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("/") + + # ------------------------------------------------------------------ + # Health + # ------------------------------------------------------------------ + + 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 + + # ------------------------------------------------------------------ + # Core cycle + # ------------------------------------------------------------------ + + 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}, + timeout=30, + ) + resp.raise_for_status() + data = resp.json() + if data.get("node_errors"): + raise ComfyUIError(f"Node errors: {json.dumps(data['node_errors'])}") + prompt_id = data.get("prompt_id") + if not prompt_id: + raise ComfyUIError(f"No prompt_id in response: {data}") + return prompt_id + + def poll( + self, + prompt_id: str, + *, + timeout: int = 600, + interval: int = 5, + ) -> dict: + """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}") + return entry + time.sleep(interval) + raise ComfyUIError( + f"Prompt {prompt_id} did not complete within {timeout}s" + ) + + def download( + self, + filename: str, + subfolder: str, + dest: Path, + ) -> Path: + """Download an output artifact from the ComfyUI server.""" + resp = requests.get( + f"{self.server_url}/view", + params={ + "filename": filename, + "subfolder": subfolder, + "type": "output", + }, + timeout=120, + ) + resp.raise_for_status() + dest.parent.mkdir(parents=True, exist_ok=True) + dest.write_bytes(resp.content) + return dest + + def upload_image(self, local_path: Path, name: str) -> str: + """Upload a local image so it can be referenced by LoadImage nodes. + + Returns the server-side filename. + """ + with open(local_path, "rb") as f: + resp = requests.post( + f"{self.server_url}/upload/image", + files={"image": (name, f, "image/png")}, + timeout=30, + ) + resp.raise_for_status() + return resp.json()["name"] + + # ------------------------------------------------------------------ + # High-level helper + # ------------------------------------------------------------------ + + def generate( + self, + workflow: dict, + output_node: str, + dest: Path, + *, + timeout: int = 600, + interval: int = 5, + ) -> list[Path]: + """Submit → poll → download. Returns list of artifact paths.""" + prompt_id = self.submit(workflow) + entry = self.poll(prompt_id, timeout=timeout, interval=interval) + + 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", []) + if not items: + raise ComfyUIError( + f"No output artifacts on node {output_node}. " + f"Available nodes: {list(outputs.keys())}" + ) + + paths: list[Path] = [] + for i, item in enumerate(items): + suffix = Path(item["filename"]).suffix + if len(items) == 1: + target = dest + else: + target = dest.with_stem(f"{dest.stem}_{i:03d}").with_suffix(suffix) + self.download(item["filename"], item.get("subfolder", ""), target) + paths.append(target) + return paths + + # ------------------------------------------------------------------ + # Workflow helpers + # ------------------------------------------------------------------ + + @staticmethod + def load_workflow(path: Path) -> dict: + """Load a workflow JSON template from disk.""" + with open(path) as f: + return json.load(f) + + @staticmethod + def patch_workflow( + workflow: dict, patches: dict[str, dict[str, Any]] + ) -> dict: + """Deep-copy *workflow* and apply *patches*. + + *patches* maps ``node_id`` → ``{input_name: value, ...}``. + """ + w = copy.deepcopy(workflow) + for node_id, values in patches.items(): + if node_id not in w: + raise ComfyUIError( + f"Node {node_id!r} not found in workflow. " + f"Available: {list(w.keys())}" + ) + for key, val in values.items(): + w[node_id]["inputs"][key] = val + return w + + @staticmethod + def random_seed() -> int: + """Return a random seed suitable for ComfyUI noise nodes.""" + return random.randint(0, 2**32 - 1) diff --git a/tools/_comfyui/workflows/ace-step-music.json b/tools/_comfyui/workflows/ace-step-music.json new file mode 100644 index 00000000..5f1af02b --- /dev/null +++ b/tools/_comfyui/workflows/ace-step-music.json @@ -0,0 +1,27 @@ +{ + "1": { + "class_type": "AceStepModelLoader", + "inputs": { + "model": "ace_step_v1_3.5b.safetensors" + } + }, + "2": { + "class_type": "AceStepSampler", + "inputs": { + "model": ["1", 0], + "prompt": "", + "lyrics": "", + "duration": 30.0, + "seed": 42, + "steps": 60, + "cfg": 3.0 + } + }, + "3": { + "class_type": "SaveAudio", + "inputs": { + "audio": ["2", 0], + "filename_prefix": "openmontage_music" + } + } +} diff --git a/tools/_comfyui/workflows/flux2-txt2img.json b/tools/_comfyui/workflows/flux2-txt2img.json new file mode 100644 index 00000000..4dbb8981 --- /dev/null +++ b/tools/_comfyui/workflows/flux2-txt2img.json @@ -0,0 +1,96 @@ +{ + "1": { + "class_type": "UNETLoader", + "inputs": { + "unet_name": "flux2-dev-nvfp4.safetensors", + "weight_dtype": "default" + } + }, + "2": { + "class_type": "CLIPLoader", + "inputs": { + "clip_name": "mistral_3_small_flux2_fp4_mixed.safetensors", + "type": "flux2", + "device": "cpu" + } + }, + "3": { + "class_type": "VAELoader", + "inputs": { + "vae_name": "flux2-vae.safetensors" + } + }, + "4": { + "class_type": "CLIPTextEncode", + "inputs": { + "clip": ["2", 0], + "text": "" + } + }, + "5": { + "class_type": "FluxGuidance", + "inputs": { + "conditioning": ["4", 0], + "guidance": 3.5 + } + }, + "6": { + "class_type": "EmptyFlux2LatentImage", + "inputs": { + "width": 1024, + "height": 1024, + "batch_size": 1 + } + }, + "7": { + "class_type": "RandomNoise", + "inputs": { + "noise_seed": 42 + } + }, + "8": { + "class_type": "BasicGuider", + "inputs": { + "model": ["1", 0], + "conditioning": ["5", 0] + } + }, + "9": { + "class_type": "KSamplerSelect", + "inputs": { + "sampler_name": "euler" + } + }, + "10": { + "class_type": "Flux2Scheduler", + "inputs": { + "steps": 20, + "width": 1024, + "height": 1024 + } + }, + "11": { + "class_type": "SamplerCustomAdvanced", + "inputs": { + "noise": ["7", 0], + "guider": ["8", 0], + "sampler": ["9", 0], + "sigmas": ["10", 0], + "latent_image": ["6", 0] + } + }, + "12": { + "class_type": "VAEDecode", + "inputs": { + "samples": ["11", 0], + "vae": ["3", 0] + } + }, + "13": { + "class_type": "SaveImage", + "inputs": { + "images": ["12", 0], + "filename_prefix": "openmontage" + } + } +} diff --git a/tools/_comfyui/workflows/wan22-i2v-4step.json b/tools/_comfyui/workflows/wan22-i2v-4step.json new file mode 100644 index 00000000..83313a67 --- /dev/null +++ b/tools/_comfyui/workflows/wan22-i2v-4step.json @@ -0,0 +1,154 @@ +{ + "84": { + "class_type": "CLIPLoader", + "inputs": { + "clip_name": "umt5_xxl_fp8_e4m3fn_scaled.safetensors", + "type": "wan", + "device": "default" + } + }, + "89": { + "class_type": "CLIPTextEncode", + "inputs": { + "clip": ["84", 0], + "text": "oversaturated, overexposed, static, blurry details, subtitles, style, artwork, painting, still frame, gray overall, worst quality, low quality, JPEG artifacts, ugly, deformed, extra fingers, poorly drawn hands, poorly drawn face, deformed limbs, fused fingers, static frame, cluttered background, three legs, many people in background, walking backwards" + } + }, + "90": { + "class_type": "VAELoader", + "inputs": { + "vae_name": "wan_2.1_vae.safetensors" + } + }, + "93": { + "class_type": "CLIPTextEncode", + "inputs": { + "clip": ["84", 0], + "text": "" + } + }, + "95": { + "class_type": "UNETLoader", + "inputs": { + "unet_name": "wan2.2_i2v_high_noise_14B_fp8_scaled.safetensors", + "weight_dtype": "default" + } + }, + "96": { + "class_type": "UNETLoader", + "inputs": { + "unet_name": "wan2.2_i2v_low_noise_14B_fp8_scaled.safetensors", + "weight_dtype": "default" + } + }, + "97": { + "class_type": "LoadImage", + "inputs": { + "image": "" + } + }, + "98": { + "class_type": "WanImageToVideo", + "inputs": { + "width": 640, + "height": 640, + "length": 81, + "batch_size": 1, + "positive": ["93", 0], + "negative": ["89", 0], + "vae": ["90", 0], + "start_image": ["97", 0] + } + }, + "101": { + "class_type": "LoraLoaderModelOnly", + "inputs": { + "model": ["95", 0], + "lora_name": "wan2.2_i2v_lightx2v_4steps_lora_v1_high_noise.safetensors", + "strength_model": 1.0 + } + }, + "102": { + "class_type": "LoraLoaderModelOnly", + "inputs": { + "model": ["96", 0], + "lora_name": "wan2.2_i2v_lightx2v_4steps_lora_v1_low_noise.safetensors", + "strength_model": 1.0 + } + }, + "103": { + "class_type": "ModelSamplingSD3", + "inputs": { + "model": ["102", 0], + "shift": 5.0 + } + }, + "104": { + "class_type": "ModelSamplingSD3", + "inputs": { + "model": ["101", 0], + "shift": 5.0 + } + }, + "86": { + "class_type": "KSamplerAdvanced", + "inputs": { + "model": ["104", 0], + "positive": ["98", 0], + "negative": ["98", 1], + "latent_image": ["98", 2], + "add_noise": "enable", + "noise_seed": 42, + "control_after_generate": "randomize", + "steps": 4, + "cfg": 1.0, + "sampler_name": "euler", + "scheduler": "simple", + "start_at_step": 0, + "end_at_step": 2, + "return_with_leftover_noise": "enable" + } + }, + "85": { + "class_type": "KSamplerAdvanced", + "inputs": { + "model": ["103", 0], + "positive": ["98", 0], + "negative": ["98", 1], + "latent_image": ["86", 0], + "add_noise": "disable", + "noise_seed": 0, + "control_after_generate": "fixed", + "steps": 4, + "cfg": 1.0, + "sampler_name": "euler", + "scheduler": "simple", + "start_at_step": 2, + "end_at_step": 4, + "return_with_leftover_noise": "disable" + } + }, + "87": { + "class_type": "VAEDecode", + "inputs": { + "samples": ["85", 0], + "vae": ["90", 0] + } + }, + "94": { + "class_type": "CreateVideo", + "inputs": { + "images": ["87", 0], + "fps": 16 + } + }, + "108": { + "class_type": "SaveVideo", + "inputs": { + "video": ["94", 0], + "filename_prefix": "openmontage_i2v", + "format": "auto", + "codec": "auto" + } + } +} diff --git a/tools/_comfyui/workflows/wan22-t2v-4step.json b/tools/_comfyui/workflows/wan22-t2v-4step.json new file mode 100644 index 00000000..f5772aca --- /dev/null +++ b/tools/_comfyui/workflows/wan22-t2v-4step.json @@ -0,0 +1,143 @@ +{ + "1": { + "class_type": "CLIPLoader", + "inputs": { + "clip_name": "umt5_xxl_fp8_e4m3fn_scaled.safetensors", + "type": "wan", + "device": "default" + } + }, + "2": { + "class_type": "CLIPTextEncode", + "inputs": { + "clip": ["1", 0], + "text": "" + } + }, + "3": { + "class_type": "CLIPTextEncode", + "inputs": { + "clip": ["1", 0], + "text": "oversaturated, overexposed, static, blurry details, subtitles, style, artwork, painting, still frame, gray overall, worst quality, low quality, JPEG artifacts, ugly, deformed, extra fingers, poorly drawn hands, poorly drawn face, deformed limbs, fused fingers, static frame, cluttered background, three legs, many people in background, walking backwards" + } + }, + "4": { + "class_type": "VAELoader", + "inputs": { + "vae_name": "wan2.2_vae.safetensors" + } + }, + "5": { + "class_type": "UNETLoader", + "inputs": { + "unet_name": "wan2.2_t2v_high_noise_14B_fp8_scaled.safetensors", + "weight_dtype": "default" + } + }, + "6": { + "class_type": "UNETLoader", + "inputs": { + "unet_name": "wan2.2_t2v_low_noise_14B_fp8_scaled.safetensors", + "weight_dtype": "default" + } + }, + "7": { + "class_type": "LoraLoaderModelOnly", + "inputs": { + "model": ["5", 0], + "lora_name": "wan2.2_t2v_lightx2v_4steps_lora_v1.1_high_noise.safetensors", + "strength_model": 1.0 + } + }, + "8": { + "class_type": "LoraLoaderModelOnly", + "inputs": { + "model": ["6", 0], + "lora_name": "wan2.2_t2v_lightx2v_4steps_lora_v1.1_low_noise.safetensors", + "strength_model": 1.0 + } + }, + "9": { + "class_type": "ModelSamplingSD3", + "inputs": { + "model": ["7", 0], + "shift": 5.0 + } + }, + "10": { + "class_type": "ModelSamplingSD3", + "inputs": { + "model": ["8", 0], + "shift": 5.0 + } + }, + "11": { + "class_type": "EmptyLatentImage", + "inputs": { + "width": 832, + "height": 480, + "batch_size": 81 + } + }, + "12": { + "class_type": "KSamplerAdvanced", + "inputs": { + "model": ["9", 0], + "positive": ["2", 0], + "negative": ["3", 0], + "latent_image": ["11", 0], + "add_noise": "enable", + "noise_seed": 42, + "control_after_generate": "randomize", + "steps": 4, + "cfg": 1.0, + "sampler_name": "euler", + "scheduler": "simple", + "start_at_step": 0, + "end_at_step": 2, + "return_with_leftover_noise": "enable" + } + }, + "13": { + "class_type": "KSamplerAdvanced", + "inputs": { + "model": ["10", 0], + "positive": ["2", 0], + "negative": ["3", 0], + "latent_image": ["12", 0], + "add_noise": "disable", + "noise_seed": 0, + "control_after_generate": "fixed", + "steps": 4, + "cfg": 1.0, + "sampler_name": "euler", + "scheduler": "simple", + "start_at_step": 2, + "end_at_step": 4, + "return_with_leftover_noise": "disable" + } + }, + "14": { + "class_type": "VAEDecode", + "inputs": { + "samples": ["13", 0], + "vae": ["4", 0] + } + }, + "15": { + "class_type": "CreateVideo", + "inputs": { + "images": ["14", 0], + "fps": 16 + } + }, + "16": { + "class_type": "SaveVideo", + "inputs": { + "video": ["15", 0], + "filename_prefix": "openmontage_t2v", + "format": "auto", + "codec": "auto" + } + } +} diff --git a/tools/audio/comfyui_music.py b/tools/audio/comfyui_music.py new file mode 100644 index 00000000..3d95672b --- /dev/null +++ b/tools/audio/comfyui_music.py @@ -0,0 +1,192 @@ +"""ComfyUI music generation via ACE-Step model. + +Generates background music and songs locally using the ACE-Step 3.5B +model running inside a ComfyUI server. Custom workflows are accepted +via the ``workflow_json`` input. +""" + +from __future__ import annotations + +import json +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 + +_WORKFLOWS = Path(__file__).resolve().parent.parent / "_comfyui" / "workflows" + +_OUTPUT_NODE = "3" + + +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 = [] + install_instructions = ( + "Start a ComfyUI server and set COMFYUI_SERVER_URL " + "(default http://localhost:8188).\n" + "Requires ACE-Step model (ace_step_v1_3.5b.safetensors) in " + "ComfyUI's checkpoints directory and the ACE-Step custom node installed." + ) + agent_skills = ["music"] + + capabilities = [ + "generate_background_music", + "generate_instrumental", + "generate_song", + "text_to_music", + ] + supports = { + "seed": True, + "duration_control": True, + "lyrics": True, + "custom_workflow": True, + "offline": True, + } + best_for = [ + "local music generation without API costs", + "background music and instrumentals for video production", + "song generation with lyrics", + ] + not_good_for = [ + "setups without a running ComfyUI server", + "highest quality commercial music (use Suno or ElevenLabs)", + ] + fallback = "suno_music" + fallback_tools = ["suno_music", "elevenlabs_music", "freesound_music"] + + input_schema = { + "type": "object", + "required": ["prompt"], + "properties": { + "prompt": { + "type": "string", + "description": "Music style / mood description (e.g. 'upbeat corporate background music')", + }, + "lyrics": { + "type": "string", + "default": "", + "description": "Optional lyrics for song generation", + }, + "duration": { + "type": "number", + "default": 30.0, + "description": "Duration in seconds", + }, + "steps": {"type": "integer", "default": 60}, + "cfg": {"type": "number", "default": 3.0}, + "seed": {"type": "integer", "description": "Random if omitted"}, + "output_path": {"type": "string", "description": "Where to save the audio"}, + "workflow_json": { + "type": "string", + "description": "Optional full ComfyUI workflow JSON (overrides default)", + }, + }, + } + + resource_profile = ResourceProfile( + cpu_cores=2, ram_mb=8000, vram_mb=6000, disk_mb=500, network_required=False, + ) + retry_policy = RetryPolicy(max_retries=1, retryable_errors=["timeout"]) + idempotency_key_fields = ["prompt", "lyrics", "duration", "steps", "seed"] + side_effects = ["writes audio file to output_path"] + user_visible_verification = ["Listen to generated audio for quality and mood match"] + + def __init__(self) -> None: + self._client = ComfyUIClient() + + def get_status(self) -> ToolStatus: + if self._client.is_available(): + return ToolStatus.AVAILABLE + return ToolStatus.UNAVAILABLE + + def estimate_cost(self, inputs: dict[str, Any]) -> float: + return 0.0 + + def estimate_runtime(self, inputs: dict[str, Any]) -> float: + duration = inputs.get("duration", 30.0) + return duration * 2.0 # rough: ~2x realtime + + def execute(self, inputs: dict[str, Any]) -> ToolResult: + if not self._client.is_available(): + return ToolResult( + success=False, + error="ComfyUI server not reachable. " + self.install_instructions, + ) + + start = time.time() + seed = inputs.get("seed") or ComfyUIClient.random_seed() + duration = inputs.get("duration", 30.0) + output_path = Path( + inputs.get("output_path", f"comfyui_music_{seed}.wav") + ) + + try: + if inputs.get("workflow_json"): + workflow = json.loads(inputs["workflow_json"]) + else: + workflow = ComfyUIClient.load_workflow( + _WORKFLOWS / "ace-step-music.json" + ) + workflow = ComfyUIClient.patch_workflow(workflow, { + "2": { + "prompt": inputs["prompt"], + "lyrics": inputs.get("lyrics", ""), + "duration": duration, + "seed": seed, + "steps": inputs.get("steps", 60), + "cfg": inputs.get("cfg", 3.0), + }, + "3": {"filename_prefix": output_path.stem}, + }) + + paths = self._client.generate( + workflow, + output_node=_OUTPUT_NODE, + dest=output_path, + timeout=int(duration * 4), # generous timeout + ) + + except ComfyUIError as exc: + return ToolResult(success=False, error=str(exc)) + except Exception as exc: + return ToolResult(success=False, error=f"ComfyUI music generation failed: {exc}") + + return ToolResult( + success=True, + data={ + "provider": "comfyui", + "model": "ace-step-v1-3.5b", + "prompt": inputs["prompt"], + "lyrics": inputs.get("lyrics", ""), + "duration": duration, + "output": str(paths[0]), + "format": output_path.suffix.lstrip("."), + }, + artifacts=[str(p) for p in paths], + cost_usd=0.0, + duration_seconds=round(time.time() - start, 2), + seed=seed, + model="ace-step-v1-3.5b", + ) diff --git a/tools/graphics/comfyui_image.py b/tools/graphics/comfyui_image.py new file mode 100644 index 00000000..002723a5 --- /dev/null +++ b/tools/graphics/comfyui_image.py @@ -0,0 +1,165 @@ +"""ComfyUI image generation via a local or remote ComfyUI server. + +Default workflow: FLUX 2 Dev (NVFP4) with Mistral text encoder. +Supports custom workflows via the ``workflow_json`` input. +""" + +from __future__ import annotations + +import json +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 + +_WORKFLOWS = Path(__file__).resolve().parent.parent / "_comfyui" / "workflows" + + +class ComfyUIImage(BaseTool): + name = "comfyui_image" + version = "0.1.0" + tier = ToolTier.GENERATE + capability = "image_generation" + provider = "comfyui" + stability = ToolStability.EXPERIMENTAL + execution_mode = ExecutionMode.SYNC + determinism = Determinism.SEEDED + runtime = ToolRuntime.LOCAL_GPU + + dependencies = [] # checked at runtime via server health + install_instructions = ( + "Start a ComfyUI server and set COMFYUI_SERVER_URL " + "(default http://localhost:8188).\n" + "See https://github.com/comfyanonymous/ComfyUI for setup." + ) + agent_skills = [] + + capabilities = ["text_to_image"] + supports = { + "seed": True, + "custom_size": True, + "custom_workflow": True, + "offline": True, + } + best_for = [ + "local GPU generation without API costs", + "Blackwell / DGX Spark hardware where diffusers is unsupported", + "full control over sampling via custom ComfyUI workflows", + ] + not_good_for = [ + "setups without a running ComfyUI server", + "CPU-only machines", + ] + fallback = "flux_image" + fallback_tools = ["flux_image", "local_diffusion", "openai_image"] + + input_schema = { + "type": "object", + "required": ["prompt"], + "properties": { + "prompt": {"type": "string", "description": "Text prompt for image generation"}, + "width": {"type": "integer", "default": 1024}, + "height": {"type": "integer", "default": 1024}, + "steps": {"type": "integer", "default": 20}, + "guidance": {"type": "number", "default": 3.5}, + "seed": {"type": "integer", "description": "Random if omitted"}, + "output_path": {"type": "string", "description": "Where to save the image"}, + "workflow_json": { + "type": "string", + "description": "Optional full ComfyUI workflow JSON (overrides default)", + }, + }, + } + + 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", "width", "height", "steps", "seed"] + side_effects = ["writes image file to output_path"] + user_visible_verification = ["Inspect generated image for quality and prompt adherence"] + + def __init__(self) -> None: + self._client = ComfyUIClient() + + def get_status(self) -> ToolStatus: + if self._client.is_available(): + return ToolStatus.AVAILABLE + return ToolStatus.UNAVAILABLE + + def estimate_cost(self, inputs: dict[str, Any]) -> float: + return 0.0 + + def estimate_runtime(self, inputs: dict[str, Any]) -> float: + return float(inputs.get("steps", 20)) * 1.5 + + def execute(self, inputs: dict[str, Any]) -> ToolResult: + if not self._client.is_available(): + return ToolResult( + success=False, + error="ComfyUI server not reachable. " + self.install_instructions, + ) + + start = time.time() + seed = inputs.get("seed") or ComfyUIClient.random_seed() + width = inputs.get("width", 1024) + height = inputs.get("height", 1024) + steps = inputs.get("steps", 20) + guidance = inputs.get("guidance", 3.5) + output_path = Path(inputs.get("output_path", f"comfyui_image_{seed}.png")) + + try: + if inputs.get("workflow_json"): + workflow = json.loads(inputs["workflow_json"]) + else: + workflow = ComfyUIClient.load_workflow(_WORKFLOWS / "flux2-txt2img.json") + workflow = ComfyUIClient.patch_workflow(workflow, { + "4": {"text": inputs["prompt"]}, + "5": {"guidance": guidance}, + "6": {"width": width, "height": height, "batch_size": 1}, + "7": {"noise_seed": seed}, + "10": {"steps": steps, "width": width, "height": height}, + "13": {"filename_prefix": output_path.stem}, + }) + + paths = self._client.generate( + workflow, output_node="13", dest=output_path, timeout=600, + ) + + except ComfyUIError as exc: + return ToolResult(success=False, error=str(exc)) + except Exception as exc: + return ToolResult(success=False, error=f"ComfyUI image generation failed: {exc}") + + return ToolResult( + success=True, + data={ + "provider": "comfyui", + "model": "flux2-dev-nvfp4", + "prompt": inputs["prompt"], + "width": width, + "height": height, + "steps": steps, + "guidance": guidance, + "output": str(paths[0]), + "format": "png", + }, + artifacts=[str(p) for p in paths], + cost_usd=0.0, + duration_seconds=round(time.time() - start, 2), + seed=seed, + model="flux2-dev-nvfp4", + ) diff --git a/tools/video/comfyui_video.py b/tools/video/comfyui_video.py new file mode 100644 index 00000000..c17f7f63 --- /dev/null +++ b/tools/video/comfyui_video.py @@ -0,0 +1,250 @@ +"""ComfyUI video generation via a local or remote ComfyUI server. + +Supports text-to-video and image-to-video using WAN 2.2 14B with +4-step LightX2V LoRA acceleration. Custom workflows are accepted +via the ``workflow_json`` input. +""" + +from __future__ import annotations + +import json +import time +from pathlib import Path +from typing import Any + +import requests + +from tools.base_tool import ( + BaseTool, + Determinism, + ExecutionMode, + ResourceProfile, + RetryPolicy, + ToolResult, + ToolRuntime, + ToolStability, + ToolStatus, + ToolTier, +) +from tools._comfyui.client import ComfyUIClient, ComfyUIError + +_WORKFLOWS = Path(__file__).resolve().parent.parent / "_comfyui" / "workflows" + +# Output node IDs in the bundled workflows +_T2V_OUTPUT_NODE = "16" +_I2V_OUTPUT_NODE = "108" + + +class ComfyUIVideo(BaseTool): + name = "comfyui_video" + version = "0.1.0" + tier = ToolTier.GENERATE + capability = "video_generation" + provider = "comfyui" + stability = ToolStability.EXPERIMENTAL + execution_mode = ExecutionMode.SYNC + determinism = Determinism.SEEDED + runtime = ToolRuntime.LOCAL_GPU + + dependencies = [] + 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." + ) + agent_skills = [] + + capabilities = ["text_to_video", "image_to_video"] + supports = { + "seed": True, + "reference_image": True, + "custom_workflow": True, + "offline": True, + } + best_for = [ + "local GPU video generation without API costs", + "Blackwell / DGX Spark hardware where diffusers is unsupported", + "image-to-video with WAN 2.2 14B (4-step accelerated)", + "text-to-video with WAN 2.2 14B (4-step accelerated)", + ] + not_good_for = [ + "setups without a running ComfyUI server", + "CPU-only machines", + ] + fallback = "wan_video" + fallback_tools = ["wan_video", "hunyuan_video", "ltx_video_local", "kling_video"] + + input_schema = { + "type": "object", + "required": ["prompt"], + "properties": { + "prompt": {"type": "string", "description": "Text prompt for video generation"}, + "operation": { + "type": "string", + "enum": ["text_to_video", "image_to_video"], + "default": "text_to_video", + }, + "reference_image_path": { + "type": "string", + "description": "Local path to reference image (for image_to_video)", + }, + "reference_image_url": { + "type": "string", + "description": "URL of reference image (for image_to_video, downloaded first)", + }, + "width": {"type": "integer", "default": 832, "description": "T2V default 832, I2V default 640"}, + "height": {"type": "integer", "default": 480, "description": "T2V default 480, I2V default 640"}, + "num_frames": {"type": "integer", "default": 81, "description": "81 frames = 5s at 16fps"}, + "seed": {"type": "integer", "description": "Random if omitted"}, + "output_path": {"type": "string", "description": "Where to save the video"}, + "workflow_json": { + "type": "string", + "description": "Optional full ComfyUI workflow JSON (overrides default)", + }, + }, + } + + resource_profile = ResourceProfile( + cpu_cores=2, ram_mb=32000, vram_mb=16000, disk_mb=2000, network_required=False, + ) + retry_policy = RetryPolicy(max_retries=1, retryable_errors=["timeout"]) + idempotency_key_fields = ["prompt", "operation", "width", "height", "num_frames", "seed"] + side_effects = ["writes video file to output_path"] + user_visible_verification = ["Watch generated clip for motion coherence and artifacts"] + + def __init__(self) -> None: + self._client = ComfyUIClient() + + def get_status(self) -> ToolStatus: + if self._client.is_available(): + return ToolStatus.AVAILABLE + return ToolStatus.UNAVAILABLE + + def estimate_cost(self, inputs: dict[str, Any]) -> float: + return 0.0 + + def estimate_runtime(self, inputs: dict[str, Any]) -> float: + operation = inputs.get("operation", "text_to_video") + if operation == "image_to_video": + return 210.0 # ~3.5 min + return 240.0 # ~4 min + + def execute(self, inputs: dict[str, Any]) -> ToolResult: + if not self._client.is_available(): + return ToolResult( + success=False, + error="ComfyUI server not reachable. " + self.install_instructions, + ) + + operation = inputs.get("operation", "text_to_video") + start = time.time() + seed = inputs.get("seed") or ComfyUIClient.random_seed() + output_path = Path( + inputs.get("output_path", f"comfyui_video_{operation}_{seed}.mp4") + ) + + try: + if inputs.get("workflow_json"): + workflow = json.loads(inputs["workflow_json"]) + output_node = _T2V_OUTPUT_NODE + elif operation == "image_to_video": + workflow, output_node = self._build_i2v(inputs, seed, output_path) + else: + workflow, output_node = self._build_t2v(inputs, seed, output_path) + + paths = self._client.generate( + workflow, + output_node=output_node, + dest=output_path, + timeout=900, + interval=10, + ) + + except ComfyUIError as exc: + return ToolResult(success=False, error=str(exc)) + except Exception as exc: + return ToolResult(success=False, error=f"ComfyUI video generation failed: {exc}") + + width = inputs.get("width", 832 if operation == "text_to_video" else 640) + height = inputs.get("height", 480 if operation == "text_to_video" else 640) + num_frames = inputs.get("num_frames", 81) + + return ToolResult( + success=True, + data={ + "provider": "comfyui", + "model": "wan2.2-14b-fp8-4step", + "prompt": inputs["prompt"], + "operation": operation, + "width": width, + "height": height, + "num_frames": num_frames, + "fps": 16, + "duration_seconds": round(num_frames / 16, 2), + "output": str(paths[0]), + "format": "mp4", + }, + artifacts=[str(p) for p in paths], + cost_usd=0.0, + duration_seconds=round(time.time() - start, 2), + seed=seed, + model="wan2.2-14b-fp8-4step", + ) + + # ------------------------------------------------------------------ + # Workflow builders + # ------------------------------------------------------------------ + + def _build_t2v( + self, inputs: dict[str, Any], seed: int, output_path: Path + ) -> tuple[dict, str]: + width = inputs.get("width", 832) + height = inputs.get("height", 480) + num_frames = inputs.get("num_frames", 81) + + workflow = ComfyUIClient.load_workflow(_WORKFLOWS / "wan22-t2v-4step.json") + workflow = ComfyUIClient.patch_workflow(workflow, { + "2": {"text": inputs["prompt"]}, + "11": {"width": width, "height": height, "batch_size": num_frames}, + "12": {"noise_seed": seed}, + "16": {"filename_prefix": output_path.stem}, + }) + return workflow, _T2V_OUTPUT_NODE + + def _build_i2v( + self, inputs: dict[str, Any], seed: int, output_path: Path + ) -> tuple[dict, str]: + width = inputs.get("width", 640) + height = inputs.get("height", 640) + num_frames = inputs.get("num_frames", 81) + + # Resolve reference image + ref_path = inputs.get("reference_image_path") + ref_url = inputs.get("reference_image_url") + + if ref_url and not ref_path: + # Download to a temp location + resp = requests.get(ref_url, timeout=60) + resp.raise_for_status() + ref_path = str(output_path.with_suffix(".ref.png")) + Path(ref_path).parent.mkdir(parents=True, exist_ok=True) + Path(ref_path).write_bytes(resp.content) + + if not ref_path: + raise ComfyUIError( + "image_to_video requires reference_image_path or reference_image_url" + ) + + # Upload to ComfyUI + upload_name = f"om_{output_path.stem}.png" + server_name = self._client.upload_image(Path(ref_path), upload_name) + + workflow = ComfyUIClient.load_workflow(_WORKFLOWS / "wan22-i2v-4step.json") + workflow = ComfyUIClient.patch_workflow(workflow, { + "93": {"text": inputs["prompt"]}, + "97": {"image": server_name}, + "98": {"width": width, "height": height, "length": num_frames}, + "86": {"noise_seed": seed}, + "108": {"filename_prefix": output_path.stem}, + }) + return workflow, _I2V_OUTPUT_NODE