Updated handler.py

handler.py

@@ -4,7 +4,7 @@ from transformers import Blip2Processor, Blip2ForConditionalGeneration
 
 from PIL import Image
 from io import BytesIO
-import torch, os, base64
+import torch, re, base64
 
 
 class EndpointHandler:
@@ -25,13 +25,11 @@ class EndpointHandler:
             A :obj:`dict`:. The object returned should be a dict of one list like {"captions": ["A hugging face at the office"]} containing :
                 - "caption": A string corresponding to the generated caption.
         """
-
-        inputs = data["inputs"]
         # parameters = data.pop("parameters", {})
         device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-        if type(inputs) is str:
-            inputs = base64.b64encode(bytes(inputs, 'utf-8'))
+
+        inputs = base64.b64decode(re.sub('^data:image/.+;base64,', '', data['inputs']))
 
         raw_images = Image.open(BytesIO(inputs))