Upload 2 files

Upload a handler.py and requirements.txt to run on inference endpoints

handler.py

@@ -0,0 +1,30 @@
+from io import BytesIO
+import base64
+
+from PIL import Image
+import torch
+from transformers import CLIPProcessor, CLIPTextModel, CLIPVisionModelWithProjection
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+class EndpointHandler():
+    def __init__(self, path=""):
+        self.text_model = CLIPTextModel.from_pretrained("rbanfield/clip-vit-large-patch14")
+        self.image_model = CLIPVisionModelWithProjection.from_pretrained("rbanfield/clip-vit-large-patch14")
+        self.processor = CLIPProcessor.from_pretrained("rbanfield/clip-vit-large-patch14")
+
+    def __call__(self, data):
+        text_input = data.pop("text", None)
+        image_input = data.pop("image", None)
+
+        if text_input:
+            processor = self.processor(text=text_input, return_tensors="pt", padding=True)
+            with torch.no_grad():
+                return self.text_model(**processor).pooler_output.tolist()
+        elif image_input:
+            image = Image.open(BytesIO(base64.b64decode(image_input)))
+            processor = self.processor(images=image, return_tensors="pt")
+            with torch.no_grad():
+                return self.image_model(**processor).image_embeds.tolist()
+        else:
+            return None

requirements.txt

@@ -0,0 +1,3 @@
+Pillow
+transformers
+torch