api.py

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())