Update handler.py with new logic

handler.py

@@ -99,7 +99,7 @@ class EndpointHandler:
             logger.error(f"Error generating embedding: {e}", exc_info=True)
             raise ValueError(f"Failed to generate embedding: {str(e)}")
 
-    def _parse_input(self, data: Union[Dict[str, Any], bytes]) -> Image.Image:
+    def _parse_input(self, data: Union[Dict[str, Any], bytes, Image.Image]) -> Image.Image:
         """
         Parse input data into PIL Image.
 
@@ -118,25 +118,48 @@ class EndpointHandler:
             ValueError: If image format is invalid
         """
         try:
-            # Case 1: Binary bytes directly
-            if isinstance(data, bytes):
+            # Case 0: Already a PIL Image
+            if isinstance(data, Image.Image):
+                return data.convert("RGB")
+
+            # Case 1: Raw binary bytes directly
+            if isinstance(data, (bytes, bytearray)):
                 return Image.open(io.BytesIO(data)).convert("RGB")
 
-            # Case 2: Dict with "inputs" key
+            # Case 2: Dict with possible variants
             if isinstance(data, dict):
-                inputs = data.get("inputs")
+                # Many endpoints pass {"inputs": <something>}
+                inputs = data.get("inputs", data)
+
+                # 2a: Inputs is already a PIL image
+                if isinstance(inputs, Image.Image):
+                    return inputs.convert("RGB")
 
-                if inputs is None:
-                    raise ValueError("Missing 'inputs' key in request data")
+                # 2b: Raw bytes
+                if isinstance(inputs, (bytes, bytearray)):
+                    return Image.open(io.BytesIO(inputs)).convert("RGB")
 
-                # Case 2a: Base64 string
+                # 2c: Base64 string (plain or data URL)
                 if isinstance(inputs, str):
-                    image_bytes = base64.b64decode(inputs)
+                    b64_str = inputs
+                    if inputs.startswith("data:"):
+                        b64_str = inputs.split(",", 1)[1]
+                    image_bytes = base64.b64decode(b64_str)
                     return Image.open(io.BytesIO(image_bytes)).convert("RGB")
 
-                # Case 2b: Binary bytes
-                if isinstance(inputs, bytes):
-                    return Image.open(io.BytesIO(inputs)).convert("RGB")
+                # 2d: Nested dict like {"image": <...>}
+                if isinstance(inputs, dict) and "image" in inputs:
+                    inner = inputs["image"]
+                    if isinstance(inner, Image.Image):
+                        return inner.convert("RGB")
+                    if isinstance(inner, (bytes, bytearray)):
+                        return Image.open(io.BytesIO(inner)).convert("RGB")
+                    if isinstance(inner, str):
+                        b64_str = inner
+                        if inner.startswith("data:"):
+                            b64_str = inner.split(",", 1)[1]
+                        image_bytes = base64.b64decode(b64_str)
+                        return Image.open(io.BytesIO(image_bytes)).convert("RGB")
 
                 raise ValueError(f"Unsupported inputs type: {type(inputs)}")