From 5dc13651b0a867e80d5cbf3deb7acdc84c5debe6 Mon Sep 17 00:00:00 2001
From: Jacob Segal <jacob.e.segal@gmail.com>
Date: Sat, 20 Apr 2024 18:12:42 -0700
Subject: [PATCH] Use custom exception types.

---
 comfy/graph.py | 12 +++++++++---
 execution.py   |  5 ++++-
 2 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/comfy/graph.py b/comfy/graph.py
index 97a759de8..bf4be4ae1 100644
--- a/comfy/graph.py
+++ b/comfy/graph.py
@@ -2,6 +2,12 @@ import nodes
 
 from comfy.graph_utils import is_link
 
+class DependencyCycleError(Exception):
+    pass
+
+class NodeInputError(Exception):
+    pass
+
 class DynamicPrompt:
     def __init__(self, original_prompt):
         # The original prompt provided by the user
@@ -76,10 +82,10 @@ class TopologicalSort:
     def make_input_strong_link(self, to_node_id, to_input):
         inputs = self.dynprompt.get_node(to_node_id)["inputs"]
         if to_input not in inputs:
-            raise Exception(f"Node {to_node_id} says it needs input {to_input}, but there is no input to that node at all")
+            raise NodeInputError(f"Node {to_node_id} says it needs input {to_input}, but there is no input to that node at all")
         value = inputs[to_input]
         if not is_link(value):
-            raise Exception(f"Node {to_node_id} says it needs input {to_input}, but that value is a constant")
+            raise NodeInputError(f"Node {to_node_id} says it needs input {to_input}, but that value is a constant")
         from_node_id, from_socket = value
         self.add_strong_link(from_node_id, from_socket, to_node_id)
 
@@ -141,7 +147,7 @@ class ExecutionList(TopologicalSort):
             return None
         available = self.get_ready_nodes()
         if len(available) == 0:
-            raise Exception("Dependency cycle detected")
+            raise DependencyCycleError("Dependency cycle detected")
         next_node = available[0]
         # If an output node is available, do that first.
         # Technically this has no effect on the overall length of execution, but it feels better as a user
diff --git a/execution.py b/execution.py
index edeb71059..ecd0850af 100644
--- a/execution.py
+++ b/execution.py
@@ -23,6 +23,9 @@ class ExecutionResult(Enum):
     FAILURE = 1
     SLEEPING = 2
 
+class DuplicateNodeError(Exception):
+    pass
+
 class IsChangedCache:
     def __init__(self, dynprompt, outputs_cache):
         self.dynprompt = dynprompt
@@ -337,7 +340,7 @@ def execute(server, dynprompt, caches, current_item, extra_data, executed, promp
                     # Check for conflicts
                     for node_id in new_graph.keys():
                         if dynprompt.get_node(node_id) is not None:
-                            raise Exception(f"Attempt to add duplicate node {node_id}. Ensure node ids are unique and deterministic or use graph_utils.GraphBuilder.")
+                            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)
                         display_id = node_info.get("override_display_id", unique_id)