Upload handler.py

handler.py

@@ -0,0 +1,33 @@
+from typing import Dict, Any
+from PIL import Image
+from transformers import CLIPProcessor, CLIPModel
+import torch
+import base64
+from io import BytesIO
+
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        self.model = CLIPModel.from_pretrained(path)
+        self.processor = CLIPProcessor.from_pretrained(path)
+        self.model.eval()
+
+    def _to_image(self, x) -> Image.Image:
+        if isinstance(x, Image.Image):
+            return x.convert("RGB")
+        if isinstance(x, (bytes, bytearray)):
+            return Image.open(BytesIO(x)).convert("RGB")
+        if isinstance(x, str):
+            return Image.open(BytesIO(base64.b64decode(x))).convert("RGB")
+        if isinstance(x, dict) and "image" in x:
+            return self._to_image(x["image"])
+        raise ValueError("Unsupported image input")
+
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        inputs = data.get("inputs", data)
+        image = self._to_image(inputs)
+        proc = self.processor(images=image, return_tensors="pt")
+        with torch.no_grad():
+            feats = self.model.get_image_features(**proc)
+        feats = feats / feats.norm(p=2, dim=-1, keepdim=True)  # L2 normalize
+        return {"embedding": feats[0].tolist(), "dim": int(feats.shape[-1])}