346 lines
13 KiB
Python
346 lines
13 KiB
Python
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import json
|
||
import logging
|
||
import mimetypes
|
||
import os
|
||
from collections.abc import Callable
|
||
from pathlib import Path
|
||
from typing import Any
|
||
from urllib.parse import quote
|
||
|
||
import httpx
|
||
|
||
try:
|
||
from websockets.asyncio.client import connect as ws_connect
|
||
except ImportError: # websockets < 13
|
||
from websockets import connect as ws_connect
|
||
|
||
from websockets.exceptions import ConnectionClosed
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
ProgressCallback = Callable[[int], None]
|
||
|
||
|
||
class ComfyError(RuntimeError):
|
||
user_message = "Ошибка при обращении к ComfyUI"
|
||
|
||
|
||
class ComfyConnectionError(ComfyError):
|
||
user_message = "Нет связи с ComfyUI"
|
||
|
||
|
||
class ComfyHTTPError(ComfyError):
|
||
user_message = "ComfyUI вернул ошибку"
|
||
|
||
|
||
class ComfyValidationError(ComfyError):
|
||
user_message = "Workflow отклонён ComfyUI"
|
||
|
||
|
||
class ComfyExecutionError(ComfyError):
|
||
user_message = "Генерация завершилась с ошибкой"
|
||
|
||
|
||
class ComfyTimeout(ComfyError):
|
||
user_message = "Генерация не завершилась за отведённое время"
|
||
|
||
|
||
def _format_api_error(payload: dict[str, Any]) -> str:
|
||
parts: list[str] = []
|
||
error = payload.get("error")
|
||
if isinstance(error, dict):
|
||
message = error.get("message") or error.get("type") or str(error)
|
||
parts.append(str(message))
|
||
elif error:
|
||
parts.append(str(error))
|
||
|
||
node_errors = payload.get("node_errors") or {}
|
||
if isinstance(node_errors, dict):
|
||
for node_id, info in node_errors.items():
|
||
if not isinstance(info, dict):
|
||
continue
|
||
for item in info.get("errors") or []:
|
||
if not isinstance(item, dict):
|
||
continue
|
||
message = item.get("message") or item.get("type") or "ошибка"
|
||
details = item.get("details")
|
||
parts.append(f"узел {node_id}: {message}" + (f" ({details})" if details else ""))
|
||
return "; ".join(parts) or "workflow отклонён"
|
||
|
||
|
||
def _execution_error_message(entry: dict[str, Any]) -> str:
|
||
status = entry.get("status") or {}
|
||
for message in status.get("messages") or []:
|
||
if not isinstance(message, (list, tuple)) or len(message) != 2:
|
||
continue
|
||
event, data = message
|
||
if event == "execution_error" and isinstance(data, dict):
|
||
text = data.get("exception_message") or data.get("exception_type")
|
||
if text:
|
||
return str(text)
|
||
return "ComfyUI сообщил об ошибке выполнения"
|
||
|
||
|
||
class ProgressSession:
|
||
def __init__(self, websocket: Any) -> None:
|
||
self._ws = websocket
|
||
self._closed = False
|
||
|
||
async def wait(
|
||
self,
|
||
prompt_id: str,
|
||
*,
|
||
timeout: float,
|
||
on_progress: ProgressCallback,
|
||
) -> bool:
|
||
"""True — выполнение завершилось; False — WS потерян, нужен polling."""
|
||
loop = asyncio.get_running_loop()
|
||
deadline = loop.time() + timeout
|
||
|
||
try:
|
||
while True:
|
||
remaining = deadline - loop.time()
|
||
if remaining <= 0:
|
||
raise ComfyTimeout("WebSocket: превышено время ожидания")
|
||
try:
|
||
raw = await asyncio.wait_for(self._ws.recv(), timeout=remaining)
|
||
except TimeoutError as exc:
|
||
raise ComfyTimeout("WebSocket: превышено время ожидания") from exc
|
||
except ConnectionClosed:
|
||
logger.warning("WebSocket ComfyUI закрыт, перехожу на polling")
|
||
return False
|
||
|
||
if isinstance(raw, bytes):
|
||
raw = raw.decode("utf-8", "replace")
|
||
try:
|
||
message = json.loads(raw)
|
||
except json.JSONDecodeError:
|
||
continue
|
||
if not isinstance(message, dict):
|
||
continue
|
||
|
||
message_type = message.get("type")
|
||
data = message.get("data") or {}
|
||
if not isinstance(data, dict):
|
||
data = {}
|
||
|
||
message_prompt = data.get("prompt_id")
|
||
if message_prompt and message_prompt != prompt_id:
|
||
continue
|
||
|
||
if message_type == "progress":
|
||
value = data.get("value", 0)
|
||
maximum = data.get("max", 0)
|
||
if isinstance(value, (int, float)) and isinstance(maximum, (int, float)) and maximum > 0:
|
||
on_progress(max(0, min(100, int(value * 100 / maximum))))
|
||
elif message_type == "execution_error":
|
||
raise ComfyExecutionError(_execution_error_message({"status": {"messages": [["execution_error", data]]}}))
|
||
elif message_type == "execution_interrupted":
|
||
raise ComfyExecutionError("Генерация прервана")
|
||
elif message_type == "execution_success" or (
|
||
message_type == "executing" and data.get("node") is None
|
||
):
|
||
return True
|
||
except ConnectionClosed:
|
||
logger.warning("WebSocket ComfyUI закрыт, перехожу на polling")
|
||
return False
|
||
|
||
async def close(self) -> None:
|
||
if self._closed:
|
||
return
|
||
self._closed = True
|
||
try:
|
||
await self._ws.close()
|
||
except Exception:
|
||
logger.debug("Ошибка закрытия WebSocket", exc_info=True)
|
||
|
||
|
||
class ComfyClient:
|
||
def __init__(self, base_url: str, ws_url: str) -> None:
|
||
self._base_url = base_url.rstrip("/")
|
||
self._ws_url = ws_url
|
||
self._client = httpx.AsyncClient(
|
||
base_url=self._base_url,
|
||
timeout=httpx.Timeout(30.0, connect=10.0),
|
||
trust_env=False,
|
||
)
|
||
|
||
async def close(self) -> None:
|
||
await self._client.aclose()
|
||
|
||
async def submit(self, workflow: dict[str, Any], client_id: str) -> str:
|
||
try:
|
||
response = await self._client.post(
|
||
"/prompt",
|
||
json={"prompt": workflow, "client_id": client_id},
|
||
)
|
||
except httpx.HTTPError as exc:
|
||
raise ComfyConnectionError(f"POST /prompt: {exc}") from exc
|
||
|
||
if response.status_code >= 400:
|
||
raise ComfyHTTPError(
|
||
f"POST /prompt: HTTP {response.status_code}: {response.text[:500]}"
|
||
)
|
||
try:
|
||
payload = response.json()
|
||
except ValueError as exc:
|
||
raise ComfyHTTPError("POST /prompt: ответ не является JSON") from exc
|
||
|
||
if payload.get("error") or payload.get("node_errors"):
|
||
raise ComfyValidationError(_format_api_error(payload))
|
||
|
||
prompt_id = payload.get("prompt_id")
|
||
if not prompt_id:
|
||
raise ComfyHTTPError("POST /prompt: в ответе нет prompt_id")
|
||
return str(prompt_id)
|
||
|
||
async def upload_image(self, path: Path, remote_name: str) -> str:
|
||
try:
|
||
content = path.read_bytes()
|
||
except OSError as exc:
|
||
raise ComfyError(f"Не удалось прочитать входной файл {path}: {exc}") from exc
|
||
|
||
mime = mimetypes.guess_type(path.name)[0] or "application/octet-stream"
|
||
try:
|
||
response = await self._client.post(
|
||
"/upload/image",
|
||
files={"image": (remote_name, content, mime)},
|
||
data={"type": "input", "overwrite": "true"},
|
||
timeout=60.0,
|
||
)
|
||
except httpx.HTTPError as exc:
|
||
raise ComfyConnectionError(f"POST /upload/image: {exc}") from exc
|
||
|
||
if response.status_code >= 400:
|
||
raise ComfyHTTPError(
|
||
f"POST /upload/image: HTTP {response.status_code}: {response.text[:500]}"
|
||
)
|
||
try:
|
||
payload = response.json()
|
||
except ValueError as exc:
|
||
raise ComfyHTTPError("POST /upload/image: ответ не является JSON") from exc
|
||
|
||
name = payload.get("name")
|
||
if not name:
|
||
raise ComfyHTTPError("POST /upload/image: в ответе нет имени файла")
|
||
subfolder = payload.get("subfolder") or ""
|
||
return f"{subfolder}/{name}" if subfolder else str(name)
|
||
|
||
async def get_history(self, prompt_id: str) -> dict[str, Any] | None:
|
||
try:
|
||
response = await self._client.get(f"/history/{quote(prompt_id)}")
|
||
except httpx.HTTPError as exc:
|
||
raise ComfyConnectionError(f"GET /history: {exc}") from exc
|
||
|
||
if response.status_code >= 400:
|
||
raise ComfyHTTPError(
|
||
f"GET /history: HTTP {response.status_code}: {response.text[:500]}"
|
||
)
|
||
try:
|
||
payload = response.json()
|
||
except ValueError as exc:
|
||
raise ComfyHTTPError("GET /history: ответ не является JSON") from exc
|
||
entry = payload.get(prompt_id)
|
||
return entry if isinstance(entry, dict) else None
|
||
|
||
async def wait_for_history(self, prompt_id: str, *, timeout: float) -> dict[str, Any]:
|
||
loop = asyncio.get_running_loop()
|
||
deadline = loop.time() + timeout
|
||
while loop.time() < deadline:
|
||
entry = await self.get_history(prompt_id)
|
||
if entry is not None:
|
||
status = entry.get("status") or {}
|
||
if status.get("status_str") == "error":
|
||
raise ComfyExecutionError(_execution_error_message(entry))
|
||
if status.get("completed") or status.get("status_str") == "success":
|
||
return entry
|
||
await asyncio.sleep(min(2.0, max(0.1, deadline - loop.time())))
|
||
raise ComfyTimeout("Не дождались результата в истории ComfyUI")
|
||
|
||
def find_output_image(self, entry: dict[str, Any]) -> dict[str, Any] | None:
|
||
outputs = entry.get("outputs") or {}
|
||
if not isinstance(outputs, dict):
|
||
return None
|
||
for node_output in outputs.values():
|
||
if not isinstance(node_output, dict):
|
||
continue
|
||
for image in node_output.get("images") or []:
|
||
if isinstance(image, dict) and image.get("filename"):
|
||
return image
|
||
return None
|
||
|
||
def find_output_text(self, entry: dict[str, Any]) -> str | None:
|
||
outputs = entry.get("outputs") or {}
|
||
if not isinstance(outputs, dict):
|
||
return None
|
||
for node_output in outputs.values():
|
||
if not isinstance(node_output, dict):
|
||
continue
|
||
value = node_output.get("text")
|
||
if isinstance(value, list):
|
||
value = value[0] if value else None
|
||
if isinstance(value, str) and value.strip():
|
||
return value.strip()
|
||
return None
|
||
|
||
async def download_output(self, entry: dict[str, Any], destination: Path) -> Path:
|
||
image = self.find_output_image(entry)
|
||
if image is None:
|
||
raise ComfyError("ComfyUI не вернул изображение")
|
||
return await self.download_image(
|
||
filename=str(image["filename"]),
|
||
subfolder=str(image.get("subfolder") or ""),
|
||
image_type=str(image.get("type") or "output"),
|
||
destination=destination,
|
||
)
|
||
|
||
async def download_image(
|
||
self,
|
||
*,
|
||
filename: str,
|
||
subfolder: str,
|
||
image_type: str,
|
||
destination: Path,
|
||
) -> Path:
|
||
destination.parent.mkdir(parents=True, exist_ok=True)
|
||
temp_path = destination.with_name(destination.name + ".part")
|
||
params = {"filename": filename, "subfolder": subfolder, "type": image_type}
|
||
try:
|
||
async with self._client.stream(
|
||
"GET", "/view", params=params, timeout=httpx.Timeout(120.0, connect=10.0)
|
||
) as response:
|
||
if response.status_code >= 400:
|
||
body = (await response.aread())[:500]
|
||
raise ComfyHTTPError(
|
||
f"GET /view: HTTP {response.status_code}: {body!r}"
|
||
)
|
||
with temp_path.open("wb") as handle:
|
||
async for chunk in response.aiter_bytes():
|
||
handle.write(chunk)
|
||
os.replace(temp_path, destination)
|
||
except httpx.HTTPError as exc:
|
||
temp_path.unlink(missing_ok=True)
|
||
raise ComfyConnectionError(f"GET /view: {exc}") from exc
|
||
except Exception:
|
||
temp_path.unlink(missing_ok=True)
|
||
raise
|
||
return destination
|
||
|
||
async def open_progress(self, client_id: str) -> ProgressSession | None:
|
||
url = f"{self._ws_url}?clientId={quote(client_id)}"
|
||
try:
|
||
websocket = await ws_connect(
|
||
url,
|
||
open_timeout=5,
|
||
close_timeout=2,
|
||
max_size=8 * 1024 * 1024,
|
||
ping_interval=20,
|
||
)
|
||
except Exception as exc: # noqa: BLE001 - WS не должен ронять генерацию
|
||
logger.warning("Не удалось подключиться к WebSocket ComfyUI: %s", exc)
|
||
return None
|
||
return ProgressSession(websocket)
|