From d52ed407a759688b50803e0ff930b206ae1721e6 Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Tue, 13 Jun 2023 13:36:47 -0400
Subject: [PATCH] Send websocket message only when prompt is actually done
 executing.

---
 execution.py |  6 ------
 main.py      | 13 +++++++++++--
 2 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/execution.py b/execution.py
index fc9578bc..f93de846 100644
--- a/execution.py
+++ b/execution.py
@@ -310,7 +310,6 @@ class PromptExecutor:
         else:
             self.server.client_id = None
 
-        execution_start_time = time.perf_counter()
         if self.server.client_id is not None:
             self.server.send_sync("execution_start", { "prompt_id": prompt_id}, self.server.client_id)
 
@@ -358,12 +357,7 @@ class PromptExecutor:
             for x in executed:
                 self.old_prompt[x] = copy.deepcopy(prompt[x])
             self.server.last_node_id = None
-            if self.server.client_id is not None:
-                self.server.send_sync("executing", { "node": None, "prompt_id": prompt_id }, self.server.client_id)
 
-        print("Prompt executed in {:.2f} seconds".format(time.perf_counter() - execution_start_time))
-        gc.collect()
-        comfy.model_management.soft_empty_cache()
 
 
 def validate_inputs(prompt, item, validated):
diff --git a/main.py b/main.py
index 8293c06f..22425d2a 100644
--- a/main.py
+++ b/main.py
@@ -3,6 +3,8 @@ import itertools
 import os
 import shutil
 import threading
+import gc
+import time
 
 from comfy.cli_args import args
 import comfy.utils
@@ -28,15 +30,22 @@ import folder_paths
 import server
 from server import BinaryEventTypes
 from nodes import init_custom_nodes
-
+import comfy.model_management
 
 def prompt_worker(q, server):
     e = execution.PromptExecutor(server)
     while True:
         item, item_id = q.get()
-        e.execute(item[2], item[1], item[3], item[4])
+        execution_start_time = time.perf_counter()
+        prompt_id = item[1]
+        e.execute(item[2], prompt_id, item[3], item[4])
         q.task_done(item_id, e.outputs_ui)
+        if server.client_id is not None:
+            server.send_sync("executing", { "node": None, "prompt_id": prompt_id }, server.client_id)
 
+        print("Prompt executed in {:.2f} seconds".format(time.perf_counter() - execution_start_time))
+        gc.collect()
+        comfy.model_management.soft_empty_cache()
 
 async def run(server, address='', port=8188, verbose=True, call_on_start=None):
     await asyncio.gather(server.start(address, port, verbose, call_on_start), server.publish_loop())