372 lines
17 KiB
Python
372 lines
17 KiB
Python
from __future__ import annotations
|
|
|
|
import base64
|
|
import asyncio
|
|
import ipaddress
|
|
import json
|
|
import mimetypes
|
|
import os
|
|
import secrets
|
|
import socket
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
from urllib.parse import quote, urlencode, urlsplit
|
|
|
|
import httpx
|
|
from fastapi import Depends, FastAPI, Header, HTTPException, Response
|
|
from fastapi.responses import StreamingResponse
|
|
from pydantic import BaseModel, Field
|
|
|
|
from workflow import MAX_VIDEO_DURATION_SECONDS, compile_workflow, load_registry, public_model, public_workflow
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Provider:
|
|
id: str
|
|
name: str
|
|
base_url: str
|
|
api_key: str = ""
|
|
|
|
|
|
class JobRequest(BaseModel):
|
|
provider_id: str = "default"
|
|
workflow_key: str
|
|
prompt: str
|
|
negative_prompt: str = ""
|
|
input_images: list[str] = Field(default_factory=list, max_length=9)
|
|
width: int | None = Field(default=None, ge=64, le=8192)
|
|
height: int | None = Field(default=None, ge=64, le=8192)
|
|
duration: float | None = Field(default=None, gt=0, le=MAX_VIDEO_DURATION_SECONDS)
|
|
generate_audio: bool | None = None
|
|
seed: int | None = Field(default=None, ge=0)
|
|
batch_size: int = Field(default=1, ge=1, le=4)
|
|
denoise: float | None = Field(default=None, ge=0, le=1)
|
|
metadata: dict[str, Any] = Field(default_factory=dict)
|
|
|
|
|
|
app = FastAPI(title="Canvas ComfyUI Workflow Adapter", version="1.0.0")
|
|
provider_selection_lock = asyncio.Lock()
|
|
provider_reservations: dict[str, int] = {}
|
|
|
|
|
|
def auto_balance_enabled() -> bool:
|
|
return os.getenv("COMFY_AUTO_BALANCE", "").strip().lower() in {"1", "true", "yes", "on"}
|
|
|
|
|
|
def providers() -> dict[str, Provider]:
|
|
raw = os.getenv("COMFY_PROVIDERS_JSON", "").strip()
|
|
rows: list[dict[str, Any]] = []
|
|
if raw:
|
|
parsed = json.loads(raw)
|
|
rows = parsed if isinstance(parsed, list) else []
|
|
elif os.getenv("COMFY_URL", "").strip():
|
|
rows = [{"id": "default", "name": "ComfyUI", "base_url": os.getenv("COMFY_URL"), "api_key": os.getenv("COMFY_API_KEY", "")}]
|
|
result: dict[str, Provider] = {}
|
|
for row in rows:
|
|
provider = Provider(id=str(row.get("id") or "").strip(), name=str(row.get("name") or row.get("id") or "ComfyUI"), base_url=str(row.get("base_url") or "").strip().rstrip("/"), api_key=str(row.get("api_key") or ""))
|
|
if provider.id and provider.base_url.startswith(("http://", "https://")):
|
|
result[provider.id] = provider
|
|
return result
|
|
|
|
|
|
def require_token(authorization: str | None = Header(default=None)) -> None:
|
|
expected = os.getenv("COMFY_ADAPTER_TOKEN", "").strip()
|
|
if expected and not secrets.compare_digest(authorization or "", f"Bearer {expected}"):
|
|
raise HTTPException(status_code=401, detail="invalid adapter token")
|
|
|
|
|
|
def provider_headers(provider: Provider) -> dict[str, str]:
|
|
return {"Authorization": f"Bearer {provider.api_key}"} if provider.api_key else {}
|
|
|
|
|
|
async def provider_queue_depth(provider: Provider) -> int | None:
|
|
try:
|
|
async with httpx.AsyncClient(timeout=5) as client:
|
|
response = await client.get(f"{provider.base_url}/queue", headers=provider_headers(provider))
|
|
response.raise_for_status()
|
|
payload = response.json()
|
|
except (httpx.HTTPError, ValueError, TypeError):
|
|
return None
|
|
return len(payload.get("queue_running") or []) + len(payload.get("queue_pending") or [])
|
|
|
|
|
|
async def select_provider(requested_provider_id: str) -> Provider | None:
|
|
configured = providers()
|
|
requested = configured.get(requested_provider_id)
|
|
if requested is None or requested_provider_id != "default" or not auto_balance_enabled() or len(configured) < 2:
|
|
return requested
|
|
|
|
candidates = list(configured.values())
|
|
depths = await asyncio.gather(*(provider_queue_depth(provider) for provider in candidates))
|
|
available = [
|
|
(depth + provider_reservations.get(provider.id, 0), 0 if provider.id == requested_provider_id else 1, provider)
|
|
for provider, depth in zip(candidates, depths)
|
|
if depth is not None
|
|
]
|
|
if not available:
|
|
return requested
|
|
return min(available, key=lambda item: (item[0], item[1]))[2]
|
|
|
|
|
|
async def reserve_provider(requested_provider_id: str) -> Provider | None:
|
|
async with provider_selection_lock:
|
|
provider = await select_provider(requested_provider_id)
|
|
if provider is not None:
|
|
provider_reservations[provider.id] = provider_reservations.get(provider.id, 0) + 1
|
|
return provider
|
|
|
|
|
|
async def release_provider(provider: Provider) -> None:
|
|
async with provider_selection_lock:
|
|
remaining = provider_reservations.get(provider.id, 0) - 1
|
|
if remaining > 0:
|
|
provider_reservations[provider.id] = remaining
|
|
else:
|
|
provider_reservations.pop(provider.id, None)
|
|
|
|
|
|
def input_host_allowlist() -> set[str]:
|
|
return {value.strip().lower().rstrip(".") for value in os.getenv("COMFY_INPUT_HOST_ALLOWLIST", "").split(",") if value.strip()}
|
|
|
|
|
|
async def validate_input_url(value: str) -> None:
|
|
parsed = urlsplit(value)
|
|
host = (parsed.hostname or "").lower().rstrip(".")
|
|
if parsed.scheme not in {"http", "https"} or not host or parsed.username or parsed.password:
|
|
raise HTTPException(status_code=400, detail="reference image URL is invalid")
|
|
if host in input_host_allowlist():
|
|
return
|
|
try:
|
|
infos = await asyncio.to_thread(socket.getaddrinfo, host, parsed.port or (443 if parsed.scheme == "https" else 80), type=socket.SOCK_STREAM)
|
|
except OSError as exc:
|
|
raise HTTPException(status_code=400, detail="reference image host cannot be resolved") from exc
|
|
if not infos:
|
|
raise HTTPException(status_code=400, detail="reference image host cannot be resolved")
|
|
for info in infos:
|
|
address = ipaddress.ip_address(info[4][0])
|
|
if not address.is_global:
|
|
raise HTTPException(status_code=400, detail="reference image URL cannot target a private address")
|
|
|
|
|
|
async def download_input_image(client: httpx.AsyncClient, value: str) -> tuple[bytes, str]:
|
|
await validate_input_url(value)
|
|
chunks: list[bytes] = []
|
|
total = 0
|
|
async with client.stream("GET", value, follow_redirects=False) as response:
|
|
response.raise_for_status()
|
|
if 300 <= response.status_code < 400:
|
|
raise HTTPException(status_code=400, detail="reference image redirects are not accepted")
|
|
mime_type = response.headers.get("content-type", "image/png").split(";", 1)[0]
|
|
if not mime_type.startswith("image/"):
|
|
raise HTTPException(status_code=400, detail="reference image URL did not return an image")
|
|
async for chunk in response.aiter_bytes():
|
|
total += len(chunk)
|
|
if total > 30 * 1024 * 1024:
|
|
raise HTTPException(status_code=413, detail="reference image exceeds 30 MB")
|
|
chunks.append(chunk)
|
|
return b"".join(chunks), mime_type
|
|
|
|
|
|
def encode_job(provider_id: str, prompt_id: str) -> str:
|
|
return base64.urlsafe_b64encode(f"{provider_id}:{prompt_id}".encode()).decode().rstrip("=")
|
|
|
|
|
|
def decode_job(job_id: str) -> tuple[Provider, str]:
|
|
try:
|
|
raw = base64.urlsafe_b64decode(job_id + "=" * (-len(job_id) % 4)).decode()
|
|
provider_id, prompt_id = raw.split(":", 1)
|
|
provider = providers()[provider_id]
|
|
except Exception as exc:
|
|
raise HTTPException(status_code=404, detail="job not found") from exc
|
|
return provider, prompt_id
|
|
|
|
|
|
@app.get("/health")
|
|
async def health() -> dict[str, Any]:
|
|
return {
|
|
"status": "ok",
|
|
"providers": len(providers()),
|
|
"workflows": len(load_registry()),
|
|
"autoBalance": auto_balance_enabled(),
|
|
}
|
|
|
|
|
|
@app.get("/v1/providers", dependencies=[Depends(require_token)])
|
|
async def list_providers() -> dict[str, Any]:
|
|
return {"providers": [{"id": item.id, "name": item.name, "baseUrl": item.base_url} for item in providers().values()]}
|
|
|
|
|
|
@app.get("/v1/workflows", dependencies=[Depends(require_token)])
|
|
async def list_workflows() -> dict[str, Any]:
|
|
return {"workflows": [public_workflow(item) for item in load_registry().values()]}
|
|
|
|
|
|
@app.get("/v1/models", dependencies=[Depends(require_token)])
|
|
async def list_models() -> dict[str, Any]:
|
|
return {"object": "list", "data": [public_model(item) for item in load_registry().values()]}
|
|
|
|
|
|
async def upload_image(client: httpx.AsyncClient, provider: Provider, value: str, index: int, client_id: str) -> str:
|
|
value = value.strip()
|
|
if value.startswith("comfy://"):
|
|
return value.removeprefix("comfy://")
|
|
if value.startswith("data:"):
|
|
header, encoded = value.split(",", 1)
|
|
mime_type = header.split(";", 1)[0].removeprefix("data:") or "image/png"
|
|
if not mime_type.startswith("image/"):
|
|
raise HTTPException(status_code=400, detail="reference data URL must contain an image")
|
|
try:
|
|
data = base64.b64decode(encoded, validate=True)
|
|
except (ValueError, base64.binascii.Error) as exc:
|
|
raise HTTPException(status_code=400, detail="reference image data URL is invalid") from exc
|
|
elif value.startswith(("http://", "https://")):
|
|
data, mime_type = await download_input_image(client, value)
|
|
else:
|
|
raise HTTPException(status_code=400, detail="input_images must contain http(s), data, or comfy URLs")
|
|
if len(data) > 30 * 1024 * 1024:
|
|
raise HTTPException(status_code=413, detail="reference image exceeds 30 MB")
|
|
extension = mimetypes.guess_extension(mime_type) or ".png"
|
|
filename = f"{client_id}-reference-{index}{extension}"
|
|
response = await client.post(f"{provider.base_url}/upload/image", files={"image": (filename, data, mime_type)}, data={"overwrite": "true"}, headers=provider_headers(provider))
|
|
response.raise_for_status()
|
|
payload = response.json()
|
|
return str(payload.get("name") or filename)
|
|
|
|
|
|
@app.post("/v1/jobs", dependencies=[Depends(require_token)])
|
|
async def create_job(request: JobRequest) -> dict[str, Any]:
|
|
registry = load_registry()
|
|
spec = registry.get(request.workflow_key)
|
|
if spec is None:
|
|
raise HTTPException(status_code=400, detail="workflow is not registered")
|
|
provider = await reserve_provider(request.provider_id)
|
|
if provider is None:
|
|
raise HTTPException(status_code=400, detail="provider is not configured")
|
|
try:
|
|
client_id = f"canvas-{secrets.token_hex(8)}"
|
|
async with httpx.AsyncClient(timeout=60) as client:
|
|
uploaded = [await upload_image(client, provider, value, index, client_id) for index, value in enumerate(request.input_images, start=1)]
|
|
try:
|
|
workflow = compile_workflow(spec, request.model_dump(), uploaded, client_id)
|
|
except ValueError as exc:
|
|
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
|
response = await client.post(f"{provider.base_url}/prompt", json={"prompt": workflow, "client_id": client_id}, headers=provider_headers(provider))
|
|
response.raise_for_status()
|
|
payload = response.json()
|
|
finally:
|
|
await release_provider(provider)
|
|
prompt_id = str(payload.get("prompt_id") or "").strip()
|
|
if not prompt_id:
|
|
raise HTTPException(status_code=502, detail="ComfyUI did not return prompt_id")
|
|
return {"id": encode_job(provider.id, prompt_id), "promptId": prompt_id, "providerId": provider.id, "workflowKey": spec.key, "workflowRevision": spec.revision, "status": "submitted"}
|
|
|
|
|
|
def extract_outputs(entry: dict[str, Any], job_id: str) -> list[dict[str, Any]]:
|
|
result: list[dict[str, Any]] = []
|
|
for node in (entry.get("outputs") or {}).values():
|
|
if not isinstance(node, dict):
|
|
continue
|
|
for field in ("images", "gifs", "audio"):
|
|
for item in node.get(field) or []:
|
|
if not isinstance(item, dict) or not item.get("filename"):
|
|
continue
|
|
filename = str(item["filename"])
|
|
mime_type = mimetypes.guess_type(filename)[0] or "application/octet-stream"
|
|
kind = "video" if mime_type.startswith("video/") or filename.lower().endswith((".gif", ".webp")) and field == "gifs" else "audio" if mime_type.startswith("audio/") else "image"
|
|
index = len(result)
|
|
result.append({"index": index, "kind": kind, "mimeType": mime_type, "filename": filename, "url": f"/jobs/{quote(job_id, safe='')}/outputs/{index}"})
|
|
return result
|
|
|
|
|
|
async def job_state(job_id: str) -> tuple[Provider, str, dict[str, Any], list[dict[str, Any]]]:
|
|
provider, prompt_id = decode_job(job_id)
|
|
async with httpx.AsyncClient(timeout=30) as client:
|
|
response = await client.get(f"{provider.base_url}/history/{quote(prompt_id, safe='')}", headers=provider_headers(provider))
|
|
response.raise_for_status()
|
|
history = response.json()
|
|
entry = history.get(prompt_id) or {}
|
|
outputs = extract_outputs(entry, job_id)
|
|
return provider, prompt_id, entry, outputs
|
|
|
|
|
|
@app.get("/v1/jobs/{job_id}", dependencies=[Depends(require_token)])
|
|
async def get_job(job_id: str) -> dict[str, Any]:
|
|
try:
|
|
provider, prompt_id, entry, outputs = await job_state(job_id)
|
|
except httpx.TimeoutException:
|
|
provider, prompt_id = decode_job(job_id)
|
|
return {
|
|
"id": job_id,
|
|
"promptId": prompt_id,
|
|
"providerId": provider.id,
|
|
"status": "running",
|
|
"outputs": [],
|
|
"error": "",
|
|
"pollDeferred": True,
|
|
}
|
|
status_payload = entry.get("status") or {}
|
|
if status_payload.get("status_str") == "error":
|
|
status = "failed"
|
|
elif outputs:
|
|
status = "succeeded"
|
|
elif status_payload.get("completed"):
|
|
status = "failed"
|
|
else:
|
|
status = "running"
|
|
error = "ComfyUI execution failed" if status == "failed" else ""
|
|
if status == "failed":
|
|
for event in status_payload.get("messages") or []:
|
|
if not isinstance(event, (list, tuple)) or len(event) < 2 or event[0] != "execution_error" or not isinstance(event[1], dict):
|
|
continue
|
|
# 只映射已识别的内核错误，不向用户返回节点输入、提示词或 traceback。
|
|
if "tensor strides exceed int32 range" in str(event[1].get("exception_message") or ""):
|
|
error = "当前时长与分辨率组合超出采样内核范围，请缩短时长后重试（最多 15 秒）"
|
|
break
|
|
return {"id": job_id, "promptId": prompt_id, "providerId": provider.id, "status": status, "outputs": outputs, "error": error}
|
|
|
|
|
|
@app.post("/v1/jobs/{job_id}/cancel", dependencies=[Depends(require_token)])
|
|
async def cancel_job(job_id: str) -> dict[str, Any]:
|
|
provider, prompt_id = decode_job(job_id)
|
|
async with httpx.AsyncClient(timeout=30) as client:
|
|
queue = (await client.get(f"{provider.base_url}/queue", headers=provider_headers(provider))).json()
|
|
running = {str(item[1]) for item in queue.get("queue_running") or [] if isinstance(item, list) and len(item) > 1}
|
|
if prompt_id in running:
|
|
await client.post(f"{provider.base_url}/interrupt", json={}, headers=provider_headers(provider))
|
|
await client.post(f"{provider.base_url}/queue", json={"delete": [prompt_id]}, headers=provider_headers(provider))
|
|
return {"id": job_id, "status": "cancelled"}
|
|
|
|
|
|
@app.get("/v1/jobs/{job_id}/outputs/{index}", dependencies=[Depends(require_token)])
|
|
async def get_output(job_id: str, index: int) -> Response:
|
|
provider, _, entry, outputs = await job_state(job_id)
|
|
if index < 0 or index >= len(outputs):
|
|
raise HTTPException(status_code=404, detail="output not found")
|
|
wanted = outputs[index]
|
|
locator: dict[str, Any] | None = None
|
|
for node in (entry.get("outputs") or {}).values():
|
|
if not isinstance(node, dict):
|
|
continue
|
|
for field in ("images", "gifs", "audio"):
|
|
for item in node.get(field) or []:
|
|
if isinstance(item, dict) and str(item.get("filename")) == wanted["filename"]:
|
|
locator = item
|
|
break
|
|
if locator: break
|
|
if locator: break
|
|
if locator is None:
|
|
raise HTTPException(status_code=404, detail="output locator not found")
|
|
query = urlencode({"filename": locator.get("filename", ""), "subfolder": locator.get("subfolder", ""), "type": locator.get("type", "output")})
|
|
client = httpx.AsyncClient(timeout=120)
|
|
response = await client.send(client.build_request("GET", f"{provider.base_url}/view?{query}", headers=provider_headers(provider)), stream=True)
|
|
if response.status_code >= 400:
|
|
await response.aclose(); await client.aclose()
|
|
raise HTTPException(status_code=502, detail="failed to read ComfyUI output")
|
|
async def stream():
|
|
try:
|
|
async for chunk in response.aiter_bytes():
|
|
yield chunk
|
|
finally:
|
|
await response.aclose(); await client.aclose()
|
|
return StreamingResponse(stream(), media_type=response.headers.get("content-type", wanted["mimeType"]))
|