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