xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-09-09 04:57:02 +08:00
Code Issues Packages Projects Releases Wiki Activity

add async worker pool

Browse Source
This commit is contained in:
Aleksey Penkov 2025-04-02 18:00:42 +03:00
parent ab5413351e
commit 9de3aac2f7
2 changed files with 22 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
comfy/cli_args.py
View File

@ -37,6 +37,7 @@ parser = argparse.ArgumentParser()
parser.add_argument("--listen", type=str, default="127.0.0.1", metavar="IP", nargs="?", const="0.0.0.0,::", help="Specify the IP address to listen on (default: 127.0.0.1). You can give a list of ip addresses by separating them with a comma like: 127.2.2.2,127.3.3.3 If --listen is provided without an argument, it defaults to 0.0.0.0,:: (listens on all ipv4 and ipv6)")
parser.add_argument("--port", type=int, default=8188, help="Set the listen port.")
parser.add_argument("--workers", type=int, default=1, help="Number of concurrent workers to process tasks (default: 1)")
parser.add_argument("--tls-keyfile", type=str, help="Path to TLS (SSL) key file. Enables TLS, makes app accessible at https://... requires --tls-certfile to function")
parser.add_argument("--tls-certfile", type=str, help="Path to TLS (SSL) certificate file. Enables TLS, makes app accessible at https://... requires --tls-keyfile to function")
parser.add_argument("--enable-cors-header", type=str, default=None, metavar="ORIGIN", nargs="?", const="*", help="Enable CORS (Cross-Origin Resource Sharing) with optional origin or allow all with default '*'.")

25
main.py
View File

@ -154,19 +154,25 @@ def cuda_malloc_warning():
logging.warning("\nWARNING: this card most likely does not support cuda-malloc, if you get \"CUDA error\" please run ComfyUI with: --disable-cuda-malloc\n")
def prompt_worker(q, server_instance):
async def async_prompt_worker(q, server_instance):
"""
Asynchronous worker that repeatedly offloads the blocking q.get() call
to a background executor, processes a prompt, and then performs necessary cleanup.
"""
current_time: float = 0.0
e = execution.PromptExecutor(server_instance, lru_size=args.cache_lru)
last_gc_collect = 0
need_gc = False
gc_collect_interval = 10.0
loop = asyncio.get_event_loop()
while True:
timeout = 1000.0
if need_gc:
timeout = max(gc_collect_interval - (current_time - last_gc_collect), 0.0)
queue_item = q.get(timeout=timeout)
queue_item = await loop.run_in_executor(None, q.get, timeout)
if queue_item is not None:
item, item_id = queue_item
execution_start_time = time.perf_counter()
@ -209,6 +215,18 @@ def prompt_worker(q, server_instance):
last_gc_collect = current_time
need_gc = False
def start_worker_pool(prompt_queue, server_instance, num_workers):
"""
Starts a pool of asynchronous worker tasks inside a dedicated event loop
running in a separate thread.
"""
def thread_main():
asyncio_loop = asyncio.new_event_loop()
asyncio.set_event_loop(asyncio_loop)
tasks = [asyncio_loop.create_task(async_prompt_worker(prompt_queue, server_instance))
for _ in range(num_workers)]
asyncio_loop.run_until_complete(asyncio.gather(*tasks))
threading.Thread(target=thread_main, daemon=True).start()
async def run(server_instance, address='', port=8188, verbose=True, call_on_start=None):
addresses = []
@ -236,7 +254,6 @@ def cleanup_temp():
if os.path.exists(temp_dir):
shutil.rmtree(temp_dir, ignore_errors=True)
def start_comfyui(asyncio_loop=None):
"""
Starts the ComfyUI server using the provided asyncio event loop or creates a new one.
@ -268,7 +285,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()
start_worker_pool(q, prompt_server, num_workers=args.workers)
if args.quick_test_for_ci:
exit(0)