Moved to 16 bit precision

handler.py

@@ -22,7 +22,7 @@ class EndpointHandler():
 
         print(" $$$$ Model Loading $$$$")
         self.processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b")
-        self.model = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-opt-2.7b", device_map="auto")
+        self.model = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-opt-2.7b", torch_dtype=torch.float16, device_map="auto")
         print(" $$$$ model loaded $$$$")
         self.model.eval()
         self.model = self.model.to(device)
@@ -66,7 +66,7 @@ class EndpointHandler():
         # img_url = 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/demo.jpg' 
         # raw_image = Image.open(requests.get(img_url, stream=True).raw).convert('RGB')
         
-        inputs = self.processor(raw_image, prompt, return_tensors="pt").to("cuda")
+        inputs = self.processor(raw_image, prompt, return_tensors="pt").to("cuda", torch.float16)
         
         out = self.model.generate(**inputs)
         captions = processor.decode(out[0], skip_special_tokens=True)