crbns/drestly-image-processing

1import os2from contextlib import asynccontextmanager3from io import BytesIO4import asyncio5 6from fastapi import FastAPI, BackgroundTasks, Depends7from pydantic import BaseModel8from PIL import Image9from rembg import remove, new_session10from supabase import create_client, Client11from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor12from telemetry import setup_telemetry13from opentelemetry import metrics, trace14 15from auth import verify_user16 17SUPABASE_URL = os.environ["SUPABASE_URL"]18SECRET_KEY = os.environ["SUPABASE_SECRET_KEY"]19ORIGINALS = "originals"20CUTOUTS = "cutouts"21 22setup_telemetry()23 24state: dict = {}25 26_meter = metrics.get_meter("drestly.inference")27_queued = _meter.create_up_down_counter(28    "inference.queue.depth",29    unit="{request}",30    description="Requests waiting to acquire an inference slot",31)32_active = _meter.create_up_down_counter(33    "inference.active",34    unit="{request}",35    description="Requests currently running inference",36)37 38INFERENCE_SLOTS = asyncio.Semaphore(2)39 40 41@asynccontextmanager42async def inference_slot():43    _queued.add(1)  # entered the line44    try:45        await INFERENCE_SLOTS.acquire()46    finally:47        _queued.add(-1)  # left the line (got in, or was cancelled)48    _active.add(1)  # running49    try:50        yield51    finally:52        _active.add(-1)53        INFERENCE_SLOTS.release()54 55 56@asynccontextmanager57async def lifespan(app: FastAPI):58    state["session"] = new_session("birefnet-general-lite")  # warm the model once59    state["supabase"] = create_client(SUPABASE_URL, SECRET_KEY)60    yield61    # Flush buffered spans/metrics before the process exits (BatchSpanProcessor62    # and the 60s metric reader would otherwise drop whatever is still queued).63    trace.get_tracer_provider().shutdown()64    metrics.get_meter_provider().shutdown()65    state.clear()66 67 68app = FastAPI(lifespan=lifespan, docs_url=None, redoc_url=None, openapi_url=None)69FastAPIInstrumentor.instrument_app(app)70 71 72class ProcessRequest(BaseModel):73    item_id: str74 75 76@app.get("/health")77def health():78    return {"ok": True}79 80 81@app.post("/process", status_code=202)82def process(83    req: ProcessRequest,84    background: BackgroundTasks,85    user_id: str = Depends(verify_user),86):87    background.add_task(run_cutout, req.item_id, user_id)88    return {"status": "processing", "item_id": req.item_id}89 90 91async def run_cutout(item_id: str, user_id: str):92    tracer = trace.get_tracer("drestly.inference")93    with tracer.start_as_current_span("run_cutout") as span:94        span.set_attribute("item.id", item_id)95        async with inference_slot():96            # rembg, Pillow, and the sync Supabase client all block; run them off the97            # event loop so /health pings and other requests stay responsive.98            await asyncio.to_thread(_do_cutout_blocking, item_id, user_id)99 100 101def _do_cutout_blocking(item_id: str, user_id: str):102    sb: Client = state["supabase"]103 104    # Service key bypasses RLS, so confirm ownership ourselves before touching anything.105    # maybe_single() yields data=None for a missing row instead of raising (unlike single()).106    row = (107        sb.table("clothing_items")108        .select("user_id, original_path")109        .eq("id", item_id)110        .maybe_single()111        .execute()112    )113    if not row or not isinstance(row.data, dict) or row.data["user_id"] != user_id:114        return115 116    original_path = row.data["original_path"]117    try:118        original_bytes = sb.storage.from_(ORIGINALS).download(original_path)119 120        img = Image.open(BytesIO(original_bytes))121        out = remove(img, session=state["session"])122        out.thumbnail((1500, 1500))  # cap resolution: in-place, preserves aspect, only downscales123        buf = BytesIO()124        out.save(buf, format="WEBP", quality=80, method=6)125 126        cutout_path = f"{user_id}/{item_id}.webp"127        sb.storage.from_(CUTOUTS).upload(128            cutout_path,129            buf.getvalue(),130            file_options={"content-type": "image/webp", "upsert": "true"},131        )132 133        buf.close()134 135        sb.table("clothing_items").update(136            {137                "status": "done",138                "cutout_path": cutout_path,139            }140        ).eq("id", item_id).execute()141 142    except Exception as e:143        sb.table("clothing_items").update(144            {145                "status": "failed",146                "error": str(e)[:500],147            }148        ).eq("id", item_id).execute()149        return150 151    try:152        sb.storage.from_(ORIGINALS).remove([original_path])153    except Exception:154        pass155