File size: 2,935 Bytes
ef6797e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | # handler.py — facebook/dinov2-large on Hugging Face Inference Endpoints
#
# Buildly's visual place-recognition embedder: one image -> one 1024-d vector.
# Pooling = mean over patch tokens (AnyLoc-style), chosen for viewpoint
# robustness on bare / near-identical rooms. This choice is PERMANENT:
# changing it silently invalidates every stored room-scan vector.
import base64
import io
from typing import Any, Dict
import torch
from PIL import Image, UnidentifiedImageError
from transformers import AutoImageProcessor, AutoModel
MODEL_ID = "facebook/dinov2-large"
class EndpointHandler:
def __init__(self, path: str = MODEL_ID):
# Eager load at startup. During this window (cold start / scale-from-
# zero) the platform answers 503 on its own; Buildly maps that to
# EmbeddingUnavailable and skips the match without corrupting state.
self.processor = AutoImageProcessor.from_pretrained(path)
self.model = AutoModel.from_pretrained(path).eval()
self.device = "cuda" if torch.cuda.is_available() else "cpu"
self.model.to(self.device)
def _to_image(self, inp: Any) -> Image.Image:
"""Tolerant of every shape Buildly's probe or HF may hand us:
PIL.Image | raw bytes | base64 str | {"image": <b64|bytes>}.
Anything unparseable raises -> the endpoint returns 4xx (NOT 503),
so Buildly's format probe falls through to the next wire shape
instead of treating it as the model being unavailable.
"""
if isinstance(inp, Image.Image):
return inp.convert("RGB")
if isinstance(inp, dict): # {"image": ...}
inp = inp.get("image") or inp.get("inputs")
if isinstance(inp, str): # base64 jpeg
try:
inp = base64.b64decode(inp)
except Exception as exc:
raise ValueError(f"input string is not valid base64: {exc}")
if isinstance(inp, (bytes, bytearray)):
try:
return Image.open(io.BytesIO(inp)).convert("RGB")
except (UnidentifiedImageError, OSError) as exc:
raise ValueError(f"bytes are not a decodable image: {exc}")
raise ValueError(f"unsupported input type: {type(inp).__name__}")
@torch.no_grad()
def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
payload = data.get("inputs", data) if isinstance(data, dict) else data
image = self._to_image(payload)
batch = self.processor(images=image, return_tensors="pt").to(self.device)
tokens = self.model(**batch).last_hidden_state # [1, 257, 1024]
vec = tokens[:, 1:, :].mean(dim=1)[0] # mean-pool patches -> [1024]
# CLS alternative (simpler, less robust for VPR): vec = tokens[:, 0, :][0]
return {"embedding": vec.float().cpu().tolist(), "dim": int(vec.shape[-1])} |