Updated with new code

handler.py

@@ -21,8 +21,8 @@ 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", torch_dtype=torch.float16, device_map="auto")
+        self.processor = Blip2Processor.from_pretrained("blip2/sharded")
+        self.model = Blip2ForConditionalGeneration.from_pretrained("blip2/sharded", device_map = "auto", load_in_8bit = True)
         print(" $$$$ model loaded $$$$")
         self.model.eval()
         self.model = self.model.to(device)
@@ -66,9 +66,13 @@ 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", torch.float16)
+        generated_ids = self.processor(raw_image, return_tensors="pt").to("cuda", torch.float16)
+        generated_text = self.processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
+        print("@@@@@@ generated_text @@@@@@@")
+        print(generated_text)
         
-        out = self.model.generate(**inputs)
-        captions = processor.decode(out[0], skip_special_tokens=True)
         
-        return {"captions": captions}
+        # out = self.model.generate(**inputs)
+        # captions = processor.decode(out[0], skip_special_tokens=True)
+        
+        return {"captions": generated_text}