main.py

from __future__ import annotations

import logging
import os
import threading
import time
from contextlib import asynccontextmanager
from io import BytesIO
from pathlib import Path
from typing import Any

import numpy as np
import torch
from fastapi import Body, FastAPI, HTTPException, Request, Response
from PIL import Image, UnidentifiedImageError
from transformers import AutoImageProcessor, AutoModelForDepthEstimation

LOGGER = logging.getLogger("spatialthings.depth_pro")

DEFAULT_HF_MODEL_ID = "apple/DepthPro-hf"
ENDPOINT_MODEL_PATH = Path("/repository")


def _env_bool(name: str, default: bool | None = None) -> bool | None:
    raw = os.getenv(name)
    if raw is None or raw == "":
        return default
    return raw.strip().lower() in {"1", "true", "yes", "y", "on"}


def _select_model_id() -> str:
    configured = os.getenv("DEPTH_PRO_MODEL_ID")
    if configured:
        return configured
    if ENDPOINT_MODEL_PATH.exists():
        return str(ENDPOINT_MODEL_PATH)
    return DEFAULT_HF_MODEL_ID


def _select_device() -> torch.device:
    configured = os.getenv("DEPTH_PRO_DEVICE", "auto").strip().lower()
    if configured == "auto":
        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
    return torch.device(configured)


def _select_dtype(device: torch.device) -> torch.dtype:
    configured = os.getenv("DEPTH_PRO_DTYPE", "auto").strip().lower()
    if configured == "auto":
        return torch.float16 if device.type == "cuda" else torch.float32
    if configured in {"float16", "fp16", "half"}:
        return torch.float16
    if configured in {"bfloat16", "bf16"}:
        return torch.bfloat16
    if configured in {"float32", "fp32"}:
        return torch.float32
    raise ValueError(f"Unsupported DEPTH_PRO_DTYPE={configured!r}")


class ModelNotReadyError(RuntimeError):
    pass


class DepthProModel:
    def __init__(self) -> None:
        self.model_id = _select_model_id()
        self.device = _select_device()
        self.dtype = _select_dtype(self.device)
        self.max_image_bytes = int(os.getenv("MAX_IMAGE_BYTES", str(16 * 1024 * 1024)))
        self.use_fov_model = _env_bool("DEPTH_PRO_USE_FOV_MODEL", None)
        self.processor: Any | None = None
        self.model: Any | None = None
        self.load_lock = threading.Lock()
        self.lock = threading.Lock()

    def load(self) -> None:
        if self.processor is not None and self.model is not None:
            return

        with self.load_lock:
            if self.processor is not None and self.model is not None:
                return

            model_kwargs: dict[str, Any] = {"torch_dtype": self.dtype}
            if self.use_fov_model is not None:
                model_kwargs["use_fov_model"] = self.use_fov_model

            started = time.perf_counter()
            LOGGER.info(
                "Loading Depth Pro model model_id=%s device=%s dtype=%s use_fov_model=%s",
                self.model_id,
                self.device,
                self.dtype,
                self.use_fov_model,
            )
            self.processor = AutoImageProcessor.from_pretrained(self.model_id, use_fast=True)
            self.model = AutoModelForDepthEstimation.from_pretrained(self.model_id, **model_kwargs)
            self.model.to(self.device)
            self.model.eval()
            LOGGER.info("Loaded Depth Pro in %.3fs", time.perf_counter() - started)

    def unload(self) -> None:
        if self.model is not None:
            self.model.to("cpu")
        self.model = None
        self.processor = None
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    def require_loaded(self) -> tuple[Any, Any]:
        if self.processor is None or self.model is None:
            raise ModelNotReadyError("Depth Pro model is not loaded")
        return self.processor, self.model

    def health(self) -> dict[str, Any]:
        return {
            "ok": self.processor is not None and self.model is not None,
            "model_id": self.model_id,
            "device": str(self.device),
            "dtype": str(self.dtype).replace("torch.", ""),
            "scale": "metric_meters",
            "max_image_bytes": self.max_image_bytes,
        }

    def estimate(self, image: Image.Image) -> tuple[np.ndarray, float]:
        self.load()
        processor, model = self.require_loaded()
        width, height = image.size

        with self.lock:
            started = time.perf_counter()
            inputs = processor(images=image, return_tensors="pt")
            inputs = {
                key: value.to(device=self.device)
                if not torch.is_floating_point(value)
                else value.to(device=self.device, dtype=self.dtype)
                for key, value in inputs.items()
            }

            with torch.inference_mode():
                outputs = model(**inputs)

            post_processed = processor.post_process_depth_estimation(
                outputs,
                target_sizes=[(height, width)],
            )
            depth = post_processed[0]["predicted_depth"]
            elapsed = time.perf_counter() - started

        if isinstance(depth, torch.Tensor):
            depth_np = depth.detach().to(dtype=torch.float32).cpu().numpy()
        else:
            depth_np = np.asarray(depth, dtype=np.float32)

        if depth_np.ndim != 2:
            raise RuntimeError(f"Depth output must be 2D, got shape={depth_np.shape}")

        return np.ascontiguousarray(depth_np.astype(np.float32, copy=False)), elapsed


depth_pro = DepthProModel()


@asynccontextmanager
async def lifespan(app: FastAPI):
    if _env_bool("DEPTH_PRO_EAGER_LOAD", True):
        depth_pro.load()
    try:
        yield
    finally:
        depth_pro.unload()


app = FastAPI(title="SpatialThings Depth Pro API", lifespan=lifespan)


@app.get("/")
def root() -> dict[str, Any]:
    return depth_pro.health()


@app.get("/health")
def health() -> dict[str, Any]:
    return depth_pro.health()


@app.post("/estimate-depth")
def estimate_depth(
    request: Request,
    body: bytes = Body(..., media_type="image/jpeg"),
) -> Response:
    content_type = request.headers.get("content-type", "").split(";", maxsplit=1)[0].strip().lower()
    if content_type not in {"image/jpeg", "image/jpg"}:
        raise HTTPException(status_code=415, detail="Content-Type must be image/jpeg")
    if not body:
        raise HTTPException(status_code=400, detail="Missing request body")
    if len(body) > depth_pro.max_image_bytes:
        raise HTTPException(status_code=413, detail=f"Request body too large: {len(body)} bytes")

    try:
        with Image.open(BytesIO(body)) as image:
            rgb = image.convert("RGB")
    except (UnidentifiedImageError, OSError) as exc:
        raise HTTPException(status_code=400, detail=f"Invalid JPEG image: {exc}") from exc

    try:
        depth, elapsed = depth_pro.estimate(rgb)
    except ModelNotReadyError as exc:
        raise HTTPException(status_code=503, detail=str(exc)) from exc
    except RuntimeError as exc:
        LOGGER.exception("Depth estimation failed")
        raise HTTPException(status_code=500, detail=f"Depth estimation failed: {exc}") from exc

    little_endian = np.ascontiguousarray(depth.astype("<f4", copy=False))
    payload = little_endian.tobytes()
    headers = {
        "Content-Length": str(len(payload)),
        "X-Depth-Width": str(depth.shape[1]),
        "X-Depth-Height": str(depth.shape[0]),
        "X-Depth-Scale": "metric_meters",
        "X-Process-Time-Sec": f"{elapsed:.6f}",
    }
    return Response(content=payload, media_type="application/octet-stream", headers=headers)


if __name__ == "__main__":
    import uvicorn

    logging.basicConfig(level=os.getenv("LOG_LEVEL", "INFO"))
    uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", "8000")))