From 8bb858e4d39f7f6a6969c584aeeaa1d606a812d6 Mon Sep 17 00:00:00 2001
From: Michael Abrahams <miabrahams@users.noreply.github.com>
Date: Wed, 21 May 2025 05:14:17 -0400
Subject: [PATCH] Improve performance with large number of queued prompts
 (#8176)

* get_current_queue_volatile

* restore get_current_queue method

* remove extra import
---
 execution.py | 9 ++++++++-
 main.py      | 3 +--
 server.py    | 5 +++--
 3 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/execution.py b/execution.py
index e5d1c69d9..15ff7567c 100644
--- a/execution.py
+++ b/execution.py
@@ -909,7 +909,6 @@ class PromptQueue:
         self.currently_running = {}
         self.history = {}
         self.flags = {}
-        server.prompt_queue = self
 
     def put(self, item):
         with self.mutex:
@@ -954,6 +953,7 @@ class PromptQueue:
             self.history[prompt[1]].update(history_result)
             self.server.queue_updated()
 
+    # Note: slow
     def get_current_queue(self):
         with self.mutex:
             out = []
@@ -961,6 +961,13 @@ class PromptQueue:
                 out += [x]
             return (out, copy.deepcopy(self.queue))
 
+    # read-safe as long as queue items are immutable
+    def get_current_queue_volatile(self):
+        with self.mutex:
+            running = [x for x in self.currently_running.values()]
+            queued = copy.copy(self.queue)
+            return (running, queued)
+
     def get_tasks_remaining(self):
         with self.mutex:
             return len(self.queue) + len(self.currently_running)
diff --git a/main.py b/main.py
index 0fde6d221..fb1f8d20b 100644
--- a/main.py
+++ b/main.py
@@ -260,7 +260,6 @@ def start_comfyui(asyncio_loop=None):
         asyncio_loop = asyncio.new_event_loop()
         asyncio.set_event_loop(asyncio_loop)
     prompt_server = server.PromptServer(asyncio_loop)
-    q = execution.PromptQueue(prompt_server)
 
     hook_breaker_ac10a0.save_functions()
     nodes.init_extra_nodes(init_custom_nodes=not args.disable_all_custom_nodes, init_api_nodes=not args.disable_api_nodes)
@@ -271,7 +270,7 @@ def start_comfyui(asyncio_loop=None):
     prompt_server.add_routes()
     hijack_progress(prompt_server)
 
-    threading.Thread(target=prompt_worker, daemon=True, args=(q, prompt_server,)).start()
+    threading.Thread(target=prompt_worker, daemon=True, args=(prompt_server.prompt_queue, prompt_server,)).start()
 
     if args.quick_test_for_ci:
         exit(0)
diff --git a/server.py b/server.py
index 16cd88d91..1b0a73601 100644
--- a/server.py
+++ b/server.py
@@ -29,6 +29,7 @@ import comfy.model_management
 import node_helpers
 from comfyui_version import __version__
 from app.frontend_management import FrontendManager
+
 from app.user_manager import UserManager
 from app.model_manager import ModelFileManager
 from app.custom_node_manager import CustomNodeManager
@@ -159,7 +160,7 @@ class PromptServer():
         self.custom_node_manager = CustomNodeManager()
         self.internal_routes = InternalRoutes(self)
         self.supports = ["custom_nodes_from_web"]
-        self.prompt_queue = None
+        self.prompt_queue = execution.PromptQueue(self)
         self.loop = loop
         self.messages = asyncio.Queue()
         self.client_session:Optional[aiohttp.ClientSession] = None
@@ -620,7 +621,7 @@ class PromptServer():
         @routes.get("/queue")
         async def get_queue(request):
             queue_info = {}
-            current_queue = self.prompt_queue.get_current_queue()
+            current_queue = self.prompt_queue.get_current_queue_volatile()
             queue_info['queue_running'] = current_queue[0]
             queue_info['queue_pending'] = current_queue[1]
             return web.json_response(queue_info)