From 5faf2e3cfd03ff1178eb6f9f0873355bd2122c30 Mon Sep 17 00:00:00 2001
From: Jedrzej Kosinski <kosinkadink1@gmail.com>
Date: Wed, 21 Jan 2026 16:36:02 -0800
Subject: [PATCH] Create helper classes for node replace registration

---
 comfy_api/latest/__init__.py             |  2 +
 comfy_api/latest/_node_replace.py        | 99 ++++++++++++++++++++++++
 comfy_api/latest/_node_replace_public.py |  1 +
 comfy_api/v0_0_2/__init__.py             |  3 +-
 4 files changed, 104 insertions(+), 1 deletion(-)
 create mode 100644 comfy_api/latest/_node_replace.py
 create mode 100644 comfy_api/latest/_node_replace_public.py

diff --git a/comfy_api/latest/__init__.py b/comfy_api/latest/__init__.py
index b0fa14ff6..24f6f438b 100644
--- a/comfy_api/latest/__init__.py
+++ b/comfy_api/latest/__init__.py
@@ -10,6 +10,7 @@ from ._input_impl import VideoFromFile, VideoFromComponents
 from ._util import VideoCodec, VideoContainer, VideoComponents, MESH, VOXEL
 from . import _io_public as io
 from . import _ui_public as ui
+from . import _node_replace_public as node_replace
 from comfy_execution.utils import get_executing_context
 from comfy_execution.progress import get_progress_state, PreviewImageTuple
 from PIL import Image
@@ -130,4 +131,5 @@ __all__ = [
     "IO",
     "ui",
     "UI",
+    "node_replace",
 ]
diff --git a/comfy_api/latest/_node_replace.py b/comfy_api/latest/_node_replace.py
new file mode 100644
index 000000000..4703937d9
--- /dev/null
+++ b/comfy_api/latest/_node_replace.py
@@ -0,0 +1,99 @@
+from __future__ import annotations
+
+from typing import Any
+
+
+class NodeReplace:
+    """
+    Defines a possible node replacement, mapping inputs and outputs of the old node to the new node.
+
+    Also supports assigning specific values to the input widgets of the new node.
+    """
+    def __init__(self,
+        new_node_id: str,
+        old_node_id: str,
+        input_mapping: list[InputMap],
+        output_mapping: list[OutputMap],
+    ):
+        self.new_node_id = new_node_id
+        self.old_node_id = old_node_id
+        self.input_mapping = input_mapping
+        self.output_mapping = output_mapping
+
+    def as_dict(self):
+        """
+        Create serializable representation of the node replacement.
+        """
+        return {
+            "new_node_id": self.new_node_id,
+            "old_node_id": self.old_node_id,
+            "input_mapping": [m.as_dict() for m in self.input_mapping],
+            "output_mapping": [m.as_dict() for m in self.output_mapping],
+        }
+
+
+class InputMap:
+    """
+    Map inputs of node replacement.
+
+    Use InputMap.OldId or InputMap.UseValue for mapping purposes.
+    """
+    class _Assign:
+        def __init__(self, assign_type: str):
+            self.assign_type = assign_type
+
+        def as_dict(self):
+            return {
+                "assign_type": self.assign_type,
+            }
+
+    class OldId(_Assign):
+        """
+        Connect the input of the old node with given id to new node when replacing.
+        """
+        def __init__(self, old_id: str):
+            super().__init__("old_id")
+            self.old_id = old_id
+
+        def as_dict(self):
+            return super().as_dict() | {
+                "old_id": self.old_id,
+            }
+
+    class UseValue(_Assign):
+        """
+        Use the given value for the input of the new node when replacing; assumes input is a widget.
+        """
+        def __init__(self, value: Any):
+            super().__init__("use_value")
+            self.value = value
+
+        def as_dict(self):
+            return super().as_dict() | {
+                "value": self.value,
+            }
+
+    def __init__(self, new_id: str, assign: OldId | UseValue):
+        self.new_id = new_id
+        self.assign = assign
+
+    def as_dict(self):
+        return {
+            "new_id": self.new_id,
+            "assign": self.assign.as_dict(),
+        }
+
+
+class OutputMap:
+    """
+    Map outputs of node replacement via indexes, as that's how outputs are stored.
+    """
+    def __init__(self, new_idx: int, old_idx: int):
+        self.new_idx = new_idx
+        self.old_idx = old_idx
+
+    def as_dict(self):
+        return {
+            "new_idx": self.new_idx,
+            "old_idx": self.old_idx,
+        }
diff --git a/comfy_api/latest/_node_replace_public.py b/comfy_api/latest/_node_replace_public.py
new file mode 100644
index 000000000..cd18168ae
--- /dev/null
+++ b/comfy_api/latest/_node_replace_public.py
@@ -0,0 +1 @@
+from ._node_replace import *  # noqa: F403
diff --git a/comfy_api/v0_0_2/__init__.py b/comfy_api/v0_0_2/__init__.py
index c4fa1d971..0d4d567da 100644
--- a/comfy_api/v0_0_2/__init__.py
+++ b/comfy_api/v0_0_2/__init__.py
@@ -6,7 +6,7 @@ from comfy_api.latest import (
 )
 from typing import Type, TYPE_CHECKING
 from comfy_api.internal.async_to_sync import create_sync_class
-from comfy_api.latest import io, ui, IO, UI, ComfyExtension  #noqa: F401
+from comfy_api.latest import io, ui, IO, UI, ComfyExtension, node_replace  #noqa: F401
 
 
 class ComfyAPIAdapter_v0_0_2(ComfyAPI_latest):
@@ -46,4 +46,5 @@ __all__ = [
     "IO",
     "ui",
     "UI",
+    "node_replace",
 ]