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
committed by
Alastair Beal
Alastair Beal
parent
80e51fd618
commit
6ec2bbb090
@@ -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
|
||||
Reference in New Issue
Block a user