diff --git a/tests/contracts/test_comfyui_tools.py b/tests/contracts/test_comfyui_tools.py index 3c68fdc..98ab5b3 100644 --- a/tests/contracts/test_comfyui_tools.py +++ b/tests/contracts/test_comfyui_tools.py @@ -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) diff --git a/tools/_comfyui/client.py b/tools/_comfyui/client.py index 62f7ca0..8362623 100644 --- a/tools/_comfyui/client.py +++ b/tools/_comfyui/client.py @@ -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 # ------------------------------------------------------------------ diff --git a/tools/audio/comfyui_music.py b/tools/audio/comfyui_music.py index 3d95672..ebff88f 100644 --- a/tools/audio/comfyui_music.py +++ b/tools/audio/comfyui_music.py @@ -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) diff --git a/tools/graphics/comfyui_image.py b/tools/graphics/comfyui_image.py index 002723a..53ce8af 100644 --- a/tools/graphics/comfyui_image.py +++ b/tools/graphics/comfyui_image.py @@ -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) diff --git a/tools/video/comfyui_video.py b/tools/video/comfyui_video.py index c17f7f6..2cf88c6 100644 --- a/tools/video/comfyui_video.py +++ b/tools/video/comfyui_video.py @@ -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(