From 2f114682e8897e3956a03a9926e47eed17955cfb Mon Sep 17 00:00:00 2001 From: Ntsako Date: Thu, 6 Aug 2026 12:49:57 +0200 Subject: [PATCH] feat: support per-capability ComfyUI server URLs for image/video MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Resolves the "multi-server" open question from the adapter plan. ComfyUIClient(capability="image"|"video") now resolves its server URL from COMFYUI_IMAGE_SERVER_URL / COMFYUI_VIDEO_SERVER_URL first, falling back to the shared COMFYUI_SERVER_URL and then the localhost default — so comfyui_image and comfyui_video can point at separate ComfyUI instances (different GPUs, different model sets) with zero extra config for single-server setups. is_default_url/unavailable_reason() and the setup_offer metadata account for the override. Co-Authored-By: Claude Sonnet 5 --- docs/comfyui-adapter-plan.md | 24 +++++++- tests/contracts/test_comfyui_tools.py | 79 +++++++++++++++++++++++++++ tools/_comfyui/client.py | 63 ++++++++++++++------- tools/_comfyui/metadata.py | 7 +++ tools/graphics/comfyui_image.py | 6 +- tools/video/comfyui_video.py | 6 +- 6 files changed, 158 insertions(+), 27 deletions(-) diff --git a/docs/comfyui-adapter-plan.md b/docs/comfyui-adapter-plan.md index c8b7f79e..c8eef2e3 100644 --- a/docs/comfyui-adapter-plan.md +++ b/docs/comfyui-adapter-plan.md @@ -373,6 +373,17 @@ COMFYUI_POLL_TIMEOUT=600 # max wait for image gen COMFYUI_VIDEO_TIMEOUT=900 # max wait for video gen ``` +**Multi-server (optional):** point `comfyui_image` and `comfyui_video` at +separate ComfyUI instances -- e.g. one GPU running FLUX 2, another running +WAN 2.2 -- by setting a per-capability override. Each takes priority over +`COMFYUI_SERVER_URL` for its own tool only; leave both unset and everything +still talks to the single shared server. + +```bash +COMFYUI_IMAGE_SERVER_URL=http://gpu-a:8188 +COMFYUI_VIDEO_SERVER_URL=http://gpu-b:8188 +``` + **For Docker Compose setups** (ComfyUI in a container): ```bash @@ -485,8 +496,17 @@ pipeline definition, or any schema. `poll()` REST loop when it isn't installed or the connection fails — `resume_prompt_id` recovery behaves identically either way. -3. **Multi-server:** Should the adapter support multiple ComfyUI instances - (e.g., one for images, one for video) via per-capability URLs? +3. ~~**Multi-server:**~~ **Resolved.** `ComfyUIClient(capability="image"|"video")` + resolves its server URL from a per-capability env var first + (`COMFYUI_IMAGE_SERVER_URL` / `COMFYUI_VIDEO_SERVER_URL`), then the shared + `COMFYUI_SERVER_URL`, then the `http://localhost:8188` default. `comfyui_image` + and `comfyui_video` pass their capability at construction, so image and video + generation can point at different ComfyUI instances (different GPUs, different + model sets) with zero code changes -- single-server setups need no extra + configuration since both env vars are optional. `client.capability`/ + `client.is_default_url`/`client.unavailable_reason()` all account for the + override, and `COMFYUI_SETUP_OFFER.per_capability_env_var_overrides` documents + it for the setup-offer surfacing in `provider_menu()`. 4. **Music generation:** ACE-Step works in ComfyUI but OpenMontage needs a dedicated music-generation routing contract before adding `comfyui_music`. diff --git a/tests/contracts/test_comfyui_tools.py b/tests/contracts/test_comfyui_tools.py index 21dc5166..44c451c0 100644 --- a/tests/contracts/test_comfyui_tools.py +++ b/tests/contracts/test_comfyui_tools.py @@ -369,6 +369,85 @@ class TestClientHelpers: assert seen["client_id"] == client.client_id +class TestMultiServer: + + def test_capability_env_var_takes_priority_over_shared(self, monkeypatch): + from tools._comfyui.client import ComfyUIClient + + monkeypatch.setenv("COMFYUI_SERVER_URL", "http://shared:8188") + monkeypatch.setenv("COMFYUI_VIDEO_SERVER_URL", "http://video-gpu:8188") + client = ComfyUIClient(capability="video") + assert client.server_url == "http://video-gpu:8188" + + def test_falls_back_to_shared_when_capability_var_unset(self, monkeypatch): + from tools._comfyui.client import ComfyUIClient + + monkeypatch.setenv("COMFYUI_SERVER_URL", "http://shared:8188") + monkeypatch.delenv("COMFYUI_IMAGE_SERVER_URL", raising=False) + client = ComfyUIClient(capability="image") + assert client.server_url == "http://shared:8188" + + def test_other_capability_env_var_does_not_leak_across_tools(self, monkeypatch): + from tools._comfyui.client import ComfyUIClient + + monkeypatch.delenv("COMFYUI_SERVER_URL", raising=False) + monkeypatch.setenv("COMFYUI_IMAGE_SERVER_URL", "http://image-gpu:8188") + monkeypatch.delenv("COMFYUI_VIDEO_SERVER_URL", raising=False) + video_client = ComfyUIClient(capability="video") + assert video_client.server_url == "http://localhost:8188" + + def test_explicit_server_url_wins_over_capability_env_var(self, monkeypatch): + from tools._comfyui.client import ComfyUIClient + + monkeypatch.setenv("COMFYUI_VIDEO_SERVER_URL", "http://video-gpu:8188") + client = ComfyUIClient("http://explicit:1234", capability="video") + assert client.server_url == "http://explicit:1234" + + def test_no_capability_behaves_as_before(self, monkeypatch): + from tools._comfyui.client import ComfyUIClient + + monkeypatch.setenv("COMFYUI_SERVER_URL", "http://shared:8188") + client = ComfyUIClient() + assert client.server_url == "http://shared:8188" + assert client.is_default_url is False + + def test_is_default_url_true_only_when_both_vars_unset(self, monkeypatch): + from tools._comfyui.client import ComfyUIClient + + monkeypatch.delenv("COMFYUI_SERVER_URL", raising=False) + monkeypatch.delenv("COMFYUI_IMAGE_SERVER_URL", raising=False) + client = ComfyUIClient(capability="image") + assert client.is_default_url is True + + monkeypatch.setenv("COMFYUI_IMAGE_SERVER_URL", "http://image-gpu:8188") + client2 = ComfyUIClient(capability="image") + assert client2.is_default_url is False + + def test_unavailable_reason_mentions_capability_and_shared_var(self, monkeypatch): + from tools._comfyui.client import ComfyUIClient + + monkeypatch.delenv("COMFYUI_SERVER_URL", raising=False) + monkeypatch.delenv("COMFYUI_VIDEO_SERVER_URL", raising=False) + client = ComfyUIClient(capability="video") + msg = client.unavailable_reason() + assert "COMFYUI_VIDEO_SERVER_URL" in msg + assert "COMFYUI_SERVER_URL" in msg + + def test_image_and_video_tools_use_independent_servers(self, monkeypatch): + from tools.graphics.comfyui_image import ComfyUIImage + from tools.video.comfyui_video import ComfyUIVideo + + monkeypatch.delenv("COMFYUI_SERVER_URL", raising=False) + monkeypatch.setenv("COMFYUI_IMAGE_SERVER_URL", "http://image-gpu:8188") + monkeypatch.setenv("COMFYUI_VIDEO_SERVER_URL", "http://video-gpu:8188") + + image_tool = ComfyUIImage() + video_tool = ComfyUIVideo() + + assert image_tool._client.server_url == "http://image-gpu:8188" + assert video_tool._client.server_url == "http://video-gpu:8188" + + class _FakeWSTimeout(Exception): pass diff --git a/tools/_comfyui/client.py b/tools/_comfyui/client.py index 85b69421..6632bfa6 100644 --- a/tools/_comfyui/client.py +++ b/tools/_comfyui/client.py @@ -42,23 +42,58 @@ class ComfyUIClient: 4. POST /upload/image → stage a local image for I2V workflows """ - def __init__(self, server_url: str | None = None) -> None: - self.server_url = ( - server_url - or os.environ.get("COMFYUI_SERVER_URL", "http://localhost:8188") - ).rstrip("/") + def __init__( + self, server_url: str | None = None, capability: str | None = None + ) -> None: + """*capability*, if given (e.g. ``"image"``, ``"video"``), lets a + per-capability env var (``COMFYUI_{CAPABILITY}_SERVER_URL``) point + this client at its own ComfyUI instance -- useful when image and + video generation are split across separate servers/GPUs. Falls back + to the shared ``COMFYUI_SERVER_URL`` when the capability-specific + var isn't set, so single-server setups need no extra configuration. + """ + self.capability = capability + self._capability_env_var = ( + f"COMFYUI_{capability.upper()}_SERVER_URL" if capability else None + ) + resolved = server_url or self._capability_url() or os.environ.get( + "COMFYUI_SERVER_URL" + ) + self.server_url = (resolved or "http://localhost:8188").rstrip("/") # Scopes websocket execution events to this client (see wait_ws) and # is echoed back on /prompt so the server targets messages to us. self.client_id = str(uuid.uuid4()) + def _capability_url(self) -> str | None: + if self._capability_env_var: + return os.environ.get(self._capability_env_var) + return None + # ------------------------------------------------------------------ # 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") + """True if neither the capability-specific nor shared env var is set.""" + return not (self._capability_url() or os.environ.get("COMFYUI_SERVER_URL")) + + def unavailable_reason(self) -> str: + """Human-readable explanation of why the server can't be reached.""" + env_var_hint = self._capability_env_var or "COMFYUI_SERVER_URL" + if self._capability_env_var: + env_var_hint += " (or the shared COMFYUI_SERVER_URL)" + if self.is_default_url: + return ( + f"No ComfyUI server found at {self.server_url} " + f"(default — no server URL configured).\n" + f"Set {env_var_hint} 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." + ) def is_available(self) -> bool: """Return True if the ComfyUI server is reachable.""" @@ -70,20 +105,6 @@ 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 # ------------------------------------------------------------------ diff --git a/tools/_comfyui/metadata.py b/tools/_comfyui/metadata.py index dcca514b..75707868 100644 --- a/tools/_comfyui/metadata.py +++ b/tools/_comfyui/metadata.py @@ -18,6 +18,13 @@ COMFYUI_SETUP_OFFER: dict[str, Any] = { "free local video generation through ComfyUI workflows", "community workflow_json/workflow_path execution", ], + # Optional: point image/video generation at separate ComfyUI instances + # (e.g. different GPUs). Each overrides COMFYUI_SERVER_URL for its own + # tool only; single-server setups can ignore this entirely. + "per_capability_env_var_overrides": { + "comfyui_image": "COMFYUI_IMAGE_SERVER_URL", + "comfyui_video": "COMFYUI_VIDEO_SERVER_URL", + }, } diff --git a/tools/graphics/comfyui_image.py b/tools/graphics/comfyui_image.py index 4c55d94b..e91e71c3 100644 --- a/tools/graphics/comfyui_image.py +++ b/tools/graphics/comfyui_image.py @@ -58,7 +58,9 @@ class ComfyUIImage(BaseTool): install_instructions = ( "Start a ComfyUI server and set COMFYUI_SERVER_URL " "(default http://localhost:8188).\n" - "See https://github.com/comfyanonymous/ComfyUI for setup." + "See https://github.com/comfyanonymous/ComfyUI for setup.\n" + "Running a separate ComfyUI instance for images? Set COMFYUI_IMAGE_SERVER_URL " + "instead -- it takes priority over COMFYUI_SERVER_URL for this tool only." ) agent_skills = ["comfyui", "flux-best-practices"] @@ -133,7 +135,7 @@ class ComfyUIImage(BaseTool): user_visible_verification = ["Inspect generated image for quality and prompt adherence"] def __init__(self) -> None: - self._client = ComfyUIClient() + self._client = ComfyUIClient(capability="image") def get_status(self) -> ToolStatus: if not self._client.is_available(): diff --git a/tools/video/comfyui_video.py b/tools/video/comfyui_video.py index 84270c27..85ed3717 100644 --- a/tools/video/comfyui_video.py +++ b/tools/video/comfyui_video.py @@ -107,7 +107,9 @@ class ComfyUIVideo(BaseTool): 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." + "Requires WAN 2.2 models and LightX2V LoRAs in ComfyUI's model directory.\n" + "Running a separate ComfyUI instance for video? Set COMFYUI_VIDEO_SERVER_URL " + "instead -- it takes priority over COMFYUI_SERVER_URL for this tool only." ) agent_skills = ["comfyui", "ai-video-gen", "ltx2"] @@ -216,7 +218,7 @@ class ComfyUIVideo(BaseTool): user_visible_verification = ["Watch generated clip for motion coherence and artifacts"] def __init__(self) -> None: - self._client = ComfyUIClient() + self._client = ComfyUIClient(capability="video") self._last_progress_log = 0.0 def _log_progress(self, data: dict) -> None: