Create handler.py

handler.py

@@ -0,0 +1,41 @@
+import base64
+from io import BytesIO
+from typing import Dict, Any
+
+import torch
+from PIL import Image
+from diffusers import StableDiffusionXLPipeline
+
+
+# helper decoder
+def decode_base64_image(image_string):
+    base64_image = base64.b64decode(image_string)
+    buffer = BytesIO(base64_image)
+    return Image.open(buffer)
+
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        self.pipe = StableDiffusionXLPipeline.from_pretrained("/repository/roses",
+            torch_dtype=torch.float16, revision="fp16")
+        self.pipe = self.pipe.to("cuda")
+
+    def __call__(self, data: Any) -> Dict[str, str]:
+        """
+        Return predict value.
+        :param data: A dictionary contains `inputs` and optional `image` field.
+        :return: A dictionary with `image` field contains image in base64.
+        """
+        prompts = data.pop("inputs", None)
+        encoded_image = data.pop("image", None)
+        init_image = None
+        if encoded_image:
+            init_image = decode_base64_image(encoded_image)
+            init_image.thumbnail((768, 768))
+
+        image = self.pipe(prompts, init_image=init_image).images[0]
+        buffered = BytesIO()
+        image.save(buffered, format="png")
+        img_str = base64.b64encode(buffered.getvalue())
+
+        return {"image": img_str.decode()}