add custom handler

Files changed (2) hide show

__pycache__/handler.cpython-38.pyc CHANGED Viewed

Binary files a/__pycache__/handler.cpython-38.pyc and b/__pycache__/handler.cpython-38.pyc differ

handler.py CHANGED Viewed

@@ -8,7 +8,7 @@ class EndpointHandler():
     def __call__(self, data: Dict[str, Any]) -> bytes:
         # get inputs
-        prompt = data.pop("prompt", data)
         # run normal prediction
         images = self.sd.text_to_image(prompt, batch_size=1)

     def __call__(self, data: Dict[str, Any]) -> bytes:
         # get inputs
+        prompt = data.pop("inputs", data)
         # run normal prediction
         images = self.sd.text_to_image(prompt, batch_size=1)