app.py

from fastapi import FastAPI, File, UploadFile, Header, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import Response
from PIL import Image
import torch
import torch.nn.functional as F
import numpy as np
import io
import sys
import os
import asyncio
from torchvision.transforms.functional import normalize
from safetensors.torch import load_file

# ── CPU Thread Tuning ──────────────────────────────────────────────
# Must be set before any model loading
# Tells PyTorch, OpenMP, and MKL how many threads to use
# Prevents thread contention inside Docker containers
torch.set_num_threads(2)       # P3 fix: 2 threads = 1 per CPU core on HF free tier
os.environ["OMP_NUM_THREADS"] = "2"  # was 4 — 32 threads fighting 2 cores caused thrashing
os.environ["MKL_NUM_THREADS"] = "2"

# ── Dynamo Safety Net ──────────────────────────────────────────────
# If torch.compile() fails for ANY reason (missing g++, ISA detection,
# etc.) fall back to eager mode instead of crashing the request.
# Primary fix is g++ in Dockerfile; this is the belt-and-suspenders.
import torch._dynamo
torch._dynamo.config.suppress_errors = True

# ── App Setup ──────────────────────────────────────────────────────
app = FastAPI(title="Lampdra API", version="1.0")

ALLOWED_ORIGINS = [
    "https://lampdra.com",
    "https://www.lampdra.com",
    "https://paul1k-lampdra-api.hf.space",
]

app.add_middleware(
    CORSMiddleware,
    allow_origins=ALLOWED_ORIGINS,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ── API Key ────────────────────────────────────────────────────────
API_KEY = os.environ.get("LAMPDRA_API_KEY", "")

def verify_key(x_api_key: str = Header(default="")):
    # Fail-closed: if LAMPDRA_API_KEY env var is not set, reject everything.
    # Previously: if API_KEY and ... → empty string is falsy → check skipped entirely.
    if not API_KEY or x_api_key != API_KEY:
        raise HTTPException(status_code=401, detail="Invalid API key")

# ── Request Limits ─────────────────────────────────────────────────
# File size checked first — cheapest possible check, no decoding needed
# Pixel count checked after decoding — catches malicious small-file/huge-pixel PNGs
MAX_FILE_SIZE = 20 * 1024 * 1024   # 20MB — stops RAM spike before it starts
MAX_PIXELS    = 10000 * 10000       # 100MP — generous for real photos

# ── Model Pool Setup ───────────────────────────────────────────────
# P1+P2 fix: pool=1 (was 4)
# HF free tier has ~2 CPU cores per Space.
# pool=4 meant 8 model copies sharing 2 cores = each got 0.25 cores = 4x slower.
# pool=1 means 1 model per uvicorn process gets all 2 cores = fastest inference.
# GIL: pool=4 in one process = models take turns anyway — pool=1 removes illusion.
# 2 uvicorn workers × pool=1 = 2 real parallel inferences per Space.
# 5 Spaces × 2 = 10 truly parallel inferences system-wide.
# RAM: 2 workers × 1 copy × 1.2GB = ~2.4GB (was 9.6GB — 4x less RAM used)
POOL_SIZE = 1
DEVICE    = "cuda" if torch.cuda.is_available() else "cpu"

sys.path.insert(0, "/app")
from rmbg.briarmbg import BriaRMBG

# Load weights from disk once — all pool copies share this read
print(f"Device: {DEVICE}")
print(f"Loading weights from disk...")
state_dict = load_file("/app/rmbg/model.safetensors")

# Pool queue — asyncio.Queue is thread-safe and blocks callers
# when all copies are busy until one becomes free
model_pool = asyncio.Queue()

def build_model_copy():
    """Build one independent copy of the model."""
    m = BriaRMBG()
    m.load_state_dict(state_dict)
    m.eval()
    m.to(DEVICE)
    # torch.compile — compiles model to native machine code
    # First inference is slower (compile happens then)
    # Every subsequent call is 20-40% faster
    try:
        m = torch.compile(m)
    except Exception:
        pass  # compile not available on all platforms, safe to skip
    return m

# ── Ready flag — Cause 3 fix ──────────────────────────────────────
# root() returns pool_available=0 until this is True.
# Guarantees router never sends requests during model loading or warmup.
_model_ready = False

# ── Startup: Warm-up THEN Fill Pool — Cause 1 fix ─────────────────
# OLD order: fill pool → warmup borrows the only slot → real requests
#            arrive, queue inside worker, router 504s before warmup ends.
# NEW order: warmup first (pool stays empty → router sees 0 → sends nothing)
#            then fill pool → router sees 1 → safe to send real requests.
@app.on_event("startup")
async def startup_event():
    global _model_ready
    dummy = torch.zeros(1, 3, 1024, 1024).to(DEVICE)
    for i in range(POOL_SIZE):
        m = build_model_copy()
        # Warmup in thread — event loop stays free, but pool is still empty
        # so router cannot send requests here even if it tries
        def _run_warmup(model=m):
            with torch.inference_mode():
                model(dummy)
        await asyncio.to_thread(_run_warmup)
        print(f"  Copy {i + 1}/{POOL_SIZE} warmed up and compiled")
        model_pool.put_nowait(m)  # slot added AFTER warmup — not before
    del dummy
    _model_ready = True           # only NOW does root() report real pool_available
    print(f"Pool ready — {POOL_SIZE} copies loaded, warmed up, accepting requests.")

# ── Preprocessing ──────────────────────────────────────────────────
def preprocess(image: Image.Image, size=(1024, 1024)):
    # All images resize to 1024×1024 here regardless of original size
    # This is why the model RAM usage is always predictable
    img    = image.convert("RGB").resize(size, Image.BOX)  # BOX=18ms vs BILINEAR=27ms for downscale
    arr    = np.array(img, dtype=np.float32) / 255.0
    tensor = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0).to(DEVICE)
    tensor = normalize(tensor, [0.5, 0.5, 0.5], [1.0, 1.0, 1.0])
    return tensor

