handler.py

handler.py

@@ -0,0 +1,50 @@
+from typing import  Dict, List, Any
+import torch
+from diffusers import DPMSolverMultistepScheduler, StableDiffusionInpaintPipeline, AutoPipelineForInpainting, AutoPipelineForImage2Image, StableDiffusionXLImg2ImgPipeline
+from PIL import Image
+import base64
+from io import BytesIO
+
+
+# set device
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+if device.type != 'cuda':
+    raise ValueError("need to run on GPU")
+
+class EndpointHandler():
+    def __init__(self, path=""):
+
+        #self.fast_pipe = AutoPipelineForInpainting.from_pretrained("diffusers/stable-diffusion-xl-1.0-inpainting-0.1", torch_dtype=torch.float16, variant="fp16").to("cuda")
+        #self.generator = torch.Generator(device="cuda").manual_seed(0)
+
+
+        self.smooth_pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
+          "stabilityai/stable-diffusion-xl-refiner-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
+        )
+        self.smooth_pipe.to("cuda")
+        
+
+
+    def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
+        """
+        :param data: A dictionary contains `inputs` and optional `image` field.
+        :return: A dictionary with `image` field contains image in base64.
+        """
+        encoded_image = data.pop("image", None)
+        
+        prompt = data.pop("prompt", "")
+
+
+        if encoded_image is not None:
+            image = self.decode_base64_image(encoded_image)
+            out = self.smooth_pipe(prompt, image=image).images[0]
+
+            return out
+    
+    # helper to decode input image
+    def decode_base64_image(self, image_string):
+        base64_image = base64.b64decode(image_string)
+        buffer = BytesIO(base64_image)
+        image = Image.open(buffer)
+        return image