From 14af129c5509d10504113a1520c45b0ebcf81f14 Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Sat, 17 Aug 2024 11:36:10 -0400
Subject: [PATCH] Improve execution UX.

Some branches with VAELoader -> VAEDecode -> Preview were being executed
last. With this change they will be executed earlier.
---
 comfy_execution/graph.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/comfy_execution/graph.py b/comfy_execution/graph.py
index 2c7e8f8ae..b53e10f3f 100644
--- a/comfy_execution/graph.py
+++ b/comfy_execution/graph.py
@@ -202,8 +202,14 @@ class ExecutionList(TopologicalSort):
                 if is_output(blocked_node_id):
                     return node_id
 
-        #Do we want to look deeper?
+        #This should handle the VAELoader -> VAEDecode -> preview case
+        for node_id in node_list:
+            for blocked_node_id in self.blocking[node_id]:
+                for blocked_node_id1 in self.blocking[blocked_node_id]:
+                    if is_output(blocked_node_id1):
+                        return node_id
 
+        #TODO: this function should be improved
         return node_list[0]
 
     def unstage_node_execution(self):