Update handler.py

handler.py

@@ -3,14 +3,13 @@ import base64
 
 from PIL import Image
 import torch
-from transformers import CLIPProcessor, CLIPTextModel, CLIPVisionModelWithProjection
+from transformers import CLIPProcessor, CLIPModel
 
 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").to(device)
-        self.image_model = CLIPVisionModelWithProjection.from_pretrained("rbanfield/clip-vit-large-patch14").to(device)
+        self.model = CLIPModel.from_pretrained("rbanfield/clip-vit-large-patch14").to("cpu")
         self.processor = CLIPProcessor.from_pretrained("rbanfield/clip-vit-large-patch14")
 
     def __call__(self, data):
@@ -19,13 +18,15 @@ class EndpointHandler():
         image_input = inputs["image"] if "image" in inputs else None
 
         if text_input:
-            processor = self.processor(text=text_input, return_tensors="pt", padding=True).to(device)
+            processor = self.processor(text=text_input, return_tensors="pt", padding=True)
+            processor.to("cpu")
             with torch.no_grad():
-                return self.text_model(**processor).pooler_output.tolist()
+                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").to(device)
+            processor = self.processor(images=image, return_tensors="pt")
+            processor.to("cpu")
             with torch.no_grad():
-                return self.image_model(**processor).image_embeds.tolist()
+                return self.model.get_image_features(**processor).tolist()
         else:
-            return None
+            return None