Update handler.py

handler.py

@@ -13,20 +13,24 @@ class EndpointHandler():
         self.processor = CLIPProcessor.from_pretrained("rbanfield/clip-vit-large-patch14")
 
     def __call__(self, data):
-        inputs = data.pop("inputs", None)
-        text_input = inputs["text"] if "text" in inputs else None
-        image_input = inputs["image"] if "image" in inputs else None
 
+        text_input = None
+        if isinstance(data, dict):
+            inputs = data.pop("inputs", None)
+            text_input = inputs.get('text',None)
+            image_data = BytesIO(base64.b64decode(inputs['image'])) if 'image' in inputs else None
+        else:
+            # assuming its an image sent via binary
+            image_data = BytesIO(data)
+        
         if text_input:
-            processor = self.processor(text=text_input, return_tensors="pt", padding=True)
-            processor.to("cpu")
+            processor = self.processor(text=text_input, return_tensors="pt", padding=True).to(device)
             with torch.no_grad():
-                return self.model.get_text_features(**processor).tolist()
-        elif image_input:
-            image = Image.open(BytesIO(base64.b64decode(image_input)))
-            processor = self.processor(images=image, return_tensors="pt")
-            processor.to("cpu")
+                return {"embeddings": self.model.get_text_features(**processor).tolist()}
+        elif image_data:
+            image = Image.open(image_data)
+            processor = self.processor(images=image, return_tensors="pt").to(device)
             with torch.no_grad():
-                return self.model.get_image_features(**processor).tolist()
+                return {"embeddings": self.model.get_image_features(**processor).tolist()}
         else:
-            return None
+            return {"embeddings": None}