From 3f186525887ebdc9b4369a78679e64bf75a93535 Mon Sep 17 00:00:00 2001 From: bymyself Date: Wed, 4 Feb 2026 19:44:02 -0800 Subject: [PATCH] refactor: process isolation support for node replacement API - Move REGISTERED_NODE_REPLACEMENTS global to NodeReplaceManager instance state - Add NodeReplacement class to ComfyAPI_latest with async register() method - Deprecate module-level register_node_replacement() function - Call register_replacements() from comfy_entrypoint() This enables pyisolate compatibility where extensions run in separate processes and communicate via RPC. The async API allows registration calls to cross process boundaries. Refs: TDD-002 Amp-Thread-ID: https://ampcode.com/threads/T-019c2b33-ac55-76a9-9c6b-0246a8625f21 --- app/node_replace_manager.py | 41 +++++++++++++++----- comfy_api/latest/__init__.py | 17 +++++++++ comfy_api/latest/_node_replace.py | 16 +++++++- comfy_extras/nodes_post_processing.py | 54 ++++++++++++++------------- 4 files changed, 90 insertions(+), 38 deletions(-) diff --git a/app/node_replace_manager.py b/app/node_replace_manager.py index 3b1b7ab36..b513d7bc2 100644 --- a/app/node_replace_manager.py +++ b/app/node_replace_manager.py @@ -6,18 +6,39 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: from comfy_api.latest._node_replace import NodeReplace -REGISTERED_NODE_REPLACEMENTS: dict[str, list[NodeReplace]] = {} - -def register_node_replacement(node_replace: NodeReplace): - REGISTERED_NODE_REPLACEMENTS.setdefault(node_replace.old_node_id, []).append(node_replace) - -def registered_as_dict(): - return { - k: [v.as_dict() for v in v_list] for k, v_list in REGISTERED_NODE_REPLACEMENTS.items() - } class NodeReplaceManager: + """ + Manages node replacement registrations. + + Stores replacements as instance state (not module-level globals) to support + process isolation via pyisolate, where extensions run in separate processes + and communicate via RPC. + """ + + def __init__(self): + self._replacements: dict[str, list[NodeReplace]] = {} + + def register(self, node_replace: NodeReplace): + """Register a node replacement mapping.""" + self._replacements.setdefault(node_replace.old_node_id, []).append(node_replace) + + def get_replacement(self, old_node_id: str) -> list[NodeReplace] | None: + """Get replacements for an old node ID.""" + return self._replacements.get(old_node_id) + + def has_replacement(self, old_node_id: str) -> bool: + """Check if a replacement exists for an old node ID.""" + return old_node_id in self._replacements + + def as_dict(self): + """Serialize all replacements to dict.""" + return { + k: [v.as_dict() for v in v_list] + for k, v_list in self._replacements.items() + } + def add_routes(self, routes): @routes.get("/node_replacements") async def get_node_replacements(request): - return web.json_response(registered_as_dict()) + return web.json_response(self.as_dict()) diff --git a/comfy_api/latest/__init__.py b/comfy_api/latest/__init__.py index f0856e126..d5b4875e3 100644 --- a/comfy_api/latest/__init__.py +++ b/comfy_api/latest/__init__.py @@ -22,6 +22,23 @@ class ComfyAPI_latest(ComfyAPIBase): VERSION = "latest" STABLE = False + class NodeReplacement(ProxiedSingleton): + async def register(self, node_replace: 'node_replace.NodeReplace') -> None: + """ + Register a node replacement mapping. + + This async method supports process isolation via pyisolate, where + extensions run in separate processes. The call is RPC'd to the host + process where PromptServer and NodeReplaceManager live. + + Args: + node_replace: A NodeReplace object defining the old->new mapping + """ + from server import PromptServer + PromptServer.instance.node_replace_manager.register(node_replace) + + node_replacement: NodeReplacement + class Execution(ProxiedSingleton): async def set_progress( self, diff --git a/comfy_api/latest/_node_replace.py b/comfy_api/latest/_node_replace.py index 8e90eebb7..b55cd8a69 100644 --- a/comfy_api/latest/_node_replace.py +++ b/comfy_api/latest/_node_replace.py @@ -1,13 +1,25 @@ from __future__ import annotations +import warnings from typing import Any -import app.node_replace_manager + def register_node_replacement(node_replace: NodeReplace): """ Register node replacement. + + .. deprecated:: + Use ``ComfyAPI.node_replacement.register()`` instead. + This synchronous function does not work with process isolation (pyisolate). """ - app.node_replace_manager.register_node_replacement(node_replace) + warnings.warn( + "register_node_replacement() is deprecated. " + "Use 'await ComfyAPI.node_replacement.register()' instead for pyisolate compatibility.", + DeprecationWarning, + stacklevel=2 + ) + from server import PromptServer + PromptServer.instance.node_replace_manager.register(node_replace) class NodeReplace: diff --git a/comfy_extras/nodes_post_processing.py b/comfy_extras/nodes_post_processing.py index 5c9f8e6df..95f9c4c68 100644 --- a/comfy_extras/nodes_post_processing.py +++ b/comfy_extras/nodes_post_processing.py @@ -656,21 +656,22 @@ class BatchImagesMasksLatentsNode(io.ComfyNode): return io.NodeOutput(batched) -from comfy_api.latest import node_replace +from comfy_api.latest import ComfyAPI, node_replace -def register_replacements(): - register_replacements_longeredge() - register_replacements_batchimages() - register_replacements_upscaleimage() - register_replacements_controlnet() - register_replacements_load3d() - register_replacements_preview3d() - register_replacements_svdimg2vid() - register_replacements_conditioningavg() +async def register_replacements(): + """Register all built-in node replacements using the async API.""" + await register_replacements_longeredge() + await register_replacements_batchimages() + await register_replacements_upscaleimage() + await register_replacements_controlnet() + await register_replacements_load3d() + await register_replacements_preview3d() + await register_replacements_svdimg2vid() + await register_replacements_conditioningavg() -def register_replacements_longeredge(): +async def register_replacements_longeredge(): # No dynamic inputs here - node_replace.register_node_replacement(node_replace.NodeReplace( + await ComfyAPI.node_replacement.register(node_replace.NodeReplace( new_node_id="ImageScaleToMaxDimension", old_node_id="ResizeImagesByLongerEdge", old_widget_ids=["longer_edge"], @@ -683,9 +684,9 @@ def register_replacements_longeredge(): output_mapping=[node_replace.OutputMap(new_idx=0, old_idx=0)], )) -def register_replacements_batchimages(): +async def register_replacements_batchimages(): # BatchImages node uses Autogrow - node_replace.register_node_replacement(node_replace.NodeReplace( + await ComfyAPI.node_replacement.register(node_replace.NodeReplace( new_node_id="BatchImagesNode", old_node_id="ImageBatch", input_mapping=[ @@ -694,9 +695,9 @@ def register_replacements_batchimages(): ], )) -def register_replacements_upscaleimage(): +async def register_replacements_upscaleimage(): # ResizeImageMaskNode uses DynamicCombo - node_replace.register_node_replacement(node_replace.NodeReplace( + await ComfyAPI.node_replacement.register(node_replace.NodeReplace( new_node_id="ResizeImageMaskNode", old_node_id="ImageScaleBy", old_widget_ids=["upscale_method", "scale_by"], @@ -708,9 +709,9 @@ def register_replacements_upscaleimage(): ], )) -def register_replacements_controlnet(): +async def register_replacements_controlnet(): # T2IAdapterLoader → ControlNetLoader - node_replace.register_node_replacement(node_replace.NodeReplace( + await ComfyAPI.node_replacement.register(node_replace.NodeReplace( new_node_id="ControlNetLoader", old_node_id="T2IAdapterLoader", input_mapping=[ @@ -718,30 +719,30 @@ def register_replacements_controlnet(): ], )) -def register_replacements_load3d(): +async def register_replacements_load3d(): # Load3DAnimation merged into Load3D - node_replace.register_node_replacement(node_replace.NodeReplace( + await ComfyAPI.node_replacement.register(node_replace.NodeReplace( new_node_id="Load3D", old_node_id="Load3DAnimation", )) -def register_replacements_preview3d(): +async def register_replacements_preview3d(): # Preview3DAnimation merged into Preview3D - node_replace.register_node_replacement(node_replace.NodeReplace( + await ComfyAPI.node_replacement.register(node_replace.NodeReplace( new_node_id="Preview3D", old_node_id="Preview3DAnimation", )) -def register_replacements_svdimg2vid(): +async def register_replacements_svdimg2vid(): # Typo fix: SDV → SVD - node_replace.register_node_replacement(node_replace.NodeReplace( + await ComfyAPI.node_replacement.register(node_replace.NodeReplace( new_node_id="SVD_img2vid_Conditioning", old_node_id="SDV_img2vid_Conditioning", )) -def register_replacements_conditioningavg(): +async def register_replacements_conditioningavg(): # Typo fix: trailing space in node name - node_replace.register_node_replacement(node_replace.NodeReplace( + await ComfyAPI.node_replacement.register(node_replace.NodeReplace( new_node_id="ConditioningAverage", old_node_id="ConditioningAverage ", )) @@ -763,4 +764,5 @@ class PostProcessingExtension(ComfyExtension): ] async def comfy_entrypoint() -> PostProcessingExtension: + await register_replacements() return PostProcessingExtension()