Update handler.py

handler.py

@@ -8,39 +8,20 @@ import json
 
 class EndpointHandler():
     def __init__(self, path=""):
-        # Preload all the elements you are going to need at inference.
-        # pseudo:
-        # self.model= load_model(path)
         self.processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
         self.model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base").to("cuda")
-        #是否能够只传handler一个文件？答案是可以的。。。
+        
     def __call__(self, data):
-        """
-        data args:
-            inputs (:obj: `str` | `PIL.Image` | `np.array`)
-            kwargs
-        Return:
-            A :obj:`list` | `dict`: will be serialized and returned
-        """
-        inputs=data.pop("inputs",data)
-        inputs=base64.b64decode(inputs)
-        raw_images = Image.open(BytesIO(inputs))
-        # raw_image = Image.open(image_path).convert('RGB')
-
-        # # conditional image captioning
-        # text = "a photography of"
-        # inputs = self.processor(raw_image, text, return_tensors="pt").to("cuda", torch.float16)
-
-        # out = self.model.generate(**inputs)
-        # print(self.processor.decode(out[0], skip_special_tokens=True))
-        # >>> a photography of a woman and her dog
-
-        # unconditional image captioning
+        info=data['inputs']
+        img=info.pop('image',data)
+        image_bytes=base64.b64decode(img)
+        raw_images = Image.open(BytesIO(image_bytes))
+        
         inputs = self.processor(raw_images, return_tensors="pt").to("cuda")
 
         out = self.model.generate(**inputs)
         # print(self.processor.decode(out[0], skip_special_tokens=True))
-        return json.dumps({'text':self.processor.decode(out[0], skip_special_tokens=True)})
+        return {'text':self.processor.decode(out[0], skip_special_tokens=True)}
     
 if __name__=="__main__":
     my_handler=EndpointHandler(path='.')