"test/git@developer.sourcefind.cn:OpenDAS/nni.git" did not exist on "d03c411c8eb865a840242d59260a942d041e2fba"
Commit 958fe963 authored by John Karabudak's avatar John Karabudak
Browse files

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.
parent 3715994c
...@@ -514,7 +514,7 @@ async def image_generations( ...@@ -514,7 +514,7 @@ async def image_generations(
data = ImageGenerationPayload(**data) data = ImageGenerationPayload(**data)
res = comfyui_generate_image( res = await comfyui_generate_image(
app.state.config.MODEL, app.state.config.MODEL,
data, data,
user.id, user.id,
......
import asyncio
import websocket # NOTE: websocket-client (https://github.com/websocket-client/websocket-client) import websocket # NOTE: websocket-client (https://github.com/websocket-client/websocket-client)
import uuid import uuid
import json import json
...@@ -328,7 +329,7 @@ class ImageGenerationPayload(BaseModel): ...@@ -328,7 +329,7 @@ class ImageGenerationPayload(BaseModel):
flux_fp8_clip: Optional[bool] = None flux_fp8_clip: Optional[bool] = None
def comfyui_generate_image( async def comfyui_generate_image(
model: str, payload: ImageGenerationPayload, client_id, base_url model: str, payload: ImageGenerationPayload, client_id, base_url
): ):
ws_url = base_url.replace("http://", "ws://").replace("https://", "wss://") ws_url = base_url.replace("http://", "ws://").replace("https://", "wss://")
...@@ -397,7 +398,7 @@ def comfyui_generate_image( ...@@ -397,7 +398,7 @@ def comfyui_generate_image(
return None return None
try: 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: except Exception as e:
log.exception(f"Error while receiving images: {e}") log.exception(f"Error while receiving images: {e}")
images = None images = None
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment