handler.py

"""
Custom Hugging Face Inference Endpoint handler for SigLIP2 image embeddings.

There is no ready-made "embeddings" task handler for google/siglip2-so400m-patch14-384
on HF Inference Endpoints (the model ships for zero-shot classification, not raw
embedding extraction), so this handler exposes model.get_image_features() directly.

Deploy: create a new Inference Endpoint from the google/siglip2-so400m-patch14-384
repo, upload this file (and requirements.txt) as the custom handler, select a GPU
instance, deploy, then copy the resulting endpoint URL into the HF_EMBEDDING_ENDPOINT_URL
Supabase secret.

Request body: raw image bytes (any content-type recognized by PIL: image/jpeg, image/png, ...)
Response body: {"embedding": [1152 floats]}
"""
import base64
import io

import torch
from PIL import Image
from transformers import AutoModel, AutoProcessor

MODEL_ID = "google/siglip2-so400m-patch14-384"


class EndpointHandler:
    def __init__(self, path=""):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model = AutoModel.from_pretrained(path or MODEL_ID).to(self.device).eval()
        self.processor = AutoProcessor.from_pretrained(path or MODEL_ID)

    def __call__(self, data):
        image_bytes = self._extract_image_bytes(data)
        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")

        inputs = self.processor(images=image, return_tensors="pt").to(self.device)

        with torch.no_grad():
            features = self.model.get_image_features(**inputs)

        embedding = features[0].cpu().to(torch.float32).tolist()
        return {"embedding": embedding}

    @staticmethod
    def _extract_image_bytes(data: dict) -> bytes:
        # HF Inference Endpoints pass raw bytes under "inputs" when the request
        # content-type is an image type; some clients instead send base64 text.
        raw = data.get("inputs", data)
        if isinstance(raw, bytes):
            return raw
        if isinstance(raw, str):
            return base64.b64decode(raw)
        raise ValueError("Expected raw image bytes or base64 string under 'inputs'")