Update handler.py

handler.py

@@ -92,13 +92,15 @@ class EndpointHandler():
         #       }
         #     })
         inputs = data.pop("inputs", data)
-        try:
-            imageBase64 = inputs["image"]
-            image = Image.open(BytesIO(base64.b64decode(imageBase64.split(",")[1].encode())))
+        parameters = data.pop("parameters", {})
 
-        except:
-            image_url = inputs['image']
-            image = Image.open(requests.get(image_url, stream=True).raw).convert('RGB')
+        # try:
+        #     imageBase64 = inputs["image"]
+        #     image = Image.open(BytesIO(base64.b64decode(imageBase64.split(",")[1].encode())))
+
+        # except:
+        #     image_url = inputs['image']
+        #     image = Image.open(requests.get(image_url, stream=True).raw).convert('RGB')
             
         question = inputs["question"]
         
@@ -107,14 +109,35 @@ class EndpointHandler():
 
         # image = Image.open(requests.get(imageBase64, stream=True).raw)
         # image = Image.open(requests.get(image_url, stream=True).raw).convert('RGB')
+#### https://huggingface.co/SlowPacer/witron-image-captioning/blob/main/handler.py
+        inputs = data.pop("inputs", data)
+        parameters = data.pop("parameters", {})
+
+        if isinstance(inputs, Image.Image):
+            image = [inputs]
+        else:
+            inputs = isinstance(inputs, str) and [inputs] or inputs
+            image = [Image.open(BytesIO(base64.b64decode(_img))) for _img in inputs]
+                                     
+        # 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)#, torch.float16)
 
         # answer = self._generate_answer(
         #     model_path, prompt, image, 
         # )
-        out = self.model.generate(**processed)
+    
+        with torch.no_grad():
+            out = self.model.generate(**processed)
         
         result = {}
         text_output = self.processor.decode(out[0], skip_special_tokens=True)