Upload handler.py

fixed image load issue

handler.py

@@ -26,7 +26,10 @@ class EndpointHandler():
         inputs = data.pop("inputs", data)
         parameters = data.pop("parameters", {})
 
-        raw_images = [Image.open(BytesIO(_img)) for _img in inputs]
+        if isinstance(inputs, bytes):  # If a single image is passed
+            raw_images = [Image.open(BytesIO(inputs))]
+        else:  # If already a list of images
+            raw_images = [Image.open(BytesIO(_img)) for _img in inputs]
 
         processed_image = self.processor(images=raw_images, return_tensors="pt")
         processed_image["pixel_values"] = processed_image["pixel_values"].to(device)
@@ -38,4 +41,6 @@ class EndpointHandler():
             )
         description = self.processor.batch_decode(out, skip_special_tokens=True)
 
-        return {"description": description}
+        return {"description": description}
+
+handler = EndpointHandler()