| """ |
| DINOv2 image embedding handler for HF Inference Endpoints. |
| Accepts base64-encoded images via {"inputs": "<base64_string>"}. |
| Returns a 768-dim CLS token embedding as a flat list of floats. |
| """ |
| import base64 |
| import os |
| from io import BytesIO |
|
|
| import torch |
| from PIL import Image |
| from transformers import AutoImageProcessor, AutoModel |
|
|
|
|
| class EndpointHandler: |
| def __init__(self, path=""): |
| |
| has_weights = path and os.path.exists(os.path.join(path, "config.json")) |
| model_id = path if has_weights else "facebook/dinov2-base" |
| self.processor = AutoImageProcessor.from_pretrained(model_id) |
| self.model = AutoModel.from_pretrained(model_id) |
| self.model.eval() |
|
|
| def __call__(self, data: dict) -> list: |
| """ |
| data: {"inputs": "<base64_string>"} |
| Returns: [float x 768] — DINOv2 CLS token embedding |
| """ |
| raw = data.get("inputs", "") |
|
|
| if isinstance(raw, str): |
| img_bytes = base64.b64decode(raw) |
| elif isinstance(raw, bytes): |
| img_bytes = raw |
| else: |
| raise ValueError(f"Unexpected input type: {type(raw)}") |
|
|
| image = Image.open(BytesIO(img_bytes)).convert("RGB") |
|
|
| with torch.no_grad(): |
| inputs = self.processor(images=image, return_tensors="pt") |
| outputs = self.model(**inputs) |
| |
| embedding = outputs.last_hidden_state[:, 0, :].squeeze().tolist() |
|
|
| return embedding |
|
|
