Updated handler.py

handler.py

@@ -4,9 +4,7 @@ from transformers import Blip2Processor, Blip2ForConditionalGeneration
 
 from PIL import Image
 from io import BytesIO
-import torch
-import os
-
+import torch, os, base64
 
 
 class EndpointHandler:
@@ -28,11 +26,14 @@ class EndpointHandler:
                 - "caption": A string corresponding to the generated caption.
         """
 
-        inputs = data.pop("inputs", data)
+        inputs = data["inputs"]
         # parameters = data.pop("parameters", {})
         device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
- 
-        raw_images = Image.open(BytesIO(inputs.encode('utf-8')))
+
+        if type(inputs) is str:
+            inputs = inputs.encode('utf-8')
+
+        raw_images = Image.open(BytesIO(inputs))
                                      
         processed_image = self.processor(images=raw_images, return_tensors="pt").to(device)
 
@@ -41,4 +42,7 @@ class EndpointHandler:
         captions = self.processor.decode(out[0], skip_special_tokens=True)
 
         # postprocess the prediction
-        return {"captions": captions}
+        return {"captions": captions}
+    
+
+EndpointHandler()