From 6ca331d890f3126426cb102962f8f70aa3a17c43 Mon Sep 17 00:00:00 2001 From: jab416171 Date: Fri, 21 Aug 2026 17:38:23 -0600 Subject: [PATCH 1/3] feat(workflows): add add_node with auto-incrementing ID Add add_node() method that auto-generates a node ID (one greater than the highest existing ID). Links in the new node's inputs cause all downstream consumers of those source outputs to be redirected to the new node. --- src/comfy_sdk/workflows.py | 74 +++++++++++++++++++++ tests/test_workflows.py | 130 +++++++++++++++++++++++++++++++++++++ 2 files changed, 204 insertions(+) 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..59d7583 100644 --- a/tests/test_workflows.py +++ b/tests/test_workflows.py @@ -209,3 +209,133 @@ def test_remove_node_redirects_preview_any(): assert "2" not in wf.json assert wf.json["3"]["inputs"]["text"] == ["1", 0] + + +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" From 5449e5be456550167e427e6fdb7fa5850b1e5877 Mon Sep 17 00:00:00 2001 From: jab416171 Date: Thu, 27 Aug 2026 23:33:49 -0600 Subject: [PATCH 2/3] fix ruff format --- src/comfy_sdk/workflows.py | 98 ++++++++++++++ tests/test_workflows.py | 260 +++++++++++++++++++++++++++++++++++++ 2 files changed, 358 insertions(+) diff --git a/src/comfy_sdk/workflows.py b/src/comfy_sdk/workflows.py index c127b10..f5c985e 100644 --- a/src/comfy_sdk/workflows.py +++ b/src/comfy_sdk/workflows.py @@ -174,6 +174,104 @@ def add_node( return node_id + 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 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 59d7583..7acc8e3 100644 --- a/tests/test_workflows.py +++ b/tests/test_workflows.py @@ -69,6 +69,266 @@ 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_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": { From ea6779c4cbb28822f127fa5d5c4b98b37389cc40 Mon Sep 17 00:00:00 2001 From: jab416171 Date: Fri, 28 Aug 2026 10:38:27 -0600 Subject: [PATCH 3/3] clean up duplicate code from merge conflict --- src/comfy_sdk/workflows.py | 98 -------------- tests/test_workflows.py | 260 ------------------------------------- 2 files changed, 358 deletions(-) diff --git a/src/comfy_sdk/workflows.py b/src/comfy_sdk/workflows.py index f5c985e..c127b10 100644 --- a/src/comfy_sdk/workflows.py +++ b/src/comfy_sdk/workflows.py @@ -174,104 +174,6 @@ def add_node( return node_id - 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 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 7acc8e3..44eb28a 100644 --- a/tests/test_workflows.py +++ b/tests/test_workflows.py @@ -199,136 +199,6 @@ def test_add_node_auto_id_on_empty_graph(): assert new_id == "1" -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": { @@ -469,133 +339,3 @@ def test_remove_node_redirects_preview_any(): assert "2" not in wf.json assert wf.json["3"]["inputs"]["text"] == ["1", 0] - - -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"