diff --git a/src/comfy_sdk/workflows.py b/src/comfy_sdk/workflows.py index d21bffe..31acbf4 100644 --- a/src/comfy_sdk/workflows.py +++ b/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)})" diff --git a/tests/test_workflows.py b/tests/test_workflows.py index 9258dea..581b60b 100644 --- a/tests/test_workflows.py +++ b/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]