| """ |
| 企业级 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 |
