Add handler

handler.py

@@ -0,0 +1,28 @@
+from diffusers import AutoPipelineForImage2Image
+import torch
+from typing import Dict, Any
+from PIL import Image
+from io import BytesIO
+import base64
+
+
+class EndpointHandler():
+
+    def __init__(self, path="."):
+        if torch.cuda.is_available():
+            device = "cuda"
+        else:
+            device = "cpu"
+        self._pipe = AutoPipelineForImage2Image.from_pretrained(path, torch_dtype=torch.float32) #.to(device)
+
+    def __call__(self, data: Dict[str, Any]) -> list[Dict[str, Any]]:
+        inputs = data.pop("inputs", data)
+
+        params = {"prompt": inputs.get("prompt", ""),
+                  "image": Image.open(BytesIO(base64.b64decode(inputs['image']))),
+                  "strength": inputs.get("strength", 0.3),
+                  "guidance_scale": inputs.get("guidance_scale", 10),
+                  "height": 768,
+                  "width": 768}
+
+        return self._pipe(**params).images[0]