Upload handler.py

handler.py

@@ -0,0 +1,46 @@
+"""
+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