Update handler.py

handler.py

@@ -1,5 +1,5 @@
 from PIL import Image
-from typing import Dict, Any, List
+from typing import Dict, Any
 import torch
 import base64
 from io import BytesIO
@@ -17,15 +17,14 @@ class EndpointHandler():
         self.max_length = 16
         self.num_beams = 4
         
-    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
-    # def __call__(self, image_data: str) -> dict:
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
         try:
             image_bytes = data.get("inputs", None)
             
             # Convert base64 encoded image string to a PIL Image
             raw_image = Image.open(BytesIO(image_bytes))
 
-            # Ensure the image is in RGB mode
+            # Ensure the image is in RGB mode (if necessary)
             if raw_image.mode != "RGB":
                 raw_image = raw_image.convert(mode="RGB")
 
@@ -40,7 +39,6 @@ class EndpointHandler():
 
             return {"caption": caption}
         except Exception as e:
+            # Log the error for better tracking
             print(f"Error during processing: {str(e)}")
             return {"caption": "", "error": str(e)}
-            
-