anujg1508 commited on
Commit
3de8d46
1 Parent(s): 15c460b

update handler

Browse files
Files changed (1) hide show
  1. handler.py +2 -1
handler.py CHANGED
@@ -19,9 +19,10 @@ class EndpointHandler():
19
 
20
  def __call__(self, data:Any) -> List[List[Dict[str, float]]]:
21
  print(data)
 
22
  print(device)
23
  with autocast(device.type):
24
- image = self.pipe(prompt=data, guidance_scale=7.5).images[0]
25
 
26
  buffered = BytesIO()
27
  image.save(buffered, format="JPEG")
 
19
 
20
  def __call__(self, data:Any) -> List[List[Dict[str, float]]]:
21
  print(data)
22
+ inputs = data.pop("inputs", data)
23
  print(device)
24
  with autocast(device.type):
25
+ image = self.pipe(inputs, guidance_scale=7.5).images[0]
26
 
27
  buffered = BytesIO()
28
  image.save(buffered, format="JPEG")