comfyui: satisfy provider contract review items

This commit is contained in:
calesthio
2026-04-23 17:49:04 -07:00
committed by Alastair Beal
parent 4c62186c95
commit 7c4bb08890
9 changed files with 1152 additions and 71 deletions
+199 -15
View File
@@ -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,
}
+33 -5
View File
@@ -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))