UniqAI/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui synced 2025-02-22 13:18:25 +00:00
Code Issues Actions 4 Packages Projects Releases Wiki Activity

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

Browse Source 
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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
backend/apps/images/main.py
View File

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

5
backend/apps/images/utils/comfyui.py
View File

@ -1,3 +1,4 @@
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):
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(
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