# ── Routes ─────────────────────────────────────────────────────────
@app.get("/")
@app.head("/")
def root():
    return {
        "status"          : "Lampdra API is running",
        "device"          : DEVICE,
        "pool_size"       : POOL_SIZE,
        # Returns 0 until warmup completes — prevents router sending
        # requests to a Space that is still loading or compiling
        "pool_available"  : model_pool.qsize() if _model_ready else 0,
    }

@app.post("/remove-background")
async def remove_background(
    file      : UploadFile = File(...),
    x_api_key : str        = Header(default=""),
    origin    : str        = Header(default="")
):
    # ── Server-side origin check ───────────────────────────────────
    # The worker should only ever be called by the router, never directly
    # by browsers or any other caller. Only the router's HF space is allowed.
    ALLOWED_POST_ORIGINS = {
        "https://paul1k-lampdra-router.hf.space",
    }
    if origin not in ALLOWED_POST_ORIGINS:
        raise HTTPException(status_code=403, detail="Forbidden")

    verify_key(x_api_key)

    # ── Guard 1: File Size ─────────────────────────────────────────
    # Checked before ANY decoding — cheapest check possible
    # Stops large files before they ever touch RAM
    contents = await file.read()
    if len(contents) > MAX_FILE_SIZE:
        raise HTTPException(
            status_code=400,
            detail=f"File too large. Maximum allowed size is 20MB."
        )

    # ── Guard 2: Readable Image ────────────────────────────────────
    # .convert("RGB") is the real test — if this works, the model can handle it
    # Avoids image.verify() which incorrectly rejects many valid real-world images
    # (phone photos with non-standard EXIF, progressive JPEGs, etc.)
    try:
        image = Image.open(io.BytesIO(contents)).convert("RGB")
    except Exception:
        raise HTTPException(
            status_code=400,
            detail="Cannot read image file. Please upload a valid JPG, PNG, or WEBP."
        )

    # ── Guard 3: Pixel Dimensions ──────────────────────────────────
    # Guards against malicious images: small file size but enormous pixel count
    # Example: a 1MB PNG that is 50000×50000 would consume ~7.5GB RAM to decode
    # This check runs AFTER decoding so we have real dimensions
    if image.width * image.height > MAX_PIXELS:
        del image  # free RAM immediately before raising
        raise HTTPException(
            status_code=400,
            detail=f"Image dimensions too large. Maximum is 10000×10000 pixels."
        )

    # ── Preprocess ─────────────────────────────────────────────────
    orig_size    = image.size
    input_tensor = preprocess(image)
    # Convert to RGBA now while image is still in memory.
    # Avoids re-opening from bytes later (~168ms saved per request).
    image_rgba   = image.convert("RGBA")
    del image    # free the large RGB decoded image immediately
                 # image_rgba is kept — it is small (same pixels, +alpha channel)

    # ── Grab a Free Model from the Pool ───────────────────────────
    # If all 4 copies are busy this line waits until one is free
    # Nobody gets rejected — they queue fairly
    model = await model_pool.get()

    try:
        # Run inference in a thread pool so the asyncio event loop
        # stays free to handle I/O (new connections, sending responses)
        # while the CPU-bound model inference runs in the background.
        # torch.inference_mode disables gradient tracking — 5-15% faster.
        def run_inference():
            with torch.inference_mode():
                return model(input_tensor)

        result = await asyncio.to_thread(run_inference)
        del input_tensor  # free tensor now that model is done with it

        mask = result[0][0].squeeze()
        mask = F.interpolate(
            mask.unsqueeze(0).unsqueeze(0),
            size=(orig_size[1], orig_size[0]),
            mode="bilinear",
            align_corners=False
        ).squeeze()

        # Move to CPU only when needed for numpy conversion
        mask        = (mask.cpu() * 255).clamp(0, 255).byte().numpy()
        mask_image  = Image.fromarray(mask)

        # image_rgba already in memory — no re-decode needed
        image_out = image_rgba
        image_out.putalpha(mask_image)

        output_buffer = io.BytesIO()
        # WebP with alpha: ~70% smaller than PNG compress_level=1, similar encode speed.
        # quality=85 is visually lossless for transparency masks at this resolution.
        # method=3 balances encode speed vs compression (0=fastest, 6=smallest).
        image_out.save(output_buffer, format="WEBP", quality=85, method=3)
        output_buffer.seek(0)

        return Response(content=output_buffer.getvalue(), media_type="image/webp")

    finally:
        # ALWAYS return the model copy to the pool
        # This runs even if the request crashes with any Python-level error
        # Pool never permanently loses a slot due to soft crashes
        await model_pool.put(model)