Update handler.py

handler.py

@@ -15,11 +15,13 @@ class EndpointHandler():
         self.model.eval()
         self.max_length = 16
         self.num_beams = 4
-
-    def __call__(self, image_data: str) -> dict:
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+    # def __call__(self, image_data: str) -> dict:
         try:
+            image_bytes = data.get("inputs", None)
+            
             # Convert base64 encoded image string to a PIL Image
-            raw_image = Image.open(BytesIO(base64.b64decode(image_data))).convert("RGB")
+            raw_image = Image.open(BytesIO(image_bytes))
 
             # Ensure the image is in RGB mode
             if raw_image.mode != "RGB":
@@ -32,7 +34,7 @@ class EndpointHandler():
             gen_kwargs = {"max_length": self.max_length, "num_beams": self.num_beams}
             output_ids = self.model.generate(pixel_values, **gen_kwargs)
 
-            caption = self.processor.batch_decode(output_ids, skip_special_tokens=True)[0]
+            caption = self.processor.batch_decode(output_ids[0], skip_special_tokens=True).strip()
 
             return {"caption": caption}
         except Exception as e: