114 lines
3.6 KiB
Python
114 lines
3.6 KiB
Python
import asyncio
|
|
import json
|
|
from pathlib import Path
|
|
|
|
from fastapi import FastAPI, File, HTTPException, UploadFile, WebSocket, WebSocketDisconnect
|
|
from fastapi.responses import FileResponse, JSONResponse
|
|
from fastapi.staticfiles import StaticFiles
|
|
|
|
from . import db, registry, store
|
|
from .runner import runner
|
|
|
|
app = FastAPI(title="MODELBEAST")
|
|
WEB_DIST = db.ROOT / "web" / "dist"
|
|
|
|
|
|
@app.on_event("startup")
|
|
async def startup():
|
|
runner.operators = registry.load_operators()
|
|
app.state.con = db.connect()
|
|
asyncio.create_task(runner.worker())
|
|
|
|
|
|
# -- operators ----------------------------------------------------------------
|
|
@app.get("/api/operators")
|
|
def list_operators():
|
|
ops = []
|
|
for op in runner.operators.values():
|
|
public = {k: v for k, v in op.items() if k != "dir"}
|
|
ops.append(public)
|
|
return ops
|
|
|
|
|
|
# -- assets ---------------------------------------------------------------------
|
|
@app.get("/api/assets")
|
|
def list_assets():
|
|
return store.list_assets(app.state.con)
|
|
|
|
|
|
@app.post("/api/assets")
|
|
async def upload_asset(file: UploadFile = File(...)):
|
|
data = await file.read()
|
|
if not data:
|
|
raise HTTPException(400, "empty file")
|
|
return store.register_upload(app.state.con, file.filename or "unnamed", data)
|
|
|
|
|
|
@app.delete("/api/assets/{asset_id}")
|
|
def delete_asset(asset_id: str):
|
|
if not store.delete_asset(app.state.con, asset_id):
|
|
raise HTTPException(404, "no such asset")
|
|
return {"ok": True}
|
|
|
|
|
|
@app.get("/api/assets/{asset_id}/file")
|
|
def asset_file(asset_id: str, member: str | None = None):
|
|
asset = store.get_asset(app.state.con, asset_id)
|
|
if not asset:
|
|
raise HTTPException(404, "no such asset")
|
|
path = Path(asset["path"])
|
|
if member: # file inside a folder asset (e.g. one frame)
|
|
target = (path / member).resolve()
|
|
if not str(target).startswith(str(path.resolve())) or not target.is_file():
|
|
raise HTTPException(404, "no such member")
|
|
return FileResponse(target)
|
|
if path.is_dir():
|
|
files = sorted(p.name for p in path.iterdir() if p.is_file())
|
|
return JSONResponse({"folder": True, "files": files})
|
|
return FileResponse(path, filename=asset["name"])
|
|
|
|
|
|
# -- jobs -----------------------------------------------------------------------
|
|
@app.get("/api/jobs")
|
|
def list_jobs():
|
|
return runner.list_jobs(app.state.con)
|
|
|
|
|
|
@app.post("/api/jobs")
|
|
async def create_job(payload: dict):
|
|
operator = payload.get("operator")
|
|
if operator not in runner.operators:
|
|
raise HTTPException(400, f"unknown operator: {operator}")
|
|
asset_id = payload.get("asset_id")
|
|
if asset_id and not store.get_asset(app.state.con, asset_id):
|
|
raise HTTPException(400, "unknown asset")
|
|
params = payload.get("params") or {}
|
|
job = runner.create_job(app.state.con, operator, asset_id, params)
|
|
await runner.broadcast({"type": "job", "job": job})
|
|
return job
|
|
|
|
|
|
@app.get("/api/jobs/{job_id}")
|
|
def get_job(job_id: str):
|
|
job = runner.get_job(app.state.con, job_id)
|
|
if not job:
|
|
raise HTTPException(404, "no such job")
|
|
return job
|
|
|
|
|
|
# -- websocket ---------------------------------------------------------------
|
|
@app.websocket("/ws")
|
|
async def ws(websocket: WebSocket):
|
|
await websocket.accept()
|
|
runner.subscribers.add(websocket)
|
|
try:
|
|
while True:
|
|
await websocket.receive_text() # keepalive pings from client
|
|
except WebSocketDisconnect:
|
|
runner.subscribers.discard(websocket)
|
|
|
|
|
|
# -- frontend (built) -----------------------------------------------------------
|
|
if WEB_DIST.exists():
|
|
app.mount("/", StaticFiles(directory=WEB_DIST, html=True), name="web")
|