From 0cdec8be5a9b13913eef691c236671cb2ec84d64 Mon Sep 17 00:00:00 2001 From: Ken Simpson Date: Mon, 17 Aug 2026 21:53:29 -0700 Subject: [PATCH] Support ComfyUI node replacements remotely --- README.md | 2 +- comfyui_modal_sync_cloud.py | 11 +++++++++-- docs/LEARNINGS.md | 1 + tests/test_modal_executor_node.py | 20 +++++++++++++++++++- 4 files changed, 30 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index c87961c..b8cbf43 100644 --- a/README.md +++ b/README.md @@ -56,7 +56,7 @@ Each ComfyUI environment receives its own Modal app name on first startup. Modal For repository development, `uv sync --extra remote --group test` installs the same pinned SDK. Remote mode uses the stable cloud entrypoint in [`comfyui_modal_sync_cloud.py`](comfyui_modal_sync_cloud.py). On first use, Modal-Sync can auto-deploy the configured Modal app if it does not exist. -The deployed image uses Python 3.11 plus an exact ComfyUI support and CUDA package set, including ComfyUI's current `comfy-aimdo==0.4.13` and `comfy-kitchen==0.2.31` pins and the import-time dependencies used by built-in extras such as Math Expression and GLSL. Remote prompt executors mirror both legacy scalar and current active/inactive RAM-pressure cache arguments from ComfyUI, and the cache adapter supports both synchronous legacy access and current coroutine-based cache operations. Remote pre-execution validation finalizes V3 dynamic input schemas against the live prompt before checking required sockets, so nodes using `io.Autogrow` accept expanded paths such as `images.image0` exactly as local ComfyUI does. The image's local build context is limited to the ComfyUI source packages, top-level Python modules, and runtime configuration needed by the headless worker; model directories, custom nodes, caches, tests, virtual environments, user data, and unknown top-level directories stay out of the image snapshot. Before every process's first remote invocation, Modal-Sync compares the deployed worker's runtime fingerprint with the local source, ComfyUI source, custom-node requirements, and runtime-shaping settings. A missing or mismatched fingerprint is treated as stale and replaced automatically when `COMFY_MODAL_AUTO_DEPLOY=true`; replacement uses Modal's non-interactive SDK stop API when available and an explicitly confirmed CLI fallback, so app shutdown cannot stall on a hidden terminal prompt. +The deployed image uses Python 3.11 plus an exact ComfyUI support and CUDA package set, including ComfyUI's current `comfy-aimdo==0.4.13` and `comfy-kitchen==0.2.31` pins and the import-time dependencies used by built-in extras such as Math Expression and GLSL. Its headless PromptServer shim also instantiates ComfyUI's `NodeReplaceManager`, allowing current built-in and custom extensions to register replacement nodes during remote initialization. Remote prompt executors mirror both legacy scalar and current active/inactive RAM-pressure cache arguments from ComfyUI, and the cache adapter supports both synchronous legacy access and current coroutine-based cache operations. Remote pre-execution validation finalizes V3 dynamic input schemas against the live prompt before checking required sockets, so nodes using `io.Autogrow` accept expanded paths such as `images.image0` exactly as local ComfyUI does. The image's local build context is limited to the ComfyUI source packages, top-level Python modules, and runtime configuration needed by the headless worker; model directories, custom nodes, caches, tests, virtual environments, user data, and unknown top-level directories stay out of the image snapshot. Before every process's first remote invocation, Modal-Sync compares the deployed worker's runtime fingerprint with the local source, ComfyUI source, custom-node requirements, and runtime-shaping settings. A missing or mismatched fingerprint is treated as stale and replaced automatically when `COMFY_MODAL_AUTO_DEPLOY=true`; replacement uses Modal's non-interactive SDK stop API when available and an explicitly confirmed CLI fallback, so app shutdown cannot stall on a hidden terminal prompt. ## Using It In ComfyUI diff --git a/comfyui_modal_sync_cloud.py b/comfyui_modal_sync_cloud.py index 5edcccb..8b00141 100644 --- a/comfyui_modal_sync_cloud.py +++ b/comfyui_modal_sync_cloud.py @@ -1900,12 +1900,13 @@ class _HeadlessPromptQueue: class _HeadlessPromptServerInstance: """Minimal PromptServer.instance replacement for custom-node import side effects.""" - def __init__(self) -> None: + def __init__(self, node_replace_manager: Any | None) -> None: """Initialize route registration and no-op websocket state.""" from aiohttp import web self.routes = web.RouteTableDef() self.app = web.Application() + self.node_replace_manager = node_replace_manager self.supports = ["custom_nodes_from_web"] self.client_id: str | None = None self.last_node_id: str | None = None @@ -2647,7 +2648,13 @@ def _ensure_headless_prompt_server_instance() -> None: if getattr(prompt_server_class, "instance", None) is not None: return - prompt_server_class.instance = _HeadlessPromptServerInstance() + node_replace_manager_class = getattr(server, "NodeReplaceManager", None) + node_replace_manager = ( + node_replace_manager_class() + if callable(node_replace_manager_class) + else None + ) + prompt_server_class.instance = _HeadlessPromptServerInstance(node_replace_manager) logger.info("Installed headless PromptServer.instance for remote custom-node initialization.") diff --git a/docs/LEARNINGS.md b/docs/LEARNINGS.md index ed827bb..1dc9c6b 100644 --- a/docs/LEARNINGS.md +++ b/docs/LEARNINGS.md @@ -364,3 +364,4 @@ - Modal does not expose authoritative per-prompt billing in real time, but the active-container poll can support a clearly labeled GPU-only estimate by integrating each running container's published per-second rate over the active prompt window. Filter cost attribution to the prompt's selected GPU because the task list is instance-wide, retain disappeared-container intervals in the accumulated total, reset on prompt changes, and disclose the pricing snapshot date plus excluded CPU, memory, storage, credits, reservations, and post-prompt warm time. - Viewer-stable canvas text must not reserve graph-space width beside a zoom-scaled progress meter. At low zoom, a fixed-screen-size iterations-per-second column can consume almost the entire bar; keep the fill full width, draw it first, then paint the rate as a tightly padded opaque overlay so progress can flow underneath without sacrificing legibility. - ComfyUI core dependency pins can become startup-critical before the remote worker reaches a payload. ComfyUI's INT8 attention import now calls `comfy_kitchen.int8_attention_is_available()`, so pairing current ComfyUI with `comfy-kitchen==0.2.26` makes every GPU-snapshot container fail during `@modal.enter(snap=True)` and crash-loop. Keep both exact support packages aligned with the copied ComfyUI checkout (`comfy-kitchen==0.2.31` and `comfy-aimdo==0.4.13` for the current revision), and let the changed runtime fingerprint force an image rebuild. +- Current ComfyUI built-in replacements register through `PromptServer.instance.node_replace_manager` during `nodes.init_extra_nodes(...)`. A remote headless PromptServer shim must instantiate the real `server.NodeReplaceManager`; supplying only routes, queue state, and websocket no-ops makes `nodes_replacements.py` fail to import and silently drops compatibility mappings for renamed nodes. diff --git a/tests/test_modal_executor_node.py b/tests/test_modal_executor_node.py index e58b9a0..ac60bfe 100644 --- a/tests/test_modal_executor_node.py +++ b/tests/test_modal_executor_node.py @@ -1844,8 +1844,22 @@ def test_modal_cloud_installs_headless_prompt_server_instance( monkeypatch: Any, ) -> None: """Remote custom-node init should get a minimal PromptServer.instance shim.""" + class FakeNodeReplaceManager: + """Record replacement registrations made through the headless server.""" + + def __init__(self) -> None: + """Initialize an empty replacement registry.""" + self.registrations: list[Any] = [] + + def register(self, replacement: Any) -> None: + """Record one replacement registered by a ComfyUI extension.""" + self.registrations.append(replacement) + fake_prompt_server_class = type("PromptServer", (), {}) - fake_server_module = types.SimpleNamespace(PromptServer=fake_prompt_server_class) + fake_server_module = types.SimpleNamespace( + NodeReplaceManager=FakeNodeReplaceManager, + PromptServer=fake_prompt_server_class, + ) monkeypatch.setitem(sys.modules, "server", fake_server_module) modal_cloud_module._ensure_headless_prompt_server_instance() @@ -1854,6 +1868,7 @@ def test_modal_cloud_installs_headless_prompt_server_instance( assert instance is not None assert hasattr(instance, "routes") assert hasattr(instance, "app") + assert isinstance(instance.node_replace_manager, FakeNodeReplaceManager) assert instance.supports == ["custom_nodes_from_web"] assert instance.client_id is None assert instance.last_node_id is None @@ -1866,6 +1881,9 @@ def test_modal_cloud_installs_headless_prompt_server_instance( instance.add_on_prompt_handler("handler") assert instance.on_prompt_handlers == ["handler"] + instance.node_replace_manager.register("replacement") + assert instance.node_replace_manager.registrations == ["replacement"] + instance.prompt_queue.put((1, "prompt-id", {}, {}, [])) assert instance.prompt_queue.get_current_queue() == ( [],