comfyui: add native ComfyUI provider for image, video, and music generation

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) <noreply@anthropic.com>
This commit is contained in:
martimramos
2026-04-16 23:59:33 +01:00
committed by Alastair Beal
parent 80e51fd618
commit 6ec2bbb090
11 changed files with 1793 additions and 0 deletions

View File

@@ -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?

View File

@@ -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

View File

@@ -0,0 +1 @@
"""ComfyUI integration — shared client and bundled workflow templates."""

207
tools/_comfyui/client.py Normal file
View File

@@ -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)

View File

@@ -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"
}
}
}

View File

@@ -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"
}
}
}

View File

@@ -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"
}
}
}

View File

@@ -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"
}
}
}

View File

@@ -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",
)

View File

@@ -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",
)

View File

@@ -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