Merge pull request #4448 from JohnTheNerd/comfyui-async-fix

fix: ComfyUI generation no longer causes FastAPI to stall for all users
This commit is contained in:
Timothy Jaeryang Baek 2024-08-08 11:36:19 +02:00 committed by GitHub
commit a7063a598d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 4 additions and 3 deletions

View File

@ -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,

View File

@ -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