Update handler.py

handler.py

@@ -4,7 +4,10 @@ from typing import Dict, List, Any
 # import torch
 from datetime import datetime
 
+import torch
 
+import logging
+logging.basicConfig(format='%(levelname)s:%(message)s', level=logging.DEBUG)
 
 
 import requests
@@ -19,6 +22,12 @@ class EndpointHandler():
         self.processor = Blip2Processor.from_pretrained(path)
         self.model = Blip2ForConditionalGeneration.from_pretrained(path, device_map="auto")
 
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        self.model.to(self.device)
+
+        logging.info('Model moved to device-' + self.device)
+        
         # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         # self.model.eval()
         # self.model.to(device=device, dtype=self.torch_dtype)
@@ -71,11 +80,14 @@ class EndpointHandler():
 
         raw_image = Image.open(requests.get(img_url, stream=True).raw).convert('RGB')
         
-        question = "how many dogs are in the picture?"
-        inputs = self.processor(raw_image, question, return_tensors="pt").to("cuda")
+        # question = "how many dogs are in the picture?"
+        # inputs = self.processor(raw_image, question, return_tensors="pt").to("cuda")
+
+        inputs = self.processor(raw_image, return_tensors="pt").to("cuda")
         
         out = self.model.generate(**inputs)
-        output_text = self.processor.decode(out[0], skip_special_tokens=True)
+
+        generated_text = self.processor.batch_decode(out, skip_special_tokens=True)[0].strip()
 
         current = datetime.now()
 
@@ -100,4 +112,4 @@ class EndpointHandler():
         # new_tokens = output_ids[0, len(input_ids[0]) :]
         # output_text = self.tokenizer.decode(new_tokens, skip_special_tokens=True)
         
-        return [{"gen_text":output_text, "time_elapsed": str(current-now)}]
+        return [{"gen_text":generated_text, "time_elapsed": str(current-now)}]