Shaman
125 lines
5.5 KiB
Python
125 lines
5.5 KiB
Python
"""Async HTTP client for the ComfyUI server (official endpoints only)."""
|
|
import asyncio
|
|
import json
|
|
import secrets
|
|
import time
|
|
import uuid
|
|
from typing import Any
|
|
|
|
import httpx
|
|
|
|
from .errors import ComfyUIError
|
|
|
|
|
|
class ComfyClient:
|
|
def __init__(self, base_url: str, timeout: float = 30.0):
|
|
self._client = httpx.AsyncClient(base_url=base_url, timeout=timeout)
|
|
|
|
async def aclose(self) -> None:
|
|
await self._client.aclose()
|
|
|
|
async def submit_prompt(self, graph: dict[str, Any]) -> str:
|
|
"""Queue a graph; returns prompt_id. ComfyUI validates synchronously."""
|
|
payload = {"prompt": graph, "client_id": str(uuid.uuid4())}
|
|
try:
|
|
resp = await self._client.post("/prompt", json=payload)
|
|
except httpx.HTTPError as e:
|
|
raise ComfyUIError(f"ComfyUI unreachable: {e}") from e
|
|
if resp.status_code == 400:
|
|
detail = _validation_detail(resp.json())
|
|
raise ComfyUIError(f"ComfyUI rejected workflow: {detail}", code="comfyui_validation_failed")
|
|
if resp.status_code != 200:
|
|
raise ComfyUIError(f"ComfyUI /prompt returned HTTP {resp.status_code}")
|
|
prompt_id = resp.json().get("prompt_id")
|
|
if not prompt_id:
|
|
raise ComfyUIError("ComfyUI /prompt response missing prompt_id")
|
|
return prompt_id
|
|
|
|
async def upload_image(self, data: bytes, content_type: str = "image/png") -> str:
|
|
"""Upload an input image; returns the server-side filename."""
|
|
filename = f"adapter_{uuid.uuid4().hex[:8]}_{secrets.token_hex(4)}.png"
|
|
try:
|
|
resp = await self._client.post(
|
|
"/upload/image",
|
|
files={"image": (filename, data, content_type)},
|
|
)
|
|
except httpx.HTTPError as e:
|
|
raise ComfyUIError(f"ComfyUI unreachable: {e}") from e
|
|
if resp.status_code != 200:
|
|
raise ComfyUIError(f"ComfyUI /upload/image returned HTTP {resp.status_code}")
|
|
name = resp.json().get("name")
|
|
if not name:
|
|
raise ComfyUIError("ComfyUI /upload/image response missing name")
|
|
return name
|
|
|
|
async def wait_for_result(self, prompt_id: str, timeout: float, interval: float) -> dict:
|
|
"""Poll /history/{prompt_id} until the entry appears; returns the history entry."""
|
|
deadline = time.monotonic() + timeout
|
|
while True:
|
|
try:
|
|
resp = await self._client.get(f"/history/{prompt_id}")
|
|
except httpx.HTTPError as e:
|
|
raise ComfyUIError(f"ComfyUI unreachable: {e}") from e
|
|
if resp.status_code != 200:
|
|
raise ComfyUIError(f"ComfyUI /history returned HTTP {resp.status_code}")
|
|
entry = resp.json().get(prompt_id)
|
|
if entry is not None:
|
|
status = entry.get("status", {})
|
|
if status.get("status_str") == "error" or status.get("status_str") == "execution_failed":
|
|
raise ComfyUIError(f"ComfyUI execution failed for prompt {prompt_id}", code="comfyui_execution_failed")
|
|
return entry
|
|
if time.monotonic() >= deadline:
|
|
raise ComfyUIError(
|
|
f"Timed out after {timeout:.0f}s waiting for prompt {prompt_id}",
|
|
code="comfyui_timeout",
|
|
)
|
|
await asyncio.sleep(interval)
|
|
|
|
async def fetch_view(self, filename: str, subfolder: str = "") -> bytes:
|
|
"""Download a produced image from the output folder."""
|
|
try:
|
|
resp = await self._client.get(
|
|
"/view", params={"filename": filename, "subfolder": subfolder, "type": "output"}
|
|
)
|
|
except httpx.HTTPError as e:
|
|
raise ComfyUIError(f"ComfyUI unreachable: {e}") from e
|
|
if resp.status_code != 200:
|
|
raise ComfyUIError(f"ComfyUI /view returned HTTP {resp.status_code} for '{filename}'")
|
|
return resp.content
|
|
|
|
async def validate_graph(self, graph: dict[str, Any]) -> None:
|
|
"""Validate a graph without queueing (POST /prompt?front=false)."""
|
|
try:
|
|
resp = await self._client.post("/prompt", params={"front": "false"}, json={"prompt": graph})
|
|
except httpx.HTTPError as e:
|
|
raise ComfyUIError(f"ComfyUI unreachable: {e}") from e
|
|
if resp.status_code == 200:
|
|
return
|
|
if resp.status_code == 400:
|
|
raise ComfyUIError(
|
|
f"ComfyUI rejected workflow: {_validation_detail(resp.json())}",
|
|
code="comfyui_validation_failed",
|
|
)
|
|
raise ComfyUIError(f"ComfyUI /prompt?front=false returned HTTP {resp.status_code}")
|
|
|
|
|
|
def extract_output_image(history_entry: dict) -> tuple[str, str]:
|
|
"""(filename, subfolder) of the first output image, preferring SaveImage nodes."""
|
|
outputs = history_entry.get("outputs", {})
|
|
images: list[dict] = []
|
|
for node_output in outputs.values():
|
|
images.extend(img for img in node_output.get("images", []) if img.get("type") == "output")
|
|
if not images:
|
|
raise ComfyUIError("ComfyUI produced no output images", code="comfyui_no_output")
|
|
return images[0]["filename"], images[0].get("subfolder", "")
|
|
|
|
|
|
def _validation_detail(body: Any) -> str:
|
|
if isinstance(body, dict) and "node_errors" in body:
|
|
parts = []
|
|
for node_id, node_err in body["node_errors"].items():
|
|
for err in node_err.get("errors", []):
|
|
parts.append(f"node {node_id}: {err.get('message')} {err.get('details', '')}".strip())
|
|
return "; ".join(parts) or "validation failed"
|
|
return json.dumps(body)[:300] if body else "validation failed"
|