Update handler.py

handler.py

@@ -3,12 +3,13 @@ from transformers import Blip2Processor, Blip2ForConditionalGeneration
 import io 
 from PIL import Image
 import base64
+import torch 
 
 class EndpointHandler:
     def __init__(self, path=""):
         # load model and processor from path
         self.processor = Blip2Processor.from_pretrained(path)
-        self.model = Blip2ForConditionalGeneration.from_pretrained(path)
+        self.model = Blip2ForConditionalGeneration.from_pretrained(path, torch_dtype=torch.float16)
         self.device = "cuda"
 
         self.model.to(self.device)
@@ -21,7 +22,7 @@ class EndpointHandler:
         image_string = base64.b64decode(data["image"])
         image = Image.open(io.BytesIO(image_string))
 
-        inputs = self.processor(images=image, text=text, return_tensors="pt").to(self.device)
+        inputs = self.processor(images=image, text=text, return_tensors="pt").to(device, torch.float16)
         generated_ids = self.model.generate(**inputs)
         generated_text = self.processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()