File size: 3,702 Bytes
c96ba94 | 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 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | # Custom inference handler for the palatshq DINOv2 image-embedding endpoint.
#
# This file is the source of truth for the Hugging Face model repo
# (palatshq/dinov2-large). It is NOT imported by the API build; it lives here so
# the deployed code is version-controlled alongside its client
# (apps/api/src/lib/dino-embedder.ts). Push it to the HF repo's `main` branch and
# redeploy the endpoint — see docs/integrations/huggingface.md.
#
# DINOv2 is image-only (no text tower). We return the CLS-token embedding
# (Dinov2 pooler_output), L2-normalized so cosine similarity is a plain dot
# product — matching how the CLIP endpoint returns normalized vectors and what
# the benchmark's cosine matcher expects.
#
# Contract (mirrors dino-embedder.ts). Single image:
# request: { "inputs": { "image": "<base64 or data-url>" } }
# response: { "image_embedding": [float, ...] }
# Batch (one forward over many images, for backfills):
# request: { "inputs": { "images": ["<base64>", ...] } }
# response: { "image_embeddings": [[float, ...], ...] } # aligned with input order
import base64
import binascii
import io
from typing import Any, Dict, List
import torch
from PIL import Image
from transformers import AutoImageProcessor, AutoModel
class EndpointHandler:
def __init__(self, path: str = "") -> None:
self.device = "cuda" if torch.cuda.is_available() else "cpu"
self.processor = AutoImageProcessor.from_pretrained(path)
self.model = AutoModel.from_pretrained(path).to(self.device)
self.model.eval()
@staticmethod
def _load_image(image_value: str) -> Image.Image:
if image_value.startswith("data:"):
image_value = image_value.split(",", 1)[1]
try:
raw = base64.b64decode(image_value, validate=True)
except (binascii.Error, ValueError) as error:
raise ValueError("`image` must be base64 or a base64 data-url") from error
return Image.open(io.BytesIO(raw)).convert("RGB")
def _embed(self, images: List[Image.Image]) -> List[List[float]]:
processed = self.processor(images=images, return_tensors="pt").to(self.device)
# fp16 autocast on GPU roughly halves the vision-encoder forward cost
# with no meaningful hit to embedding quality; CPU stays fp32 (half is
# slow/unsupported for some ops there).
if self.device == "cuda":
with torch.autocast("cuda", dtype=torch.float16):
outputs = self.model(**processed)
else:
outputs = self.model(**processed)
# pooler_output is the CLS token after the final layernorm — the standard
# global descriptor for DINOv2 retrieval.
embeddings = outputs.pooler_output.float()
embeddings = embeddings / embeddings.norm(dim=-1, keepdim=True)
return embeddings.cpu().tolist()
@torch.inference_mode()
def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
inputs = data.get("inputs", data)
if not isinstance(inputs, dict):
raise ValueError("Request must contain `inputs`")
raw_images = inputs.get("images")
if raw_images is not None:
if not isinstance(raw_images, list) or len(raw_images) == 0:
raise ValueError("`inputs.images` must be a non-empty list")
images = [self._load_image(value) for value in raw_images]
return {"image_embeddings": self._embed(images)}
if inputs.get("image"):
image = self._load_image(inputs["image"])
return {"image_embedding": self._embed([image])[0]}
raise ValueError("Request must contain `inputs.image` or `inputs.images`")
|