feat(workflows): add add_node

src/comfy_sdk/workflows.py

@@ -100,6 +100,80 @@ def remove_node(self, node_id: str) -> None:
                 for key in to_delete:
                     del inputs[key]
 
+    def add_node(
+        self,
+        class_type: str,
+        *,
+        before: str | None = None,
+        after: str | None = None,
+        inputs: dict[str, Any] | None = None,
+    ) -> str:
+        """Insert a new node, redirecting downstream connections through it.
+
+        The new node is assigned an auto-incremented ID (one greater than the
+        highest existing node ID).  ``class_type`` and optional ``inputs`` are
+        stored on the node.  Any link in ``inputs`` (e.g.
+        ``{"model": ["2", 0]}``) that points to an existing node causes *all*
+        downstream consumers of that source output to be redirected to the new
+        node's corresponding output.
+
+        ``before`` / ``after`` are informational — they document which existing
+        node the new node is placed relative to, but do not affect the
+        redirection logic (which is driven entirely by the links in ``inputs``).
+
+        Args:
+            class_type: ComfyUI class type (e.g. ``"KSampler"``).
+            before: If set, the new node is inserted before this node.
+            after: If set, the new node is inserted after this node.
+            inputs: Input dict for the new node. Links in this dict drive
+                downstream redirection.
+
+        Returns:
+            The auto-generated node ID.
+
+        Raises:
+            ValueError: If both ``before`` and ``after`` are given.
+        """
+        if before and after:
+            raise ValueError("Specify either 'before' or 'after', not both")
+
+        # Auto-generate node_id: one greater than the highest existing ID
+        if self.json:
+            max_id = max(int(nid) for nid in self.json)
+            node_id = str(max_id + 1)
+        else:
+            node_id = "1"
+
+        node_entry: dict[str, Any] = {"class_type": class_type}
+        if inputs:
+            node_entry["inputs"] = inputs
+        self.json[node_id] = node_entry
+
+        new_inputs = inputs or {}
+
+        # Collect (upstream_node_id, output_index) pairs from links in inputs
+        upstream_outputs: dict[str, set[int]] = {}
+        for value in new_inputs.values():
+            if _is_link(value):
+                src_node = value[0]
+                src_output = int(value[1])
+                if src_node in self.json:
+                    upstream_outputs.setdefault(src_node, set()).add(src_output)
+
+        # Redirect downstream consumers of those upstream outputs
+        for src_node, output_indices in upstream_outputs.items():
+            for nid, node in self.json.items():
+                if nid == node_id:
+                    continue
+                node_inputs = node.get("inputs")
+                if not node_inputs:
+                    continue
+                for key, value in list(node_inputs.items()):
+                    if _is_link(value) and value[0] == src_node and int(value[1]) in output_indices:
+                        node_inputs[key] = [node_id, value[1]]
+
+        return node_id
+
     def __repr__(self) -> str:
         return f"Workflow(nodes={len(self.json)})"
 

tests/test_workflows.py

@@ -69,6 +69,136 @@ def test_plain_graph_passes_through_the_walk_unchanged():
     assert substitute_asset_handles(graph, {}) == graph
 
 
+def test_add_node_redirects_downstream_single_consumer():
+    graph = {
+        "1": {
+            "class_type": "UNETLoader",
+            "inputs": {"unet_name": "model.safetensors"},
+        },
+        "2": {
+            "class_type": "LoraLoaderModelOnly",
+            "inputs": {
+                "lora_name": "lora.safetensors",
+                "strength_model": 1,
+                "model": ["1", 0],
+            },
+        },
+        "4": {
+            "class_type": "KSampler",
+            "inputs": {
+                "seed": 0,
+                "model": ["2", 0],
+            },
+        },
+    }
+    wf = Workflow(graph)
+    new_id = wf.add_node(
+        "ModelAttentionBackend",
+        before="4",
+        inputs={
+            "attention": "pytorch attention",
+            "model": ["2", 0],
+        },
+    )
+    assert new_id == "5"
+    assert wf.json[new_id]["class_type"] == "ModelAttentionBackend"
+    assert wf.json[new_id]["inputs"]["model"] == ["2", 0]
+    assert wf.json["4"]["inputs"]["model"] == [new_id, 0]
+
+
+def test_add_node_redirects_multiple_downstream():
+    graph = {
+        "1": {
+            "class_type": "LoadImage",
+            "inputs": {"image": "example.png"},
+        },
+        "2": {
+            "class_type": "PreviewImage",
+            "inputs": {"images": ["1", 0]},
+        },
+        "3": {
+            "class_type": "PreviewImage",
+            "inputs": {"images": ["1", 0]},
+        },
+        "4": {
+            "class_type": "PreviewImage",
+            "inputs": {"images": ["1", 0]},
+        },
+    }
+    wf = Workflow(graph)
+    new_id = wf.add_node(
+        "ImageScaleToTotalPixels",
+        after="1",
+        inputs={
+            "upscale_method": "nearest-exact",
+            "megapixels": 1,
+            "image": ["1", 0],
+        },
+    )
+    assert new_id == "5"
+    assert wf.json[new_id]["inputs"]["image"] == ["1", 0]
+    assert wf.json["2"]["inputs"]["images"] == [new_id, 0]
+    assert wf.json["3"]["inputs"]["images"] == [new_id, 0]
+    assert wf.json["4"]["inputs"]["images"] == [new_id, 0]
+
+
+def test_add_node_no_redirect_when_upstream_not_in_graph():
+    graph = {
+        "1": {
+            "class_type": "CLIPTextEncode",
+            "inputs": {"text": "hello"},
+        },
+    }
+    wf = Workflow(graph)
+    new_id = wf.add_node(
+        "KSampler",
+        inputs={"model": ["999", 0]},
+    )
+    assert new_id == "2"
+    assert wf.json[new_id]["inputs"]["model"] == ["999", 0]
+
+
+def test_add_node_no_redirect_when_no_downstream_consumers():
+    graph = {
+        "1": {
+            "class_type": "CheckpointLoader",
+            "inputs": {"ckpt_name": "model.safetensors"},
+        },
+    }
+    wf = Workflow(graph)
+    new_id = wf.add_node(
+        "LoraLoader",
+        inputs={"model": ["1", 0], "clip": ["1", 1]},
+    )
+    assert new_id == "2"
+    assert wf.json[new_id]["inputs"]["model"] == ["1", 0]
+    assert wf.json[new_id]["inputs"]["clip"] == ["1", 1]
+
+
+def test_add_node_both_before_and_after_raises():
+    wf = Workflow({"1": {"class_type": "X", "inputs": {}}})
+    try:
+        wf.add_node("Y", before="1", after="1", inputs={})
+    except ValueError as e:
+        assert "not both" in str(e)
+    else:
+        assert False, "Expected ValueError"
+
+
+def test_add_node_no_inputs_no_redirect():
+    wf = Workflow({"1": {"class_type": "X", "inputs": {}}})
+    new_id = wf.add_node("Y")
+    assert new_id == "2"
+    assert wf.json[new_id]["class_type"] == "Y"
+    assert "inputs" not in wf.json[new_id]
+
+
+def test_add_node_auto_id_on_empty_graph():
+    wf = Workflow({})
+    new_id = wf.add_node("X")
+    assert new_id == "1"
+
+
 def test_remove_node_model_attention_backend():
     graph = {
         "1": {