File size: 2,012 Bytes
b65ecab a951e8d b65ecab a951e8d b65ecab a951e8d b65ecab a951e8d b65ecab a951e8d | 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 55 56 57 58 59 60 61 | from transformers import AutoProcessor, AutoModel
from PIL import Image
import torch
import requests
import io
import json
class EndpointHandler:
def __init__(self, path=""):
model_name_or_path = "oliverlabs/siglip256"
self.processor = AutoProcessor.from_pretrained(model_name_or_path)
self.model = AutoModel.from_pretrained(model_name_or_path)
self.model.eval()
def _load_image(self, image_input):
"""Load image from URL or bytes"""
if isinstance(image_input, str): # URL
response = requests.get(image_input)
image = Image.open(io.BytesIO(response.content)).convert("RGB")
elif isinstance(image_input, bytes): # raw bytes
image = Image.open(io.BytesIO(image_input)).convert("RGB")
else:
raise ValueError("Unsupported image input format")
return image
def __call__(self, data):
"""
Hugging Face Inference Endpoint calls this method with JSON input
"""
image_input = data.get("image")
texts = data.get("texts", [])
if not image_input or not texts:
return {"error": "Missing image or texts in payload."}
image = self._load_image(image_input)
inputs = self.processor(text=texts, images=image, return_tensors="pt", padding=True)
with torch.no_grad():
outputs = self.model(**inputs)
image_emb = outputs.image_embeds[0].tolist()
text_embs = [emb.tolist() for emb in outputs.text_embeds]
return {
"image_embedding": image_emb,
"text_embeddings": text_embs,
"num_texts": len(texts)
}
if __name__ == "__main__":
handler = EndpointHandler()
test_payload = {
"image": "http://images.cocodataset.org/val2017/000000039769.jpg",
"texts": ["a photo of 2 cats", "a photo of 2 dogs"]
}
result = handler(test_payload)
print(json.dumps(result, indent=2)[:1000] + "\n... (truncated)")
|