Support ComfyUI node replacements remotely

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
 

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.")
 
 

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.

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() == (
         [],