Update handler.py

handler.py

@@ -12,36 +12,13 @@ class EndpointHandler():
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         self.model_base = "Salesforce/blip2-opt-2.7b"
         self.model_name = "sooh-j/blip2-vizwizqa"
-        # self.base_model = Blip2ForConditionalGeneration.from_pretrained(self.model_base, load_in_8bit=True)
         # self.pipe = Blip2ForConditionalGeneration.from_pretrained(self.model_base, load_in_8bit=True, torch_dtype=torch.float16)
 
-        quantization_config = BitsAndBytesConfig(load_in_8bit=True)
-
-        # self.processor = Blip2Processor.from_pretrained(self.model_name)
         self.processor = AutoProcessor.from_pretrained(self.model_name)
-        # self.model = BlipForQuestionAnswering.from_pretrained(self.model_name,
-        # self.model = AutoModelForCausalLM.from_pretrained(self.model_name,
         self.model = Blip2ForConditionalGeneration.from_pretrained(self.model_name,
                                                               device_map="auto", 
-                                                              # load_in_8bit=True,
-                                                              # quantization_config=quantization_config,
                                                              ).to(self.device)
-        # self.model = PeftModel.from_pretrained(self.model_name, self.base_model_name).to(self.device)
 
-# inputs = data.get("inputs")
-#         imageBase64 = inputs.get("image")
-#         # imageURL = inputs.get("image")
-#         text = inputs.get("text")
-#         # print(imageURL)
-#         # print(text)
-#         # image = Image.open(requests.get(imageBase64, stream=True).raw)
-        
-#         image = Image.open(BytesIO(base64.b64decode(imageBase64.split(",")[1].encode())))
-#         inputs = self.processor(text=text, images=image, return_tensors="pt", padding=True)
-#         outputs = self.model(**inputs)
-#         embeddings = outputs.image_embeds.detach().numpy().flatten().tolist()
-#         return { "embeddings": embeddings }
-      
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
         data args:
@@ -89,15 +66,12 @@ class EndpointHandler():
         # processed_images = self.processor(images=raw_images, return_tensors="pt")
         # processed_images["pixel_values"] = processed_images["pixel_values"].to(device)
         # processed_images = {**processed_images, **parameters}
-        
-        # with torch.no_grad():
-        #     out = self.model.generate(**processed_images)
-        # captions = self.processor.batch_decode(out, skip_special_tokens=True)
+
 ####
 
         
         prompt = f"Question: {question}, Answer:"
-        processed = self.processor(images=image, text=prompt, return_tensors="pt").to(self.device)#, torch.float16)
+        processed = self.processor(images=image, text=prompt, return_tensors="pt").to(self.device)
 
         # answer = self._generate_answer(
         #     model_path, prompt, image,