handler.py

"""
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=""):
        # If the repo contains model weights, use them; otherwise load from HF Hub.
        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)
            # last_hidden_state: [1, seq_len, 768] — take CLS token at index 0
            embedding = outputs.last_hidden_state[:, 0, :].squeeze().tolist()

        return embedding