From 958fe9639a08732c04b102ca494d11f522b33b34 Mon Sep 17 00:00:00 2001 From: John Karabudak Date: Wed, 7 Aug 2024 21:45:20 -0230 Subject: [PATCH] fix: ComfyUI generation no longer causes FastAPI to stall for all users as the get_images() function involves a `while True` loop while waiting for a response from ComfyUI and is not async, when image generation is running the entire UI becomes unresponsive for all users. furthermore, when image generation takes too long, the Docker health check starts failing. this is certainly a bad fix as it does not convert everything to async, but rather just puts the blocking loop in a separate thread. however, it works and it at least fixes the problem for now. --- backend/apps/images/main.py | 2 +- backend/apps/images/utils/comfyui.py | 5 +++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/backend/apps/images/main.py b/backend/apps/images/main.py index 4239f3f45..a418f2693 100644 --- a/backend/apps/images/main.py +++ b/backend/apps/images/main.py @@ -514,7 +514,7 @@ async def image_generations( data = ImageGenerationPayload(**data) - res = comfyui_generate_image( + res = await comfyui_generate_image( app.state.config.MODEL, data, user.id, diff --git a/backend/apps/images/utils/comfyui.py b/backend/apps/images/utils/comfyui.py index 6c37f0c49..f82076809 100644 --- a/backend/apps/images/utils/comfyui.py +++ b/backend/apps/images/utils/comfyui.py @@ -1,3 +1,4 @@ +import asyncio import websocket # NOTE: websocket-client (https://github.com/websocket-client/websocket-client) import uuid import json @@ -328,7 +329,7 @@ class ImageGenerationPayload(BaseModel): flux_fp8_clip: Optional[bool] = None -def comfyui_generate_image( +async def comfyui_generate_image( model: str, payload: ImageGenerationPayload, client_id, base_url ): ws_url = base_url.replace("http://", "ws://").replace("https://", "wss://") @@ -397,7 +398,7 @@ def comfyui_generate_image( return None try: - images = get_images(ws, comfyui_prompt, client_id, base_url) + images = await asyncio.to_thread(get_images, ws, comfyui_prompt, client_id, base_url) except Exception as e: log.exception(f"Error while receiving images: {e}") images = None