feat(workflows): add remove_node and _is_link helper

src/comfy_sdk/workflows.py

@@ -14,6 +14,20 @@
 from typing import Any
 
 
+def _is_link(obj: Any) -> bool:
+    """Return ``True`` if ``obj`` is a ComfyUI API-format connection link
+    (``[node_id: str, output_index: int]``)."""
+    if not isinstance(obj, list):
+        return False
+    if len(obj) != 2:
+        return False
+    if not isinstance(obj[0], str):
+        return False
+    if not isinstance(obj[1], int) and not isinstance(obj[1], float):
+        return False
+    return True
+
+
 class Workflow:
     """An API-format ComfyUI graph, ready to submit.
 
@@ -37,6 +51,55 @@ def set_input(self, node_id: str, field: str, value: Any) -> None:
         inputs = node.setdefault("inputs", {})
         inputs[field] = value
 
+    def remove_node(self, node_id: str) -> None:
+        """Remove a node and redirect links through it back to their sources.
+
+        Deletes the node identified by ``node_id`` from the graph. Any input
+        connections (links) in other nodes that reference this node's outputs
+        are redirected to the source that fed into the removed node, effectively
+        unwinding any insertion point.
+
+        If the removed node has exactly one input that is a link, all downstream
+        consumers of its outputs are redirected to that source. Otherwise (zero
+        or multiple link inputs), downstream links are simply deleted.
+        """
+        removed = self.json.pop(node_id, None)
+        if removed is None:
+            return
+
+        # Collect link inputs from the removed node
+        link_inputs: list[tuple[str, int]] = []
+        removed_inputs = removed.get("inputs") or {}
+        for value in removed_inputs.values():
+            if _is_link(value):
+                link_inputs.append((value[0], int(value[1])))
+
+        if len(link_inputs) == 1:
+            # Single link input: redirect all downstream consumers to that source
+            src_node, src_output = link_inputs[0]
+            for node in self.json.values():
+                inputs = node.get("inputs")
+                if not inputs:
+                    continue
+                for key, value in list(inputs.items()):
+                    if _is_link(value) and value[0] == node_id:
+                        if src_node in self.json:
+                            inputs[key] = [src_node, src_output]
+                        else:
+                            del inputs[key]
+        else:
+            # Zero or multiple link inputs: just delete downstream links
+            for node in self.json.values():
+                inputs = node.get("inputs")
+                if not inputs:
+                    continue
+                to_delete = []
+                for key, value in inputs.items():
+                    if _is_link(value) and value[0] == node_id:
+                        to_delete.append(key)
+                for key in to_delete:
+                    del inputs[key]
+
     def __repr__(self) -> str:
         return f"Workflow(nodes={len(self.json)})"
 

tests/test_workflows.py

@@ -67,3 +67,145 @@ def test_plain_graph_passes_through_the_walk_unchanged():
     graph = {"1": {"inputs": {"seed": 42, "model": "x.safetensors"}}}
     assert find_asset_handles(graph) == []
     assert substitute_asset_handles(graph, {}) == graph
+
+
+def test_remove_node_model_attention_backend():
+    graph = {
+        "1": {
+            "inputs": {
+                "unet_name": "z.safetensors",
+                "weight_dtype": "default",
+            },
+            "class_type": "UNETLoader",
+            "_meta": {"title": "Load Diffusion Model"},
+        },
+        "2": {
+            "inputs": {
+                "lora_name": "e.safetensors",
+                "strength_model": 1,
+                "model": ["1", 0],
+            },
+            "class_type": "LoraLoaderModelOnly",
+            "_meta": {"title": "Load LoRA"},
+        },
+        "3": {
+            "inputs": {
+                "attention": "pytorch attention",
+                "model": ["2", 0],
+            },
+            "class_type": "ModelAttentionBackend",
+            "_meta": {"title": "ModelAttentionBackend"},
+        },
+        "4": {
+            "inputs": {
+                "seed": 0,
+                "steps": 20,
+                "cfg": 8,
+                "sampler_name": "euler",
+                "scheduler": "simple",
+                "denoise": 1,
+                "model": ["3", 0],
+            },
+            "class_type": "KSampler",
+            "_meta": {"title": "KSampler"},
+        },
+    }
+    wf = Workflow(graph)
+    wf.remove_node("3")
+
+    assert "3" not in wf.json
+    assert wf.json["4"]["inputs"]["model"] == ["2", 0]
+
+
+def test_remove_node_redirects_image_scale():
+    graph = {
+        "1": {
+            "inputs": {"image": "example.png"},
+            "class_type": "LoadImage",
+            "_meta": {"title": "Load Image (A)"},
+        },
+        "2": {
+            "inputs": {"images": ["5", 0]},
+            "class_type": "PreviewImage",
+            "_meta": {"title": "Preview Image (B)"},
+        },
+        "3": {
+            "inputs": {"images": ["5", 0]},
+            "class_type": "PreviewImage",
+            "_meta": {"title": "Preview Image (C)"},
+        },
+        "4": {
+            "inputs": {"images": ["5", 0]},
+            "class_type": "PreviewImage",
+            "_meta": {"title": "Preview Image (D)"},
+        },
+        "5": {
+            "inputs": {
+                "upscale_method": "nearest-exact",
+                "megapixels": 1,
+                "resolution_steps": 1,
+                "image": ["1", 0],
+            },
+            "class_type": "ImageScaleToTotalPixels",
+            "_meta": {"title": "Scale Image to Total Pixels (E)"},
+        },
+    }
+    wf = Workflow(graph)
+    wf.remove_node("5")
+
+    assert "5" not in wf.json
+    assert wf.json["2"]["inputs"]["images"] == ["1", 0]
+    assert wf.json["3"]["inputs"]["images"] == ["1", 0]
+    assert wf.json["4"]["inputs"]["images"] == ["1", 0]
+
+
+def test_remove_node_redirects_preview_any():
+    graph = {
+        "1": {
+            "inputs": {
+                "prompt": "test",
+                "max_length": 512,
+                "sampling_mode": "on",
+                "sampling_mode.temperature": 0.7,
+                "sampling_mode.top_k": 64,
+                "sampling_mode.top_p": 0.95,
+                "sampling_mode.min_p": 0.05,
+                "sampling_mode.repetition_penalty": 1.05,
+                "sampling_mode.seed": 0,
+                "sampling_mode.presence_penalty": 0,
+                "thinking": False,
+                "use_default_template": True,
+                "clip": ["4", 0],
+            },
+            "class_type": "TextGenerate",
+            "_meta": {"title": "Generate Text"},
+        },
+        "2": {
+            "inputs": {"source": ["1", 0]},
+            "class_type": "PreviewAny",
+            "_meta": {"title": "Preview as Text"},
+        },
+        "3": {
+            "inputs": {
+                "filename_prefix": "ComfyUI",
+                "format": "txt",
+                "text": ["2", 0],
+            },
+            "class_type": "SaveText",
+            "_meta": {"title": "Save Text"},
+        },
+        "4": {
+            "inputs": {
+                "clip_name": "qwen3vl_4b_bf16.safetensors",
+                "type": "stable_diffusion",
+                "device": "default",
+            },
+            "class_type": "CLIPLoader",
+            "_meta": {"title": "Load CLIP"},
+        },
+    }
+    wf = Workflow(graph)
+    wf.remove_node("2")
+
+    assert "2" not in wf.json
+    assert wf.json["3"]["inputs"]["text"] == ["1", 0]