diff --git a/src/comfy_sdk/workflows.py b/src/comfy_sdk/workflows.py index 31acbf4..c127b10 100644 --- a/src/comfy_sdk/workflows.py +++ b/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)})" diff --git a/tests/test_workflows.py b/tests/test_workflows.py index 581b60b..44eb28a 100644 --- a/tests/test_workflows.py +++ b/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": {