|
@@ -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://")
|
|
@@ -377,9 +378,9 @@ def comfyui_generate_image(
|
|
comfyui_prompt["12"]["inputs"]["weight_dtype"] = payload.flux_weight_dtype
|
|
comfyui_prompt["12"]["inputs"]["weight_dtype"] = payload.flux_weight_dtype
|
|
|
|
|
|
if payload.flux_fp8_clip:
|
|
if payload.flux_fp8_clip:
|
|
- comfyui_prompt["11"]["inputs"][
|
|
|
|
- "clip_name2"
|
|
|
|
- ] = "t5xxl_fp8_e4m3fn.safetensors"
|
|
|
|
|
|
+ comfyui_prompt["11"]["inputs"]["clip_name2"] = (
|
|
|
|
+ "t5xxl_fp8_e4m3fn.safetensors"
|
|
|
|
+ )
|
|
|
|
|
|
comfyui_prompt["5"]["inputs"]["batch_size"] = payload.n
|
|
comfyui_prompt["5"]["inputs"]["batch_size"] = payload.n
|
|
comfyui_prompt["5"]["inputs"]["width"] = payload.width
|
|
comfyui_prompt["5"]["inputs"]["width"] = payload.width
|
|
@@ -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
|