comfyui: add model discovery and actionable error messages
- Client queries ComfyUI /object_info to discover installed models (checkpoints, diffusion models, VAE, CLIP, LoRAs) - Each tool declares its required models and checks them on execute() - get_status() returns DEGRADED when server is up but models are missing - Clear error messages tell the user exactly which models to download - When COMFYUI_SERVER_URL is not set, error message tells the user to configure it in .env instead of silently failing on localhost:8188 - 8 new tests covering URL config, error messages, and model requirements 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
6ec2bbb090
commit
e3947e1f11
@@ -171,3 +171,58 @@ class TestClientHelpers:
|
||||
for _ in range(100):
|
||||
s = ComfyUIClient.random_seed()
|
||||
assert 0 <= s < 2**32
|
||||
|
||||
def test_is_default_url_when_env_not_set(self, monkeypatch):
|
||||
from tools._comfyui.client import ComfyUIClient
|
||||
monkeypatch.delenv("COMFYUI_SERVER_URL", raising=False)
|
||||
client = ComfyUIClient()
|
||||
assert client.is_default_url is True
|
||||
|
||||
def test_is_not_default_url_when_env_set(self, monkeypatch):
|
||||
from tools._comfyui.client import ComfyUIClient
|
||||
monkeypatch.setenv("COMFYUI_SERVER_URL", "http://myhost:9999")
|
||||
client = ComfyUIClient()
|
||||
assert client.is_default_url is False
|
||||
|
||||
def test_unavailable_reason_default_url(self, monkeypatch):
|
||||
from tools._comfyui.client import ComfyUIClient
|
||||
monkeypatch.delenv("COMFYUI_SERVER_URL", raising=False)
|
||||
client = ComfyUIClient()
|
||||
msg = client.unavailable_reason()
|
||||
assert "COMFYUI_SERVER_URL" in msg
|
||||
assert ".env" in msg
|
||||
|
||||
def test_unavailable_reason_custom_url(self, monkeypatch):
|
||||
from tools._comfyui.client import ComfyUIClient
|
||||
monkeypatch.setenv("COMFYUI_SERVER_URL", "http://myhost:9999")
|
||||
client = ComfyUIClient()
|
||||
msg = client.unavailable_reason()
|
||||
assert "myhost:9999" in msg
|
||||
assert "COMFYUI_SERVER_URL" not in msg
|
||||
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Model discovery (offline, no server needed)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
class TestModelRequirements:
|
||||
|
||||
def test_image_tool_has_required_models(self):
|
||||
from tools.graphics.comfyui_image import _REQUIRED_MODELS
|
||||
assert len(_REQUIRED_MODELS) > 0
|
||||
assert any("flux" in m.lower() for m in _REQUIRED_MODELS)
|
||||
|
||||
def test_video_tool_has_required_models_i2v(self):
|
||||
from tools.video.comfyui_video import _REQUIRED_MODELS_I2V
|
||||
assert len(_REQUIRED_MODELS_I2V) > 0
|
||||
assert any("i2v" in m.lower() for m in _REQUIRED_MODELS_I2V)
|
||||
|
||||
def test_video_tool_has_required_models_t2v(self):
|
||||
from tools.video.comfyui_video import _REQUIRED_MODELS_T2V
|
||||
assert len(_REQUIRED_MODELS_T2V) > 0
|
||||
assert any("t2v" in m.lower() for m in _REQUIRED_MODELS_T2V)
|
||||
|
||||
def test_music_tool_has_required_models(self):
|
||||
from tools.audio.comfyui_music import _REQUIRED_MODELS
|
||||
assert len(_REQUIRED_MODELS) > 0
|
||||
assert any("ace" in m.lower() for m in _REQUIRED_MODELS)
|
||||
|
||||
@@ -41,6 +41,11 @@ class ComfyUIClient:
|
||||
# Health
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@property
|
||||
def is_default_url(self) -> bool:
|
||||
"""True if using the fallback URL (user didn't set COMFYUI_SERVER_URL)."""
|
||||
return not os.environ.get("COMFYUI_SERVER_URL")
|
||||
|
||||
def is_available(self) -> bool:
|
||||
"""Return True if the ComfyUI server is reachable."""
|
||||
try:
|
||||
@@ -51,6 +56,79 @@ class ComfyUIClient:
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def unavailable_reason(self) -> str:
|
||||
"""Human-readable explanation of why the server can't be reached."""
|
||||
if self.is_default_url:
|
||||
return (
|
||||
f"No ComfyUI server found at {self.server_url} "
|
||||
f"(default — no COMFYUI_SERVER_URL configured).\n"
|
||||
f"Set COMFYUI_SERVER_URL in your .env file to the address of "
|
||||
f"your ComfyUI server (e.g. http://localhost:8188)."
|
||||
)
|
||||
return (
|
||||
f"ComfyUI server not reachable at {self.server_url}.\n"
|
||||
f"Check that ComfyUI is running and the URL is correct."
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Model discovery
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def list_models(self) -> dict[str, list[str]]:
|
||||
"""Query ComfyUI for available models, grouped by type.
|
||||
|
||||
Returns a dict like::
|
||||
|
||||
{
|
||||
"checkpoints": ["sd_xl_base.safetensors", ...],
|
||||
"diffusion_models": ["flux2-dev-nvfp4.safetensors", ...],
|
||||
"vae": ["ae.safetensors", ...],
|
||||
"clip": ["clip_l.safetensors", ...],
|
||||
"loras": ["my_lora.safetensors", ...],
|
||||
}
|
||||
"""
|
||||
node_to_key = {
|
||||
"CheckpointLoaderSimple": ("ckpt_name", "checkpoints"),
|
||||
"UNETLoader": ("unet_name", "diffusion_models"),
|
||||
"VAELoader": ("vae_name", "vae"),
|
||||
"CLIPLoader": ("clip_name", "clip"),
|
||||
"LoraLoaderModelOnly": ("lora_name", "loras"),
|
||||
}
|
||||
result: dict[str, list[str]] = {}
|
||||
for node_class, (field, group) in node_to_key.items():
|
||||
try:
|
||||
resp = requests.get(
|
||||
f"{self.server_url}/object_info/{node_class}", timeout=10
|
||||
)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
options = (
|
||||
data.get(node_class, {})
|
||||
.get("input", {})
|
||||
.get("required", {})
|
||||
.get(field, [[]])[0]
|
||||
)
|
||||
if isinstance(options, list):
|
||||
result[group] = options
|
||||
except Exception:
|
||||
result[group] = []
|
||||
return result
|
||||
|
||||
def check_models(
|
||||
self, required: list[str]
|
||||
) -> tuple[list[str], list[str]]:
|
||||
"""Check which of *required* model filenames are available.
|
||||
|
||||
Returns ``(found, missing)`` — two lists of filenames.
|
||||
"""
|
||||
all_models: set[str] = set()
|
||||
for names in self.list_models().values():
|
||||
all_models.update(names)
|
||||
|
||||
found = [m for m in required if m in all_models]
|
||||
missing = [m for m in required if m not in all_models]
|
||||
return found, missing
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Core cycle
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@@ -30,6 +30,10 @@ _WORKFLOWS = Path(__file__).resolve().parent.parent / "_comfyui" / "workflows"
|
||||
|
||||
_OUTPUT_NODE = "3"
|
||||
|
||||
_REQUIRED_MODELS = [
|
||||
"ace_step_v1_3.5b.safetensors",
|
||||
]
|
||||
|
||||
|
||||
class ComfyUIMusic(BaseTool):
|
||||
name = "comfyui_music"
|
||||
@@ -117,9 +121,12 @@ class ComfyUIMusic(BaseTool):
|
||||
self._client = ComfyUIClient()
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
if self._client.is_available():
|
||||
return ToolStatus.AVAILABLE
|
||||
return ToolStatus.UNAVAILABLE
|
||||
if not self._client.is_available():
|
||||
return ToolStatus.UNAVAILABLE
|
||||
_, missing = self._client.check_models(_REQUIRED_MODELS)
|
||||
if missing:
|
||||
return ToolStatus.DEGRADED
|
||||
return ToolStatus.AVAILABLE
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
return 0.0
|
||||
@@ -132,9 +139,21 @@ class ComfyUIMusic(BaseTool):
|
||||
if not self._client.is_available():
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="ComfyUI server not reachable. " + self.install_instructions,
|
||||
error=self._client.unavailable_reason(),
|
||||
)
|
||||
|
||||
if not inputs.get("workflow_json"):
|
||||
_, missing = self._client.check_models(_REQUIRED_MODELS)
|
||||
if missing:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=(
|
||||
f"ComfyUI server is running but missing required models: "
|
||||
f"{', '.join(missing)}.\n"
|
||||
f"Download them to your ComfyUI checkpoints directory."
|
||||
),
|
||||
)
|
||||
|
||||
start = time.time()
|
||||
seed = inputs.get("seed") or ComfyUIClient.random_seed()
|
||||
duration = inputs.get("duration", 30.0)
|
||||
|
||||
@@ -27,6 +27,13 @@ from tools._comfyui.client import ComfyUIClient, ComfyUIError
|
||||
|
||||
_WORKFLOWS = Path(__file__).resolve().parent.parent / "_comfyui" / "workflows"
|
||||
|
||||
# Models required by the bundled flux2-txt2img workflow
|
||||
_REQUIRED_MODELS = [
|
||||
"flux2-dev-nvfp4.safetensors",
|
||||
"mistral_3_small_flux2_fp4_mixed.safetensors",
|
||||
"flux2-vae.safetensors",
|
||||
]
|
||||
|
||||
|
||||
class ComfyUIImage(BaseTool):
|
||||
name = "comfyui_image"
|
||||
@@ -96,9 +103,12 @@ class ComfyUIImage(BaseTool):
|
||||
self._client = ComfyUIClient()
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
if self._client.is_available():
|
||||
return ToolStatus.AVAILABLE
|
||||
return ToolStatus.UNAVAILABLE
|
||||
if not self._client.is_available():
|
||||
return ToolStatus.UNAVAILABLE
|
||||
_, missing = self._client.check_models(_REQUIRED_MODELS)
|
||||
if missing:
|
||||
return ToolStatus.DEGRADED
|
||||
return ToolStatus.AVAILABLE
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
return 0.0
|
||||
@@ -110,9 +120,21 @@ class ComfyUIImage(BaseTool):
|
||||
if not self._client.is_available():
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="ComfyUI server not reachable. " + self.install_instructions,
|
||||
error=self._client.unavailable_reason(),
|
||||
)
|
||||
|
||||
if not inputs.get("workflow_json"):
|
||||
_, missing = self._client.check_models(_REQUIRED_MODELS)
|
||||
if missing:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=(
|
||||
f"ComfyUI server is running but missing required models: "
|
||||
f"{', '.join(missing)}.\n"
|
||||
f"Download them to your ComfyUI models directory."
|
||||
),
|
||||
)
|
||||
|
||||
start = time.time()
|
||||
seed = inputs.get("seed") or ComfyUIClient.random_seed()
|
||||
width = inputs.get("width", 1024)
|
||||
|
||||
@@ -34,6 +34,27 @@ _WORKFLOWS = Path(__file__).resolve().parent.parent / "_comfyui" / "workflows"
|
||||
_T2V_OUTPUT_NODE = "16"
|
||||
_I2V_OUTPUT_NODE = "108"
|
||||
|
||||
# Models required by the bundled WAN 2.2 workflows
|
||||
_REQUIRED_MODELS_COMMON = [
|
||||
"umt5_xxl_fp8_e4m3fn_scaled.safetensors",
|
||||
]
|
||||
_REQUIRED_MODELS_I2V = [
|
||||
*_REQUIRED_MODELS_COMMON,
|
||||
"wan2.2_i2v_high_noise_14B_fp8_scaled.safetensors",
|
||||
"wan2.2_i2v_low_noise_14B_fp8_scaled.safetensors",
|
||||
"wan_2.1_vae.safetensors",
|
||||
"wan2.2_i2v_lightx2v_4steps_lora_v1_high_noise.safetensors",
|
||||
"wan2.2_i2v_lightx2v_4steps_lora_v1_low_noise.safetensors",
|
||||
]
|
||||
_REQUIRED_MODELS_T2V = [
|
||||
*_REQUIRED_MODELS_COMMON,
|
||||
"wan2.2_t2v_high_noise_14B_fp8_scaled.safetensors",
|
||||
"wan2.2_t2v_low_noise_14B_fp8_scaled.safetensors",
|
||||
"wan2.2_vae.safetensors",
|
||||
"wan2.2_t2v_lightx2v_4steps_lora_v1.1_high_noise.safetensors",
|
||||
"wan2.2_t2v_lightx2v_4steps_lora_v1.1_low_noise.safetensors",
|
||||
]
|
||||
|
||||
|
||||
class ComfyUIVideo(BaseTool):
|
||||
name = "comfyui_video"
|
||||
@@ -116,9 +137,14 @@ class ComfyUIVideo(BaseTool):
|
||||
self._client = ComfyUIClient()
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
if self._client.is_available():
|
||||
return ToolStatus.AVAILABLE
|
||||
return ToolStatus.UNAVAILABLE
|
||||
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:
|
||||
return ToolStatus.DEGRADED
|
||||
return ToolStatus.AVAILABLE
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
return 0.0
|
||||
@@ -133,10 +159,23 @@ class ComfyUIVideo(BaseTool):
|
||||
if not self._client.is_available():
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="ComfyUI server not reachable. " + self.install_instructions,
|
||||
error=self._client.unavailable_reason(),
|
||||
)
|
||||
|
||||
operation = inputs.get("operation", "text_to_video")
|
||||
|
||||
if not inputs.get("workflow_json"):
|
||||
required = _REQUIRED_MODELS_I2V if operation == "image_to_video" else _REQUIRED_MODELS_T2V
|
||||
_, missing = self._client.check_models(required)
|
||||
if missing:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=(
|
||||
f"ComfyUI server is running but missing models for {operation}: "
|
||||
f"{', '.join(missing)}.\n"
|
||||
f"Download them to your ComfyUI models directory."
|
||||
),
|
||||
)
|
||||
start = time.time()
|
||||
seed = inputs.get("seed") or ComfyUIClient.random_seed()
|
||||
output_path = Path(
|
||||
|
||||
Reference in New Issue
Block a user