Call extension.on_load() in load_custom_node function, arttempt at using the public interface's node_replacement (not working currently, pushing to share)

comfy_api/latest/_node_replace.py

@@ -1,12 +1,6 @@
 from __future__ import annotations
 
 from typing import Any, TypedDict
-from server import PromptServer
-
-
-def register(node_replace: NodeReplace):
-    """Register a node replacement mapping."""
-    PromptServer.instance.node_replace_manager.register(node_replace)
 
 
 class InputMapOldId(TypedDict):

comfy_extras/nodes_replacements.py

@@ -1,20 +1,20 @@
-from comfy_api.latest import ComfyExtension, io, node_replace
+from comfy_api.latest import ComfyExtension, io, node_replace, ComfyAPI
 
 
 async def register_replacements():
     """Register all built-in node 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()
+    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_replace.NodeReplace(
+    await ComfyAPI.node_replacement.register(node_replace.NodeReplace(
             new_node_id="ImageScaleToMaxDimension",
             old_node_id="ResizeImagesByLongerEdge",
             old_widget_ids=["longer_edge"],
@@ -27,9 +27,9 @@ def register_replacements_longeredge():
             output_mapping=[{"new_idx": 0, "old_idx": 0}],
         ))
 
-def register_replacements_batchimages():
+async def register_replacements_batchimages():
     # BatchImages node uses Autogrow
-    node_replace.register(node_replace.NodeReplace(
+    await ComfyAPI.node_replacement.register(node_replace.NodeReplace(
             new_node_id="BatchImagesNode",
             old_node_id="ImageBatch",
             input_mapping=[
@@ -38,9 +38,9 @@ def register_replacements_batchimages():
             ],
         ))
 
-def register_replacements_upscaleimage():
+async def register_replacements_upscaleimage():
     # ResizeImageMaskNode uses DynamicCombo
-    node_replace.register(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"],
@@ -52,9 +52,9 @@ def register_replacements_upscaleimage():
             ],
         ))
 
-def register_replacements_controlnet():
+async def register_replacements_controlnet():
     # T2IAdapterLoader → ControlNetLoader
-    node_replace.register(node_replace.NodeReplace(
+    await ComfyAPI.node_replacement.register(node_replace.NodeReplace(
             new_node_id="ControlNetLoader",
             old_node_id="T2IAdapterLoader",
             input_mapping=[
@@ -62,38 +62,40 @@ def register_replacements_controlnet():
             ],
         ))
 
-def register_replacements_load3d():
+async def register_replacements_load3d():
     # Load3DAnimation merged into Load3D
-    node_replace.register(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_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_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_replace.NodeReplace(
+    await ComfyAPI.node_replacement.register(node_replace.NodeReplace(
             new_node_id="ConditioningAverage",
             old_node_id="ConditioningAverage ",
         ))
 
 class NodeReplacementsExtension(ComfyExtension):
+    async def on_load(self) -> None:
+        await register_replacements()
+
     async def get_node_list(self) -> list[type[io.ComfyNode]]:
         return []
 
 async def comfy_entrypoint() -> NodeReplacementsExtension:
-    await register_replacements()
     return NodeReplacementsExtension()

nodes.py

@@ -2276,6 +2276,7 @@ async def load_custom_node(module_path: str, ignore=set(), module_parent="custom
                         node_cls.RELATIVE_PYTHON_MODULE = "{}.{}".format(module_parent, get_module_name(module_path))
                     if schema.display_name is not None:
                         NODE_DISPLAY_NAME_MAPPINGS[schema.node_id] = schema.display_name
+                await extension.on_load()
                 return True
             except Exception as e:
                 logging.warning(f"Error while calling comfy_entrypoint in {module_path}: {e}")