comfyui: satisfy provider contract review items
This commit is contained in:
+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