Skip to content

Commit 8bb858e

Browse files
authored
Improve performance with large number of queued prompts (Comfy-Org#8176)
* get_current_queue_volatile * restore get_current_queue method * remove extra import
1 parent 57893c8 commit 8bb858e

3 files changed

Lines changed: 12 additions & 5 deletions

File tree

execution.py

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -909,7 +909,6 @@ def __init__(self, server):
909909
self.currently_running = {}
910910
self.history = {}
911911
self.flags = {}
912-
server.prompt_queue = self
913912

914913
def put(self, item):
915914
with self.mutex:
@@ -954,13 +953,21 @@ def task_done(self, item_id, history_result,
954953
self.history[prompt[1]].update(history_result)
955954
self.server.queue_updated()
956955

956+
# Note: slow
957957
def get_current_queue(self):
958958
with self.mutex:
959959
out = []
960960
for x in self.currently_running.values():
961961
out += [x]
962962
return (out, copy.deepcopy(self.queue))
963963

964+
# read-safe as long as queue items are immutable
965+
def get_current_queue_volatile(self):
966+
with self.mutex:
967+
running = [x for x in self.currently_running.values()]
968+
queued = copy.copy(self.queue)
969+
return (running, queued)
970+
964971
def get_tasks_remaining(self):
965972
with self.mutex:
966973
return len(self.queue) + len(self.currently_running)

main.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -260,7 +260,6 @@ def start_comfyui(asyncio_loop=None):
260260
asyncio_loop = asyncio.new_event_loop()
261261
asyncio.set_event_loop(asyncio_loop)
262262
prompt_server = server.PromptServer(asyncio_loop)
263-
q = execution.PromptQueue(prompt_server)
264263

265264
hook_breaker_ac10a0.save_functions()
266265
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):
271270
prompt_server.add_routes()
272271
hijack_progress(prompt_server)
273272

274-
threading.Thread(target=prompt_worker, daemon=True, args=(q, prompt_server,)).start()
273+
threading.Thread(target=prompt_worker, daemon=True, args=(prompt_server.prompt_queue, prompt_server,)).start()
275274

276275
if args.quick_test_for_ci:
277276
exit(0)

server.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
import node_helpers
3030
from comfyui_version import __version__
3131
from app.frontend_management import FrontendManager
32+
3233
from app.user_manager import UserManager
3334
from app.model_manager import ModelFileManager
3435
from app.custom_node_manager import CustomNodeManager
@@ -159,7 +160,7 @@ def __init__(self, loop):
159160
self.custom_node_manager = CustomNodeManager()
160161
self.internal_routes = InternalRoutes(self)
161162
self.supports = ["custom_nodes_from_web"]
162-
self.prompt_queue = None
163+
self.prompt_queue = execution.PromptQueue(self)
163164
self.loop = loop
164165
self.messages = asyncio.Queue()
165166
self.client_session:Optional[aiohttp.ClientSession] = None
@@ -620,7 +621,7 @@ async def get_history_prompt_id(request):
620621
@routes.get("/queue")
621622
async def get_queue(request):
622623
queue_info = {}
623-
current_queue = self.prompt_queue.get_current_queue()
624+
current_queue = self.prompt_queue.get_current_queue_volatile()
624625
queue_info['queue_running'] = current_queue[0]
625626
queue_info['queue_pending'] = current_queue[1]
626627
return web.json_response(queue_info)

0 commit comments

Comments
 (0)