File size: 2,229 Bytes
63ca2a0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 | """
企业级 ComfyUI 异步 TaskWorker
Author: XiaoZhe (Commercial Contact: janejulius119@gmail.com / WeChat: julius119)
"""
import asyncio
import json
import logging
import uuid
import aiohttp
from typing import Dict, Any, List
logger = logging.getLogger("VerseFlow.ComfyUIWorker")
class ComfyUIExecutionError(Exception):
pass
class ComfyUITaskWorker:
def __init__(self, host: str = "127.0.0.1", port: int = 8188):
self.base_url = f"http://{host}:{port}"
self.ws_url = f"ws://{host}:{port}/ws"
self.client_id = str(uuid.uuid4())
async def execute_prompt(self, workflow_prompt: Dict[str, Any], task_id: str) -> List[Dict[str, Any]]:
async with aiohttp.ClientSession() as session:
async with session.ws_connect(f"{self.ws_url}?clientId={self.client_id}") as ws:
payload = {"prompt": workflow_prompt, "client_id": self.client_id}
async with session.post(f"{self.base_url}/prompt", json=payload) as resp:
if resp.status != 200:
raise ComfyUIExecutionError(f"提交失败: {await resp.text()}")
prompt_id = (await resp.json())["prompt_id"]
while True:
msg = await ws.receive(timeout=120.0)
if msg.type == aiohttp.WSMsgType.TEXT:
event = json.loads(msg.data)
if event.get("type") == "executing" and event.get("data", {}).get("prompt_id") == prompt_id:
if event.get("data", {}).get("node") is None:
break
async with session.get(f"{self.base_url}/history/{prompt_id}") as hist_resp:
history = (await hist_resp.json()).get(prompt_id, {})
return self._extract_outputs(history)
def _extract_outputs(self, history: Dict[str, Any]) -> List[Dict[str, Any]]:
outputs = []
for node_id, node_out in history.get("outputs", {}).items():
if "images" in node_out:
for img in node_out["images"]:
outputs.append({"node_id": node_id, "filename": img["filename"], "type": "image"})
return outputs |