checking-model-1

anujg1508 commited on Sep 9, 2023

Commit

3de8d46

1 Parent(s): 15c460b

update handler

Files changed (1) hide show

handler.py CHANGED Viewed

@@ -19,9 +19,10 @@ class EndpointHandler():
   def __call__(self, data:Any) -> List[List[Dict[str, float]]]:
     print(data)
     print(device)
     with autocast(device.type):
-        image = self.pipe(prompt=data, guidance_scale=7.5).images[0]
     buffered = BytesIO()
     image.save(buffered, format="JPEG")

   def __call__(self, data:Any) -> List[List[Dict[str, float]]]:
     print(data)
+    inputs = data.pop("inputs", data)
     print(device)
     with autocast(device.type):
+        image = self.pipe(inputs, guidance_scale=7.5).images[0]
     buffered = BytesIO()
     image.save(buffered, format="JPEG")