Upload handler.py (#5)

- Upload handler.py (cb666b7a0c24f5a2ee8497a4176d94d869aaa20f)

handler.py

@@ -1,3 +1,4 @@
+import base64
 from typing import Any, Dict
 from transformers import BlipProcessor, BlipForConditionalGeneration
 from PIL import Image
@@ -31,20 +32,15 @@ class EndpointHandler():
         """
         logger.debug(f"Received data keys: {data.keys()}")
 
-        images = data.pop("inputs", data)
-        text = data.get("text", "")
-        parameters = data.pop("parameters", {})
+        image_base64 = data["inputs"].get("image")
+        image_data = base64.b64decode(image_base64)
+
+        # Convert image data to PIL Image
+        images = Image.open(BytesIO(image_data))
 
-        try:
-            # Ensure inputs is a list of image bytes, even if only a single image is provided
-            if isinstance(images, bytes):  # Single image as bytes
-                raw_images = [Image.open(BytesIO(images))]
-            elif isinstance(images, list):  # Multiple images as list of bytes
-                raw_images = [Image.open(BytesIO(_img)) for _img in images if isinstance(_img, bytes)]
-            else:
-                raise ValueError("Invalid image input format. Expected bytes or list of bytes.")
-        except Exception as e:
-            return {"error": f"Error fetching or processing image: {str(e)}"}
+        # Optional text input
+        text = data["inputs"].get("text", "")
+        parameters = data.pop("parameters", {})
 
         processed_image = self.processor(images=images, text=text, return_tensors="pt")
         processed_image["pixel_values"] = processed_image["pixel_values"].to(device)