Files changed (1) hide show
  1. handler.py +3 -8
handler.py CHANGED
@@ -19,7 +19,7 @@ class EndpointHandler():
19
  self.pipe = self.pipe.to(device)
20
 
21
 
22
- def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
23
  """
24
  Args:
25
  data (:obj:):
@@ -33,10 +33,5 @@ class EndpointHandler():
33
  with autocast(device.type):
34
  image = self.pipe(inputs, guidance_scale=7.5)["sample"][0]
35
 
36
- # encode image as base 64
37
- buffered = BytesIO()
38
- image.save(buffered, format="JPEG")
39
- img_str = base64.b64encode(buffered.getvalue())
40
-
41
- # postprocess the prediction
42
- return {"image": img_str.decode()}
 
19
  self.pipe = self.pipe.to(device)
20
 
21
 
22
+ def __call__(self, data: Any) -> "PIL.Image":
23
  """
24
  Args:
25
  data (:obj:):
 
33
  with autocast(device.type):
34
  image = self.pipe(inputs, guidance_scale=7.5)["sample"][0]
35
 
36
+ # encoding image as base 64 is done by the default toolkit
37
+ return image