comfyui: satisfy provider contract review items
This commit is contained in:
calesthio
@@ -140,10 +140,15 @@ class ComfyUIClient:
|
||||
json={"prompt": workflow},
|
||||
timeout=30,
|
||||
)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
try:
|
||||
data = resp.json()
|
||||
except ValueError:
|
||||
data = {}
|
||||
if data.get("node_errors"):
|
||||
raise ComfyUIError(f"Node errors: {json.dumps(data['node_errors'])}")
|
||||
if data.get("error"):
|
||||
raise ComfyUIError(f"Prompt error: {json.dumps(data['error'])}")
|
||||
resp.raise_for_status()
|
||||
prompt_id = data.get("prompt_id")
|
||||
if not prompt_id:
|
||||
raise ComfyUIError(f"No prompt_id in response: {data}")
|
||||
@@ -181,6 +186,7 @@ class ComfyUIClient:
|
||||
filename: str,
|
||||
subfolder: str,
|
||||
dest: Path,
|
||||
folder_type: str = "output",
|
||||
) -> Path:
|
||||
"""Download an output artifact from the ComfyUI server."""
|
||||
resp = requests.get(
|
||||
@@ -188,7 +194,7 @@ class ComfyUIClient:
|
||||
params={
|
||||
"filename": filename,
|
||||
"subfolder": subfolder,
|
||||
"type": "output",
|
||||
"type": folder_type,
|
||||
},
|
||||
timeout=120,
|
||||
)
|
||||
@@ -246,7 +252,12 @@ class ComfyUIClient:
|
||||
target = dest
|
||||
else:
|
||||
target = dest.with_stem(f"{dest.stem}_{i:03d}").with_suffix(suffix)
|
||||
self.download(item["filename"], item.get("subfolder", ""), target)
|
||||
self.download(
|
||||
item["filename"],
|
||||
item.get("subfolder", ""),
|
||||
target,
|
||||
item.get("type", "output"),
|
||||
)
|
||||
paths.append(target)
|
||||
return paths
|
||||
|
||||
|
||||
@@ -0,0 +1,222 @@
|
||||
"""Shared metadata helpers for ComfyUI provider tools."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
|
||||
COMFYUI_SETUP_OFFER: dict[str, Any] = {
|
||||
"kind": "local_server",
|
||||
"fix_complexity": "1-minute env-var if ComfyUI is already running; otherwise local install",
|
||||
"env_var": "COMFYUI_SERVER_URL",
|
||||
"default_url": "http://localhost:8188",
|
||||
"health_check": "GET /system_stats",
|
||||
"what_it_unlocks": [
|
||||
"free local image generation through ComfyUI workflows",
|
||||
"free local video generation through ComfyUI workflows",
|
||||
"community workflow_json/workflow_path execution",
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
BUNDLED_MODEL_STACKS: dict[str, list[dict[str, Any]]] = {
|
||||
"flux2-txt2img": [
|
||||
{
|
||||
"role": "diffusion_model",
|
||||
"name": "flux2-dev-nvfp4.safetensors",
|
||||
"quantization": "NVFP4",
|
||||
"destination_hint": "ComfyUI/models/diffusion_models/",
|
||||
"download_url": (
|
||||
"https://huggingface.co/black-forest-labs/FLUX.2-dev-NVFP4"
|
||||
),
|
||||
},
|
||||
{
|
||||
"role": "text_encoder",
|
||||
"name": "mistral_3_small_flux2_fp4_mixed.safetensors",
|
||||
"quantization": "FP4 mixed",
|
||||
"destination_hint": "ComfyUI/models/text_encoders/",
|
||||
"download_url": (
|
||||
"https://huggingface.co/Comfy-Org/flux2-dev/tree/main/"
|
||||
"split_files/text_encoders"
|
||||
),
|
||||
},
|
||||
{
|
||||
"role": "vae",
|
||||
"name": "flux2-vae.safetensors",
|
||||
"destination_hint": "ComfyUI/models/vae/",
|
||||
"download_url": (
|
||||
"https://huggingface.co/Comfy-Org/flux2-dev/blob/main/"
|
||||
"split_files/vae/flux2-vae.safetensors"
|
||||
),
|
||||
},
|
||||
],
|
||||
"wan22-t2v-4step": [
|
||||
{
|
||||
"role": "text_encoder",
|
||||
"name": "umt5_xxl_fp8_e4m3fn_scaled.safetensors",
|
||||
"quantization": "FP8",
|
||||
"destination_hint": "ComfyUI/models/text_encoders/",
|
||||
"download_url": (
|
||||
"https://huggingface.co/Comfy-Org/Wan_2.2_ComfyUI_Repackaged/"
|
||||
"tree/main/split_files/text_encoders"
|
||||
),
|
||||
},
|
||||
{
|
||||
"role": "diffusion_model_high_noise",
|
||||
"name": "wan2.2_t2v_high_noise_14B_fp8_scaled.safetensors",
|
||||
"quantization": "FP8",
|
||||
"destination_hint": "ComfyUI/models/diffusion_models/",
|
||||
"download_url": (
|
||||
"https://huggingface.co/Comfy-Org/Wan_2.2_ComfyUI_Repackaged/"
|
||||
"blob/main/split_files/diffusion_models/"
|
||||
"wan2.2_t2v_high_noise_14B_fp8_scaled.safetensors"
|
||||
),
|
||||
},
|
||||
{
|
||||
"role": "diffusion_model_low_noise",
|
||||
"name": "wan2.2_t2v_low_noise_14B_fp8_scaled.safetensors",
|
||||
"quantization": "FP8",
|
||||
"destination_hint": "ComfyUI/models/diffusion_models/",
|
||||
"download_url": (
|
||||
"https://huggingface.co/Comfy-Org/Wan_2.2_ComfyUI_Repackaged/"
|
||||
"tree/main/split_files/diffusion_models"
|
||||
),
|
||||
},
|
||||
{
|
||||
"role": "vae",
|
||||
"name": "wan2.2_vae.safetensors",
|
||||
"destination_hint": "ComfyUI/models/vae/",
|
||||
"download_url": (
|
||||
"https://huggingface.co/Comfy-Org/Wan_2.2_ComfyUI_Repackaged/"
|
||||
"tree/main/split_files/vae"
|
||||
),
|
||||
},
|
||||
{
|
||||
"role": "lora",
|
||||
"name": "wan2.2_t2v_lightx2v_4steps_lora_v1.1_high_noise.safetensors",
|
||||
"strength_model": 1.0,
|
||||
"destination_hint": "ComfyUI/models/loras/",
|
||||
"download_url": (
|
||||
"https://huggingface.co/Comfy-Org/Wan_2.2_ComfyUI_Repackaged/"
|
||||
"tree/main/split_files/loras"
|
||||
),
|
||||
},
|
||||
{
|
||||
"role": "lora",
|
||||
"name": "wan2.2_t2v_lightx2v_4steps_lora_v1.1_low_noise.safetensors",
|
||||
"strength_model": 1.0,
|
||||
"destination_hint": "ComfyUI/models/loras/",
|
||||
"download_url": (
|
||||
"https://huggingface.co/Comfy-Org/Wan_2.2_ComfyUI_Repackaged/"
|
||||
"tree/main/split_files/loras"
|
||||
),
|
||||
},
|
||||
],
|
||||
"wan22-i2v-4step": [
|
||||
{
|
||||
"role": "text_encoder",
|
||||
"name": "umt5_xxl_fp8_e4m3fn_scaled.safetensors",
|
||||
"quantization": "FP8",
|
||||
"destination_hint": "ComfyUI/models/text_encoders/",
|
||||
"download_url": (
|
||||
"https://huggingface.co/Comfy-Org/Wan_2.2_ComfyUI_Repackaged/"
|
||||
"tree/main/split_files/text_encoders"
|
||||
),
|
||||
},
|
||||
{
|
||||
"role": "diffusion_model_high_noise",
|
||||
"name": "wan2.2_i2v_high_noise_14B_fp8_scaled.safetensors",
|
||||
"quantization": "FP8",
|
||||
"destination_hint": "ComfyUI/models/diffusion_models/",
|
||||
"download_url": (
|
||||
"https://huggingface.co/Comfy-Org/Wan_2.2_ComfyUI_Repackaged/"
|
||||
"blob/main/split_files/diffusion_models/"
|
||||
"wan2.2_i2v_high_noise_14B_fp8_scaled.safetensors"
|
||||
),
|
||||
},
|
||||
{
|
||||
"role": "diffusion_model_low_noise",
|
||||
"name": "wan2.2_i2v_low_noise_14B_fp8_scaled.safetensors",
|
||||
"quantization": "FP8",
|
||||
"destination_hint": "ComfyUI/models/diffusion_models/",
|
||||
"download_url": (
|
||||
"https://huggingface.co/Comfy-Org/Wan_2.2_ComfyUI_Repackaged/"
|
||||
"tree/main/split_files/diffusion_models"
|
||||
),
|
||||
},
|
||||
{
|
||||
"role": "vae",
|
||||
"name": "wan_2.1_vae.safetensors",
|
||||
"destination_hint": "ComfyUI/models/vae/",
|
||||
"download_url": (
|
||||
"https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/"
|
||||
"tree/main/split_files/vae"
|
||||
),
|
||||
},
|
||||
{
|
||||
"role": "lora",
|
||||
"name": "wan2.2_i2v_lightx2v_4steps_lora_v1_high_noise.safetensors",
|
||||
"strength_model": 1.0,
|
||||
"destination_hint": "ComfyUI/models/loras/",
|
||||
"download_url": (
|
||||
"https://huggingface.co/Comfy-Org/Wan_2.2_ComfyUI_Repackaged/"
|
||||
"tree/main/split_files/loras"
|
||||
),
|
||||
},
|
||||
{
|
||||
"role": "lora",
|
||||
"name": "wan2.2_i2v_lightx2v_4steps_lora_v1_low_noise.safetensors",
|
||||
"strength_model": 1.0,
|
||||
"destination_hint": "ComfyUI/models/loras/",
|
||||
"download_url": (
|
||||
"https://huggingface.co/Comfy-Org/Wan_2.2_ComfyUI_Repackaged/"
|
||||
"tree/main/split_files/loras"
|
||||
),
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def workflow_hash(workflow: dict[str, Any]) -> str:
|
||||
"""Return a stable hash of the final workflow JSON submitted to ComfyUI."""
|
||||
payload = json.dumps(workflow, sort_keys=True, separators=(",", ":"))
|
||||
return hashlib.sha256(payload.encode("utf-8")).hexdigest()
|
||||
|
||||
|
||||
def model_stack(workflow_key: str | None, inputs: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
"""Return bundled or caller-supplied model stack metadata."""
|
||||
if workflow_key:
|
||||
return [dict(item) for item in BUNDLED_MODEL_STACKS[workflow_key]]
|
||||
stack = inputs.get("workflow_model_stack")
|
||||
return stack if isinstance(stack, list) else []
|
||||
|
||||
|
||||
def missing_models_payload(
|
||||
missing: list[str],
|
||||
*,
|
||||
workflow_key: str,
|
||||
workflow_name: str,
|
||||
operation: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Build a machine-readable missing-model error payload."""
|
||||
stack_by_name = {
|
||||
item["name"]: item for item in BUNDLED_MODEL_STACKS.get(workflow_key, [])
|
||||
}
|
||||
items = []
|
||||
for name in missing:
|
||||
meta = dict(stack_by_name.get(name, {}))
|
||||
meta.setdefault("name", name)
|
||||
meta.setdefault("role", "unknown")
|
||||
meta.setdefault("destination_hint", "ComfyUI/models/ matching the workflow node")
|
||||
meta.setdefault("download_url", None)
|
||||
items.append(meta)
|
||||
|
||||
return {
|
||||
"provider": "comfyui",
|
||||
"workflow": workflow_name,
|
||||
"operation": operation,
|
||||
"missing_models": items,
|
||||
"setup_offer": COMFYUI_SETUP_OFFER,
|
||||
}
|
||||
@@ -24,6 +24,13 @@ from tools.base_tool import (
|
||||
ToolTier,
|
||||
)
|
||||
from tools._comfyui.client import ComfyUIClient, ComfyUIError
|
||||
from tools._comfyui.metadata import (
|
||||
BUNDLED_MODEL_STACKS,
|
||||
COMFYUI_SETUP_OFFER,
|
||||
missing_models_payload,
|
||||
model_stack,
|
||||
workflow_hash,
|
||||
)
|
||||
|
||||
_WORKFLOWS = Path(__file__).resolve().parent.parent / "_comfyui" / "workflows"
|
||||
|
||||
@@ -47,18 +54,20 @@ class ComfyUIImage(BaseTool):
|
||||
runtime = ToolRuntime.LOCAL_GPU
|
||||
|
||||
dependencies = [] # checked at runtime via server health
|
||||
setup_offer = COMFYUI_SETUP_OFFER
|
||||
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 = []
|
||||
agent_skills = ["comfyui", "flux-best-practices"]
|
||||
|
||||
capabilities = ["text_to_image"]
|
||||
supports = {
|
||||
"seed": True,
|
||||
"custom_size": True,
|
||||
"custom_workflow": True,
|
||||
"custom_output_node": True,
|
||||
"offline": True,
|
||||
}
|
||||
best_for = [
|
||||
@@ -86,7 +95,31 @@ class ComfyUIImage(BaseTool):
|
||||
"output_path": {"type": "string", "description": "Where to save the image"},
|
||||
"workflow_json": {
|
||||
"type": "string",
|
||||
"description": "Optional full ComfyUI workflow JSON (overrides default)",
|
||||
"description": "Optional full ComfyUI workflow JSON. Requires output_node.",
|
||||
},
|
||||
"workflow_path": {
|
||||
"type": "string",
|
||||
"description": "Optional path to a ComfyUI workflow JSON file. Requires output_node.",
|
||||
},
|
||||
"output_node": {
|
||||
"type": "string",
|
||||
"description": "ComfyUI output node ID for custom workflow_json/workflow_path.",
|
||||
},
|
||||
"workflow_name": {
|
||||
"type": "string",
|
||||
"description": "Optional human-readable provenance label for a custom workflow.",
|
||||
},
|
||||
"workflow_model": {
|
||||
"type": "string",
|
||||
"description": "Optional model/provenance label for a custom workflow.",
|
||||
},
|
||||
"workflow_model_stack": {
|
||||
"type": "array",
|
||||
"description": (
|
||||
"Optional provenance metadata for custom workflow dependencies. "
|
||||
"Items should include name, role, quantization, and LoRA strengths when known."
|
||||
),
|
||||
"items": {"type": "object"},
|
||||
},
|
||||
},
|
||||
}
|
||||
@@ -116,22 +149,43 @@ class ComfyUIImage(BaseTool):
|
||||
def estimate_runtime(self, inputs: dict[str, Any]) -> float:
|
||||
return float(inputs.get("steps", 20)) * 1.5
|
||||
|
||||
def get_info(self) -> dict[str, Any]:
|
||||
info = super().get_info()
|
||||
info["setup_offer"] = self.setup_offer
|
||||
info["bundled_model_stack"] = BUNDLED_MODEL_STACKS["flux2-txt2img"]
|
||||
return info
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
custom_workflow = bool(inputs.get("workflow_json") or inputs.get("workflow_path"))
|
||||
if custom_workflow and not inputs.get("output_node"):
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=(
|
||||
"Custom ComfyUI workflows require output_node so OpenMontage "
|
||||
"knows which ComfyUI node to download artifacts from."
|
||||
),
|
||||
)
|
||||
|
||||
if not self._client.is_available():
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=self._client.unavailable_reason(),
|
||||
)
|
||||
|
||||
if not inputs.get("workflow_json"):
|
||||
if not custom_workflow:
|
||||
_, missing = self._client.check_models(_REQUIRED_MODELS)
|
||||
if missing:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
data=missing_models_payload(
|
||||
missing,
|
||||
workflow_key="flux2-txt2img",
|
||||
workflow_name="flux2-txt2img.json",
|
||||
),
|
||||
error=(
|
||||
f"ComfyUI server is running but missing required models: "
|
||||
f"{', '.join(missing)}.\n"
|
||||
f"Download them to your ComfyUI models directory."
|
||||
f"See data.missing_models for destination hints and download URLs."
|
||||
),
|
||||
)
|
||||
|
||||
@@ -144,8 +198,9 @@ class ComfyUIImage(BaseTool):
|
||||
output_path = Path(inputs.get("output_path", f"comfyui_image_{seed}.png"))
|
||||
|
||||
try:
|
||||
if inputs.get("workflow_json"):
|
||||
workflow = json.loads(inputs["workflow_json"])
|
||||
if custom_workflow:
|
||||
workflow = self._load_custom_workflow(inputs)
|
||||
output_node = str(inputs["output_node"])
|
||||
else:
|
||||
workflow = ComfyUIClient.load_workflow(_WORKFLOWS / "flux2-txt2img.json")
|
||||
workflow = ComfyUIClient.patch_workflow(workflow, {
|
||||
@@ -156,9 +211,13 @@ class ComfyUIImage(BaseTool):
|
||||
"10": {"steps": steps, "width": width, "height": height},
|
||||
"13": {"filename_prefix": output_path.stem},
|
||||
})
|
||||
output_node = "13"
|
||||
|
||||
provenance = self._workflow_provenance(
|
||||
inputs, custom_workflow, output_node, workflow
|
||||
)
|
||||
paths = self._client.generate(
|
||||
workflow, output_node="13", dest=output_path, timeout=600,
|
||||
workflow, output_node=output_node, dest=output_path, timeout=600,
|
||||
)
|
||||
|
||||
except ComfyUIError as exc:
|
||||
@@ -166,11 +225,12 @@ class ComfyUIImage(BaseTool):
|
||||
except Exception as exc:
|
||||
return ToolResult(success=False, error=f"ComfyUI image generation failed: {exc}")
|
||||
|
||||
model_name = self._model_name(inputs, custom_workflow)
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"provider": "comfyui",
|
||||
"model": "flux2-dev-nvfp4",
|
||||
"model": model_name,
|
||||
"prompt": inputs["prompt"],
|
||||
"width": width,
|
||||
"height": height,
|
||||
@@ -178,10 +238,58 @@ class ComfyUIImage(BaseTool):
|
||||
"guidance": guidance,
|
||||
"output": str(paths[0]),
|
||||
"format": "png",
|
||||
"workflow_provenance": provenance,
|
||||
},
|
||||
artifacts=[str(p) for p in paths],
|
||||
cost_usd=0.0,
|
||||
duration_seconds=round(time.time() - start, 2),
|
||||
seed=seed,
|
||||
model="flux2-dev-nvfp4",
|
||||
model=model_name,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _load_custom_workflow(inputs: dict[str, Any]) -> dict:
|
||||
if inputs.get("workflow_json"):
|
||||
return json.loads(inputs["workflow_json"])
|
||||
return ComfyUIClient.load_workflow(Path(inputs["workflow_path"]))
|
||||
|
||||
@staticmethod
|
||||
def _model_name(inputs: dict[str, Any], custom_workflow: bool) -> str:
|
||||
if not custom_workflow:
|
||||
return "flux2-dev-nvfp4"
|
||||
return (
|
||||
inputs.get("workflow_model")
|
||||
or inputs.get("model")
|
||||
or inputs.get("workflow_name")
|
||||
or "custom-comfyui-workflow"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _workflow_provenance(
|
||||
inputs: dict[str, Any],
|
||||
custom_workflow: bool,
|
||||
output_node: str,
|
||||
workflow: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
if not custom_workflow:
|
||||
return {
|
||||
"source": "bundled",
|
||||
"workflow": "flux2-txt2img.json",
|
||||
"workflow_hash_sha256": workflow_hash(workflow),
|
||||
"model_stack": model_stack("flux2-txt2img", inputs),
|
||||
"output_node": output_node,
|
||||
}
|
||||
return {
|
||||
"source": "user_supplied",
|
||||
"workflow_name": inputs.get("workflow_name"),
|
||||
"workflow_path": inputs.get("workflow_path"),
|
||||
"model": inputs.get("workflow_model") or inputs.get("model"),
|
||||
"workflow_hash_sha256": workflow_hash(workflow),
|
||||
"model_stack": model_stack(None, inputs),
|
||||
"model_stack_source": (
|
||||
"caller_supplied"
|
||||
if inputs.get("workflow_model_stack")
|
||||
else "unknown_custom_workflow"
|
||||
),
|
||||
"output_node": output_node,
|
||||
}
|
||||
|
||||
@@ -282,6 +282,7 @@ class ToolRegistry:
|
||||
"provider": tool.provider,
|
||||
"runtime": tool.runtime.value,
|
||||
"best_for": tool.best_for,
|
||||
"dependencies": info.get("dependencies", []),
|
||||
"install_instructions": tool.install_instructions,
|
||||
"status": status.value,
|
||||
}
|
||||
@@ -291,6 +292,10 @@ class ToolRegistry:
|
||||
"render_engines",
|
||||
"remotion_note",
|
||||
"provider_matrix",
|
||||
"setup_offer",
|
||||
"operation_statuses",
|
||||
"resource_profiles",
|
||||
"resource_profile_note",
|
||||
):
|
||||
if extra_key in info:
|
||||
entry[extra_key] = info[extra_key]
|
||||
@@ -398,6 +403,40 @@ class ToolRegistry:
|
||||
setup_offers: list[dict[str, Any]] = []
|
||||
for cap, bucket in menu.items():
|
||||
for entry in bucket.get("unavailable", []):
|
||||
offer = entry.get("setup_offer")
|
||||
if offer:
|
||||
setup_offers.append(
|
||||
{
|
||||
"capability": cap,
|
||||
"tool": entry.get("name"),
|
||||
"provider": entry.get("provider"),
|
||||
"runtime": entry.get("runtime"),
|
||||
"install_instructions": entry.get("install_instructions") or "",
|
||||
**offer,
|
||||
}
|
||||
)
|
||||
continue
|
||||
|
||||
env_vars = [
|
||||
dep[4:]
|
||||
for dep in entry.get("dependencies", [])
|
||||
if isinstance(dep, str) and dep.startswith("env:")
|
||||
]
|
||||
if env_vars:
|
||||
setup_offers.append(
|
||||
{
|
||||
"capability": cap,
|
||||
"tool": entry.get("name"),
|
||||
"provider": entry.get("provider"),
|
||||
"runtime": entry.get("runtime"),
|
||||
"kind": "env_var",
|
||||
"fix_complexity": "1-minute env-var",
|
||||
"env_vars": env_vars,
|
||||
"install_instructions": entry.get("install_instructions") or "",
|
||||
}
|
||||
)
|
||||
continue
|
||||
|
||||
hint = entry.get("install_instructions") or ""
|
||||
# Heuristic: 1-minute fixes mention an env var or API key.
|
||||
if any(k in hint.lower() for k in ["api key", "env", "_key=", "_api"]):
|
||||
@@ -406,10 +445,17 @@ class ToolRegistry:
|
||||
"capability": cap,
|
||||
"tool": entry.get("name"),
|
||||
"provider": entry.get("provider"),
|
||||
"runtime": entry.get("runtime"),
|
||||
"install_instructions": hint,
|
||||
}
|
||||
)
|
||||
|
||||
for entry in bucket.get("available", []) + bucket.get("unavailable", []):
|
||||
if entry.get("resource_profile_note"):
|
||||
runtime_warnings.append(
|
||||
f"{entry.get('name')}: {entry.get('resource_profile_note')}"
|
||||
)
|
||||
|
||||
result = {
|
||||
"composition_runtimes": comp_runtimes,
|
||||
"capabilities": capabilities,
|
||||
|
||||
+199
-15
@@ -27,6 +27,13 @@ from tools.base_tool import (
|
||||
ToolTier,
|
||||
)
|
||||
from tools._comfyui.client import ComfyUIClient, ComfyUIError
|
||||
from tools._comfyui.metadata import (
|
||||
BUNDLED_MODEL_STACKS,
|
||||
COMFYUI_SETUP_OFFER,
|
||||
missing_models_payload,
|
||||
model_stack,
|
||||
workflow_hash,
|
||||
)
|
||||
|
||||
_WORKFLOWS = Path(__file__).resolve().parent.parent / "_comfyui" / "workflows"
|
||||
|
||||
@@ -55,6 +62,34 @@ _REQUIRED_MODELS_T2V = [
|
||||
"wan2.2_t2v_lightx2v_4steps_lora_v1.1_low_noise.safetensors",
|
||||
]
|
||||
|
||||
_RESOURCE_PROFILES = {
|
||||
"provider_floor": {
|
||||
"vram_mb": 8000,
|
||||
"ram_mb": 16000,
|
||||
"applies_to": (
|
||||
"ComfyUI provider availability and low-VRAM custom workflows. "
|
||||
"Actual requirements depend on workflow_json/workflow_path."
|
||||
),
|
||||
},
|
||||
"bundled_wan22_14b_fp8": {
|
||||
"vram_mb": 16000,
|
||||
"ram_mb": 32000,
|
||||
"applies_to": (
|
||||
"Bundled WAN 2.2 14B FP8 T2V/I2V workflows. This is not a "
|
||||
"ComfyUI provider-wide requirement."
|
||||
),
|
||||
},
|
||||
"low_vram_custom_workflows": {
|
||||
"vram_mb": "8000-12000",
|
||||
"ram_mb": "16000-32000",
|
||||
"examples": [
|
||||
"Wan 2.1 1.3B",
|
||||
"LTX-Video / LTXV FP8 or quantized workflows",
|
||||
"Wan 2.2 GGUF / quantized community workflows",
|
||||
],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class ComfyUIVideo(BaseTool):
|
||||
name = "comfyui_video"
|
||||
@@ -68,18 +103,20 @@ class ComfyUIVideo(BaseTool):
|
||||
runtime = ToolRuntime.LOCAL_GPU
|
||||
|
||||
dependencies = []
|
||||
setup_offer = COMFYUI_SETUP_OFFER
|
||||
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 = []
|
||||
agent_skills = ["comfyui", "ai-video-gen", "ltx2"]
|
||||
|
||||
capabilities = ["text_to_video", "image_to_video"]
|
||||
supports = {
|
||||
"seed": True,
|
||||
"reference_image": True,
|
||||
"custom_workflow": True,
|
||||
"custom_output_node": True,
|
||||
"offline": True,
|
||||
}
|
||||
best_for = [
|
||||
@@ -87,10 +124,12 @@ class ComfyUIVideo(BaseTool):
|
||||
"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)",
|
||||
"custom low-VRAM ComfyUI workflows on 8GB-12GB GPUs",
|
||||
]
|
||||
not_good_for = [
|
||||
"setups without a running ComfyUI server",
|
||||
"CPU-only machines",
|
||||
"running the bundled WAN 2.2 14B FP8 workflows on GPUs below 16GB VRAM",
|
||||
]
|
||||
fallback = "wan_video"
|
||||
fallback_tools = ["wan_video", "hunyuan_video", "ltx_video_local", "kling_video"]
|
||||
@@ -120,13 +159,38 @@ class ComfyUIVideo(BaseTool):
|
||||
"output_path": {"type": "string", "description": "Where to save the video"},
|
||||
"workflow_json": {
|
||||
"type": "string",
|
||||
"description": "Optional full ComfyUI workflow JSON (overrides default)",
|
||||
"description": "Optional full ComfyUI workflow JSON. Requires output_node.",
|
||||
},
|
||||
"workflow_path": {
|
||||
"type": "string",
|
||||
"description": "Optional path to a ComfyUI workflow JSON file. Requires output_node.",
|
||||
},
|
||||
"output_node": {
|
||||
"type": "string",
|
||||
"description": "ComfyUI output node ID for custom workflow_json/workflow_path.",
|
||||
},
|
||||
"workflow_name": {
|
||||
"type": "string",
|
||||
"description": "Optional human-readable provenance label for a custom workflow.",
|
||||
},
|
||||
"workflow_model": {
|
||||
"type": "string",
|
||||
"description": "Optional model/provenance label for a custom workflow.",
|
||||
},
|
||||
"workflow_model_stack": {
|
||||
"type": "array",
|
||||
"description": (
|
||||
"Optional provenance metadata for custom workflow dependencies. "
|
||||
"Items should include name, role, quantization, scheduler, "
|
||||
"and LoRA strengths when known."
|
||||
),
|
||||
"items": {"type": "object"},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(
|
||||
cpu_cores=2, ram_mb=32000, vram_mb=16000, disk_mb=2000, network_required=False,
|
||||
cpu_cores=2, ram_mb=16000, vram_mb=8000, 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"]
|
||||
@@ -139,12 +203,49 @@ class ComfyUIVideo(BaseTool):
|
||||
def get_status(self) -> ToolStatus:
|
||||
if not self._client.is_available():
|
||||
return ToolStatus.UNAVAILABLE
|
||||
# Check that at least one operation has its models
|
||||
_, missing_i2v = self._client.check_models(_REQUIRED_MODELS_I2V)
|
||||
_, missing_t2v = self._client.check_models(_REQUIRED_MODELS_T2V)
|
||||
if missing_i2v and missing_t2v:
|
||||
statuses = self.operation_statuses()
|
||||
if any(status == "available" for status in statuses.values()):
|
||||
return ToolStatus.AVAILABLE
|
||||
if statuses:
|
||||
return ToolStatus.DEGRADED
|
||||
return ToolStatus.AVAILABLE
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def operation_statuses(self) -> dict[str, str]:
|
||||
"""Return per-operation readiness for selector routing and preflight."""
|
||||
if not self._client.is_available():
|
||||
return {
|
||||
"text_to_video": "unavailable",
|
||||
"image_to_video": "unavailable",
|
||||
}
|
||||
|
||||
_, missing_t2v = self._client.check_models(_REQUIRED_MODELS_T2V)
|
||||
_, missing_i2v = self._client.check_models(_REQUIRED_MODELS_I2V)
|
||||
return {
|
||||
"text_to_video": "available" if not missing_t2v else "degraded",
|
||||
"image_to_video": "available" if not missing_i2v else "degraded",
|
||||
}
|
||||
|
||||
def is_operation_available(self, operation: str) -> bool:
|
||||
if operation not in {"text_to_video", "image_to_video"}:
|
||||
return False
|
||||
return self.operation_statuses().get(operation) == "available"
|
||||
|
||||
def get_info(self) -> dict[str, Any]:
|
||||
info = super().get_info()
|
||||
info["operation_statuses"] = self.operation_statuses()
|
||||
info["resource_profiles"] = _RESOURCE_PROFILES
|
||||
info["setup_offer"] = self.setup_offer
|
||||
info["bundled_model_stacks"] = {
|
||||
"text_to_video": BUNDLED_MODEL_STACKS["wan22-t2v-4step"],
|
||||
"image_to_video": BUNDLED_MODEL_STACKS["wan22-i2v-4step"],
|
||||
}
|
||||
info["resource_profile_note"] = (
|
||||
"The top-level resource_profile is a ComfyUI provider floor, not a "
|
||||
"promise that every workflow fits 8GB VRAM. Bundled WAN 2.2 14B FP8 "
|
||||
"workflows recommend 16GB VRAM; custom low-VRAM workflows can target "
|
||||
"8GB-12GB depending on model, quantization, resolution, and frame count."
|
||||
)
|
||||
return info
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
return 0.0
|
||||
@@ -156,6 +257,16 @@ class ComfyUIVideo(BaseTool):
|
||||
return 240.0 # ~4 min
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
custom_workflow = bool(inputs.get("workflow_json") or inputs.get("workflow_path"))
|
||||
if custom_workflow and not inputs.get("output_node"):
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=(
|
||||
"Custom ComfyUI workflows require output_node so OpenMontage "
|
||||
"knows which ComfyUI node to download artifacts from."
|
||||
),
|
||||
)
|
||||
|
||||
if not self._client.is_available():
|
||||
return ToolResult(
|
||||
success=False,
|
||||
@@ -164,16 +275,27 @@ class ComfyUIVideo(BaseTool):
|
||||
|
||||
operation = inputs.get("operation", "text_to_video")
|
||||
|
||||
if not inputs.get("workflow_json"):
|
||||
if not custom_workflow:
|
||||
required = _REQUIRED_MODELS_I2V if operation == "image_to_video" else _REQUIRED_MODELS_T2V
|
||||
_, missing = self._client.check_models(required)
|
||||
if missing:
|
||||
workflow_key = (
|
||||
"wan22-i2v-4step"
|
||||
if operation == "image_to_video"
|
||||
else "wan22-t2v-4step"
|
||||
)
|
||||
return ToolResult(
|
||||
success=False,
|
||||
data=missing_models_payload(
|
||||
missing,
|
||||
workflow_key=workflow_key,
|
||||
workflow_name=f"{workflow_key}.json",
|
||||
operation=operation,
|
||||
),
|
||||
error=(
|
||||
f"ComfyUI server is running but missing models for {operation}: "
|
||||
f"{', '.join(missing)}.\n"
|
||||
f"Download them to your ComfyUI models directory."
|
||||
f"See data.missing_models for destination hints and download URLs."
|
||||
),
|
||||
)
|
||||
start = time.time()
|
||||
@@ -183,14 +305,17 @@ class ComfyUIVideo(BaseTool):
|
||||
)
|
||||
|
||||
try:
|
||||
if inputs.get("workflow_json"):
|
||||
workflow = json.loads(inputs["workflow_json"])
|
||||
output_node = _T2V_OUTPUT_NODE
|
||||
if custom_workflow:
|
||||
workflow = self._load_custom_workflow(inputs)
|
||||
output_node = str(inputs["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)
|
||||
|
||||
provenance = self._workflow_provenance(
|
||||
inputs, custom_workflow, output_node, operation, workflow
|
||||
)
|
||||
paths = self._client.generate(
|
||||
workflow,
|
||||
output_node=output_node,
|
||||
@@ -208,11 +333,12 @@ class ComfyUIVideo(BaseTool):
|
||||
height = inputs.get("height", 480 if operation == "text_to_video" else 640)
|
||||
num_frames = inputs.get("num_frames", 81)
|
||||
|
||||
model_name = self._model_name(inputs, custom_workflow)
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"provider": "comfyui",
|
||||
"model": "wan2.2-14b-fp8-4step",
|
||||
"model": model_name,
|
||||
"prompt": inputs["prompt"],
|
||||
"operation": operation,
|
||||
"width": width,
|
||||
@@ -222,12 +348,13 @@ class ComfyUIVideo(BaseTool):
|
||||
"duration_seconds": round(num_frames / 16, 2),
|
||||
"output": str(paths[0]),
|
||||
"format": "mp4",
|
||||
"workflow_provenance": provenance,
|
||||
},
|
||||
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",
|
||||
model=model_name,
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
@@ -287,3 +414,60 @@ class ComfyUIVideo(BaseTool):
|
||||
"108": {"filename_prefix": output_path.stem},
|
||||
})
|
||||
return workflow, _I2V_OUTPUT_NODE
|
||||
|
||||
@staticmethod
|
||||
def _load_custom_workflow(inputs: dict[str, Any]) -> dict:
|
||||
if inputs.get("workflow_json"):
|
||||
return json.loads(inputs["workflow_json"])
|
||||
return ComfyUIClient.load_workflow(Path(inputs["workflow_path"]))
|
||||
|
||||
@staticmethod
|
||||
def _model_name(inputs: dict[str, Any], custom_workflow: bool) -> str:
|
||||
if not custom_workflow:
|
||||
return "wan2.2-14b-fp8-4step"
|
||||
return (
|
||||
inputs.get("workflow_model")
|
||||
or inputs.get("model")
|
||||
or inputs.get("workflow_name")
|
||||
or "custom-comfyui-workflow"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _workflow_provenance(
|
||||
inputs: dict[str, Any],
|
||||
custom_workflow: bool,
|
||||
output_node: str,
|
||||
operation: str,
|
||||
workflow: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
if not custom_workflow:
|
||||
workflow_key = (
|
||||
"wan22-i2v-4step"
|
||||
if operation == "image_to_video"
|
||||
else "wan22-t2v-4step"
|
||||
)
|
||||
return {
|
||||
"source": "bundled",
|
||||
"workflow": (
|
||||
"wan22-i2v-4step.json"
|
||||
if operation == "image_to_video"
|
||||
else "wan22-t2v-4step.json"
|
||||
),
|
||||
"workflow_hash_sha256": workflow_hash(workflow),
|
||||
"model_stack": model_stack(workflow_key, inputs),
|
||||
"output_node": output_node,
|
||||
}
|
||||
return {
|
||||
"source": "user_supplied",
|
||||
"workflow_name": inputs.get("workflow_name"),
|
||||
"workflow_path": inputs.get("workflow_path"),
|
||||
"model": inputs.get("workflow_model") or inputs.get("model"),
|
||||
"workflow_hash_sha256": workflow_hash(workflow),
|
||||
"model_stack": model_stack(None, inputs),
|
||||
"model_stack_source": (
|
||||
"caller_supplied"
|
||||
if inputs.get("workflow_model_stack")
|
||||
else "unknown_custom_workflow"
|
||||
),
|
||||
"output_node": output_node,
|
||||
}
|
||||
|
||||
@@ -54,6 +54,12 @@ class VideoSelector(BaseTool):
|
||||
"enum": ["text_to_video", "image_to_video", "reference_to_video", "rank"],
|
||||
"default": "text_to_video",
|
||||
},
|
||||
"target_operation": {
|
||||
"type": "string",
|
||||
"enum": ["text_to_video", "image_to_video", "reference_to_video"],
|
||||
"description": "Operation to score when operation='rank'.",
|
||||
"default": "text_to_video",
|
||||
},
|
||||
"aspect_ratio": {
|
||||
"type": "string",
|
||||
"enum": ["16:9", "9:16", "1:1"],
|
||||
@@ -137,11 +143,13 @@ class VideoSelector(BaseTool):
|
||||
def execute(self, inputs: dict[str, object]) -> ToolResult:
|
||||
from lib.scoring import rank_providers
|
||||
|
||||
task_context = self._prepare_task_context(inputs)
|
||||
candidates = self._providers()
|
||||
|
||||
# Rank mode — return scored provider rankings without generating
|
||||
if inputs.get("operation") == "rank":
|
||||
rank_inputs = self._rank_inputs(inputs)
|
||||
task_context = self._prepare_task_context(rank_inputs)
|
||||
candidates = self._filter_candidates(rank_inputs, candidates)
|
||||
rankings = rank_providers(candidates, task_context)
|
||||
return ToolResult(
|
||||
success=True,
|
||||
@@ -153,6 +161,7 @@ class VideoSelector(BaseTool):
|
||||
)
|
||||
|
||||
# Normal generation — use scored selection
|
||||
task_context = self._prepare_task_context(inputs)
|
||||
tool, score = self._select_best_tool(inputs, candidates, task_context)
|
||||
if tool is None:
|
||||
return ToolResult(success=False, error="No video generation provider available.")
|
||||
@@ -252,6 +261,12 @@ class VideoSelector(BaseTool):
|
||||
operation=str(inputs.get("operation", "text_to_video")),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _rank_inputs(inputs: dict[str, object]) -> dict[str, object]:
|
||||
rank_inputs = dict(inputs)
|
||||
rank_inputs["operation"] = inputs.get("target_operation", "text_to_video")
|
||||
return rank_inputs
|
||||
|
||||
@staticmethod
|
||||
def _tool_context_payload(tool: BaseTool) -> dict[str, object]:
|
||||
info = tool.get_info()
|
||||
@@ -285,23 +300,36 @@ class VideoSelector(BaseTool):
|
||||
) -> list[BaseTool]:
|
||||
operation = inputs.get("operation", "text_to_video")
|
||||
if operation == "rank":
|
||||
return candidates
|
||||
operation = inputs.get("target_operation", "text_to_video")
|
||||
|
||||
filtered: list[BaseTool] = []
|
||||
matched_operation = False
|
||||
for tool in candidates:
|
||||
supports = getattr(tool, "supports", {})
|
||||
props = getattr(tool, "input_schema", {}).get("properties", {})
|
||||
|
||||
if operation == "image_to_video":
|
||||
if supports.get("image_to_video") or "image_url" in props or "reference_image_url" in props:
|
||||
filtered.append(tool)
|
||||
matched_operation = True
|
||||
if self._operation_ready(tool, "image_to_video"):
|
||||
filtered.append(tool)
|
||||
continue
|
||||
|
||||
if operation == "reference_to_video":
|
||||
if supports.get("reference_to_video") or "reference_image_urls" in props:
|
||||
matched_operation = True
|
||||
filtered.append(tool)
|
||||
continue
|
||||
|
||||
filtered.append(tool)
|
||||
matched_operation = True
|
||||
if self._operation_ready(tool, str(operation)):
|
||||
filtered.append(tool)
|
||||
|
||||
return filtered or candidates
|
||||
return filtered if matched_operation else candidates
|
||||
|
||||
@staticmethod
|
||||
def _operation_ready(tool: BaseTool, operation: str) -> bool:
|
||||
checker = getattr(tool, "is_operation_available", None)
|
||||
if not callable(checker):
|
||||
return True
|
||||
return bool(checker(operation))
|
||||
|
||||
Reference in New Issue
Block a user