171 lines
6.6 KiB
Python
171 lines
6.6 KiB
Python
"""HTTP surface of the engine (FastAPI). Runs on its own thread; every call
|
|
hands work to the Engine and returns, except /generate which streams the
|
|
job's frames for as long as the generation runs."""
|
|
|
|
import asyncio
|
|
import os
|
|
import re
|
|
import socket
|
|
import threading
|
|
import time
|
|
import uuid
|
|
from typing import Optional, Tuple
|
|
|
|
from fastapi import FastAPI, HTTPException, Request
|
|
from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
|
|
|
|
from toolkit.models.registry import describe_archs
|
|
|
|
from .engine import END, Engine
|
|
|
|
_SAFE_NAME = re.compile(r"[^A-Za-z0-9._-]+")
|
|
|
|
|
|
def create_app(engine: Engine, token: Optional[str] = None) -> FastAPI:
|
|
app = FastAPI(title="AI Toolkit Inference Engine", docs_url=None, redoc_url=None)
|
|
|
|
def check_auth(request: Request):
|
|
if not token:
|
|
return
|
|
supplied = request.headers.get("x-engine-token")
|
|
if supplied is None:
|
|
auth = request.headers.get("authorization", "")
|
|
if auth.lower().startswith("bearer "):
|
|
supplied = auth[7:]
|
|
if supplied != token:
|
|
raise HTTPException(status_code=401, detail="bad engine token")
|
|
|
|
@app.get("/health")
|
|
async def health():
|
|
return engine.health()
|
|
|
|
@app.get("/models")
|
|
async def models(request: Request):
|
|
check_auth(request)
|
|
return {"archs": await asyncio.to_thread(describe_archs)}
|
|
|
|
@app.get("/queue")
|
|
async def queue_(request: Request):
|
|
check_auth(request)
|
|
return {"queued": engine.queued(), "recent": engine.recent()}
|
|
|
|
@app.post("/cancel/{request_id}")
|
|
async def cancel(request: Request, request_id: str):
|
|
check_auth(request)
|
|
return {"ok": engine.cancel(request_id)}
|
|
|
|
@app.post("/unload")
|
|
async def unload(request: Request):
|
|
check_auth(request)
|
|
if engine.current is not None:
|
|
raise HTTPException(status_code=409, detail="busy")
|
|
await asyncio.to_thread(engine.unload)
|
|
return {"ok": True}
|
|
|
|
@app.post("/assets")
|
|
async def upload_asset(request: Request, name: str = "asset.png"):
|
|
"""Raw-body upload (no multipart dependency). Returns the stored path,
|
|
usable as ctrl_img in a generate request."""
|
|
check_auth(request)
|
|
body = await request.body()
|
|
if not body:
|
|
raise HTTPException(status_code=400, detail="empty body")
|
|
base = _SAFE_NAME.sub("_", os.path.basename(name)) or "asset"
|
|
stem, ext = os.path.splitext(base)
|
|
asset_id = f"{stem}_{uuid.uuid4().hex[:8]}{ext or '.bin'}"
|
|
path = os.path.join(engine.assets_folder, asset_id)
|
|
with open(path, "wb") as f:
|
|
f.write(body)
|
|
return {"id": asset_id, "path": path, "bytes": len(body)}
|
|
|
|
@app.get("/outputs/{relpath:path}")
|
|
async def get_output(request: Request, relpath: str):
|
|
check_auth(request)
|
|
root = os.path.realpath(engine.output_folder)
|
|
path = os.path.realpath(os.path.join(root, relpath))
|
|
if not path.startswith(root + os.sep) or not os.path.isfile(path):
|
|
raise HTTPException(status_code=404, detail="not found")
|
|
return FileResponse(path)
|
|
|
|
@app.post("/generate")
|
|
async def generate(request: Request):
|
|
check_auth(request)
|
|
body = await request.json()
|
|
model = body.get("model")
|
|
sample = body.get("sample") or {}
|
|
if not isinstance(model, dict) or not model.get("arch"):
|
|
raise HTTPException(status_code=400, detail="model.arch is required")
|
|
if body.get("wait"):
|
|
# plain JSON mode for scripts: block until the job is done, drop frames
|
|
job = engine.submit(model, sample, {"latents": "none"})
|
|
q = job.subscribe(replay=False)
|
|
while True:
|
|
frame = await asyncio.to_thread(q.get)
|
|
if frame is END:
|
|
break
|
|
return JSONResponse(job.info(), status_code=200 if job.status == "done" else 500)
|
|
|
|
job = engine.submit(model, sample, body.get("stream"))
|
|
# a dropped connection (page reload) must not cancel: the client
|
|
# reattaches via /stream/{id}; cancelling is explicit (/cancel)
|
|
return _stream_job(job, cancel_on_disconnect=False)
|
|
|
|
@app.get("/stream/{request_id}")
|
|
async def stream(request: Request, request_id: str):
|
|
"""(Re)attach to a request's frame stream: replays everything so far
|
|
(status/progress/result frames and the newest latent) then follows
|
|
live until the end frame. Disconnecting does not cancel the job."""
|
|
check_auth(request)
|
|
job = engine.get_job(request_id)
|
|
if job is None:
|
|
raise HTTPException(status_code=404, detail="unknown request")
|
|
return _stream_job(job, cancel_on_disconnect=False)
|
|
|
|
def _stream_job(job, cancel_on_disconnect: bool):
|
|
q = job.subscribe(replay=True)
|
|
|
|
async def frames():
|
|
try:
|
|
while True:
|
|
frame = await asyncio.to_thread(q.get)
|
|
if frame is END:
|
|
break
|
|
yield frame
|
|
finally:
|
|
if cancel_on_disconnect and job.status in ("queued", "running"):
|
|
engine.cancel(job.request_id)
|
|
|
|
return StreamingResponse(
|
|
frames(),
|
|
media_type="application/octet-stream",
|
|
headers={"X-Request-Id": job.request_id, "Cache-Control": "no-store"},
|
|
)
|
|
|
|
return app
|
|
|
|
|
|
def serve_in_thread(app, host: str = "127.0.0.1", port: int = 0, timeout: float = 30.0) -> Tuple[object, threading.Thread, str, int]:
|
|
"""Bind the socket ourselves (port 0 = ephemeral) and run uvicorn on a
|
|
daemon thread. Returns (server, thread, host, port); the caller stops it
|
|
with server.should_exit = True. uvicorn only installs signal handlers on
|
|
the main thread, so the process keeps its own."""
|
|
import uvicorn
|
|
|
|
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
|
sock.bind((host, port))
|
|
sock.listen(64)
|
|
host, port = sock.getsockname()[:2]
|
|
config = uvicorn.Config(app, log_level="warning", access_log=False, timeout_keep_alive=600)
|
|
server = uvicorn.Server(config)
|
|
thread = threading.Thread(target=server.run, kwargs={"sockets": [sock]}, daemon=True, name="engine-http")
|
|
thread.start()
|
|
deadline = time.time() + timeout
|
|
while not server.started and time.time() < deadline:
|
|
if not thread.is_alive():
|
|
raise RuntimeError("engine http server died on startup")
|
|
time.sleep(0.05)
|
|
if not server.started:
|
|
raise RuntimeError("engine http server failed to start")
|
|
return server, thread, host, port
|