freedak/OpenMontage
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files
tools/_comfyui/__init__.py
+1
View File
@@ -0,0 +1 @@
"""ComfyUI integration — shared client and bundled workflow templates."""
tools/_comfyui/client.py
+207
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)
tools/_comfyui/workflows/ace-step-music.json
+27
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"
}
}
}
tools/_comfyui/workflows/flux2-txt2img.json
+96
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"
}
}
}
tools/_comfyui/workflows/wan22-i2v-4step.json
+154
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"
}
}
}
tools/_comfyui/workflows/wan22-t2v-4step.json
+143
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"
}
}
}