Update API

handler.py

@@ -19,7 +19,7 @@ class EndpointHandler():
         self.pipe = self.pipe.to(device)
 
 
-    def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
+    def __call__(self, data: Any) -> "PIL.Image":
         """
         Args:
             data (:obj:):
@@ -33,10 +33,5 @@ class EndpointHandler():
         with autocast(device.type):
             image = self.pipe(inputs, guidance_scale=7.5)["sample"][0]  
             
-        # encode image as base 64
-        buffered = BytesIO()
-        image.save(buffered, format="JPEG")
-        img_str = base64.b64encode(buffered.getvalue())
-
-        # postprocess the prediction
-        return {"image": img_str.decode()}
+        # encoding image as base 64 is done by the default toolkit
+        return image