From 06f3ce9200734b54f123ff81e5a28757cc9ce5cb Mon Sep 17 00:00:00 2001 From: Jacob Segal Date: Sun, 21 Apr 2024 16:10:01 -0700 Subject: [PATCH] Raise exception for bad get_node calls. --- comfy/graph.py | 8 +++++++- execution.py | 2 +- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/comfy/graph.py b/comfy/graph.py index 9149ca4a2..6671bc3a0 100644 --- a/comfy/graph.py +++ b/comfy/graph.py @@ -8,6 +8,9 @@ class DependencyCycleError(Exception): class NodeInputError(Exception): pass +class NodeNotFoundError(Exception): + pass + class DynamicPrompt: def __init__(self, original_prompt): # The original prompt provided by the user @@ -22,7 +25,10 @@ class DynamicPrompt: return self.ephemeral_prompt[node_id] if node_id in self.original_prompt: return self.original_prompt[node_id] - return None + raise NodeNotFoundError(f"Node {node_id} not found") + + def has_node(self, node_id): + return node_id in self.original_prompt or node_id in self.ephemeral_prompt def add_ephemeral_node(self, node_id, node_info, parent_id, display_id): self.ephemeral_prompt[node_id] = node_info diff --git a/execution.py b/execution.py index 9db69cbd2..4688be2ad 100644 --- a/execution.py +++ b/execution.py @@ -349,7 +349,7 @@ def execute(server, dynprompt, caches, current_item, extra_data, executed, promp else: # Check for conflicts for node_id in new_graph.keys(): - if dynprompt.get_node(node_id) is not None: + if dynprompt.has_node(node_id): raise DuplicateNodeError(f"Attempt to add duplicate node {node_id}. Ensure node ids are unique and deterministic or use graph_utils.GraphBuilder.") for node_id, node_info in new_graph.items(): new_node_ids.append(node_id)