File size: 2,119 Bytes
febe160 | 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 | """
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'")
|