From 9e758b5b0ce89847afcb52a25866b2f55c119fd8 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Wed, 11 Feb 2026 01:02:55 -0800 Subject: [PATCH] Refactored _node_replace.py InputMap/OutputMap to use a TypedDict instead of objects, simplified the schema sent to the frontend, updated nodes_post_processing.py replacements to use new schema --- comfy_api/latest/_node_replace.py | 93 +++++++++------------------ comfy_extras/nodes_post_processing.py | 22 +++---- 2 files changed, 40 insertions(+), 75 deletions(-) diff --git a/comfy_api/latest/_node_replace.py b/comfy_api/latest/_node_replace.py index 89c8c144f..723863c06 100644 --- a/comfy_api/latest/_node_replace.py +++ b/comfy_api/latest/_node_replace.py @@ -1,6 +1,32 @@ from __future__ import annotations -from typing import Any +from typing import Any, TypedDict + + +class InputMapOldId(TypedDict): + """Map an old node input to a new node input by ID.""" + new_id: str + old_id: str + + +class InputMapSetValue(TypedDict): + """Set a specific value for a new node input.""" + new_id: str + set_value: Any + + +InputMap = InputMapOldId | InputMapSetValue +""" +Input mapping for node replacement. Type is inferred by dictionary keys: +- {"new_id": str, "old_id": str} - maps old input to new input +- {"new_id": str, "set_value": Any} - sets a specific value for new input +""" + + +class OutputMap(TypedDict): + """Map outputs of node replacement via indexes.""" + new_idx: int + old_idx: int class NodeReplace: @@ -28,67 +54,6 @@ class NodeReplace: "new_node_id": self.new_node_id, "old_node_id": self.old_node_id, "old_widget_ids": self.old_widget_ids, - "input_mapping": [m.as_dict() for m in self.input_mapping] if self.input_mapping else None, - "output_mapping": [m.as_dict() for m in self.output_mapping] if self.output_mapping else None, - } - - -class InputMap: - """ - Map inputs of node replacement. - - Use InputMap.OldId or InputMap.SetValue 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 SetValue(_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__("set_value") - self.value = value - - def as_dict(self): - return super().as_dict() | { - "value": self.value, - } - - def __init__(self, new_id: str, assign: OldId | SetValue): - 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, + "input_mapping": list(self.input_mapping) if self.input_mapping else None, + "output_mapping": list(self.output_mapping) if self.output_mapping else None, } diff --git a/comfy_extras/nodes_post_processing.py b/comfy_extras/nodes_post_processing.py index 0aa80f45c..32b0a1da0 100644 --- a/comfy_extras/nodes_post_processing.py +++ b/comfy_extras/nodes_post_processing.py @@ -681,12 +681,12 @@ def register_replacements_longeredge(): old_node_id="ResizeImagesByLongerEdge", old_widget_ids=["longer_edge"], input_mapping=[ - node_replace.InputMap(new_id="image", assign=node_replace.InputMap.OldId("images")), - node_replace.InputMap(new_id="largest_size", assign=node_replace.InputMap.OldId("longer_edge")), - node_replace.InputMap(new_id="upscale_method", assign=node_replace.InputMap.SetValue("lanczos")), + {"new_id": "image", "old_id": "images"}, + {"new_id": "largest_size", "old_id": "longer_edge"}, + {"new_id": "upscale_method", "set_value": "lanczos"}, ], # just to test the frontend output_mapping code, does nothing really here - output_mapping=[node_replace.OutputMap(new_idx=0, old_idx=0)], + output_mapping=[{"new_idx": 0, "old_idx": 0}], )) def register_replacements_batchimages(): @@ -695,8 +695,8 @@ def register_replacements_batchimages(): new_node_id="BatchImagesNode", old_node_id="ImageBatch", input_mapping=[ - node_replace.InputMap(new_id="images.image0", assign=node_replace.InputMap.OldId("image1")), - node_replace.InputMap(new_id="images.image1", assign=node_replace.InputMap.OldId("image2")), + {"new_id": "images.image0", "old_id": "image1"}, + {"new_id": "images.image1", "old_id": "image2"}, ], )) @@ -707,10 +707,10 @@ def register_replacements_upscaleimage(): old_node_id="ImageScaleBy", old_widget_ids=["upscale_method", "scale_by"], input_mapping=[ - node_replace.InputMap(new_id="input", assign=node_replace.InputMap.OldId("image")), - node_replace.InputMap(new_id="resize_type", assign=node_replace.InputMap.SetValue("scale by multiplier")), - node_replace.InputMap(new_id="resize_type.multiplier", assign=node_replace.InputMap.OldId("scale_by")), - node_replace.InputMap(new_id="scale_method", assign=node_replace.InputMap.OldId("upscale_method")), + {"new_id": "input", "old_id": "image"}, + {"new_id": "resize_type", "set_value": "scale by multiplier"}, + {"new_id": "resize_type.multiplier", "old_id": "scale_by"}, + {"new_id": "scale_method", "old_id": "upscale_method"}, ], )) @@ -720,7 +720,7 @@ def register_replacements_controlnet(): new_node_id="ControlNetLoader", old_node_id="T2IAdapterLoader", input_mapping=[ - node_replace.InputMap(new_id="control_net_name", assign=node_replace.InputMap.OldId("t2i_adapter_name")), + {"new_id": "control_net_name", "old_id": "t2i_adapter_name"}, ], ))