File size: 3,028 Bytes
f7d8394 6f7704a f7d8394 6f7704a f7d8394 6f7704a f7d8394 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | import hmac
import io
import zipfile
from fastapi import APIRouter, File, HTTPException, Query, UploadFile
from fastapi.responses import JSONResponse, Response
import config
import inference as inf # tests monkeypatch api.inf
router = APIRouter()
def _validate_image(upload: UploadFile, data: bytes):
if not (upload.content_type or "").startswith("image/"):
raise HTTPException(status_code=400, detail="file must be an image")
if len(data) > config.MAX_FILE_BYTES:
raise HTTPException(
status_code=400,
detail=f"file exceeds {config.MAX_FILE_MB} MB limit",
)
@router.get("/health")
async def health():
return {
"status": "ok",
"model": config.MODEL_NAME,
"queue": inf.queue_depth(),
"in_flight": inf.in_flight(),
}
@router.get("/jobs")
async def list_jobs():
job_list = inf.jobs()
return {
"total": len(job_list),
"queue": inf.queue_depth(),
"in_flight": inf.in_flight(),
"jobs": job_list,
}
@router.post("/remove")
async def remove(file: UploadFile = File(...)):
data = await file.read()
_validate_image(file, data)
out = await inf.remove_bg(data, label=file.filename)
return Response(content=out, media_type="image/png")
@router.post("/remove/batch")
async def remove_batch(files: list[UploadFile] = File(...)):
if len(files) > config.MAX_BATCH_FILES:
raise HTTPException(
status_code=400,
detail=f"too many files (max {config.MAX_BATCH_FILES})",
)
buf = io.BytesIO()
errors = []
seen = {}
with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
for upload in files:
data = await upload.read()
name = upload.filename or "image.png"
try:
_validate_image(upload, data)
out = await inf.remove_bg(data, label=name)
except HTTPException as exc:
errors.append(f"{name}: {exc.detail}")
continue
except Exception as exc: # noqa: BLE001
errors.append(f"{name}: {exc!r}")
continue
# ensure unique entry names + .png extension
stem = name.rsplit(".", 1)[0]
out_name = f"{stem}.png"
n = seen.get(out_name, 0)
seen[out_name] = n + 1
if n:
out_name = f"{stem}_{n}.png"
zf.writestr(out_name, out)
if errors:
zf.writestr("_errors.txt", "\n".join(errors))
return Response(content=buf.getvalue(), media_type="application/zip")
@router.get("/kill")
async def kill(key: str | None = Query(default=None)):
if not config.KILL_API_KEY:
raise HTTPException(
status_code=503, detail="kill endpoint not configured"
)
if not key or not hmac.compare_digest(key, config.KILL_API_KEY):
raise HTTPException(status_code=401, detail="invalid key")
return JSONResponse(inf.kill_all())
|