소스 검색

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.
John Karabudak 9 달 전
부모
커밋
958fe9639a
2개의 변경된 파일4개의 추가작업 그리고 3개의 파일을 삭제
  1. 1 1
      backend/apps/images/main.py
  2. 3 2
      backend/apps/images/utils/comfyui.py

+ 1 - 1
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,

+ 3 - 2
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