Update handler.py

handler.py

@@ -1,6 +1,6 @@
 from typing import  Dict, List, Any
 import torch
-from diffusers import DPMSolverMultistepScheduler, StableDiffusionInpaintPipeline, AutoPipelineForInpainting, AutoPipelineForImage2Image
+from diffusers import DPMSolverMultistepScheduler, StableDiffusionInpaintPipeline, AutoPipelineForInpainting, AutoPipelineForImage2Image, StableDiffusionXLImg2ImgPipeline
 from PIL import Image
 import base64
 from io import BytesIO
@@ -18,6 +18,12 @@ class EndpointHandler():
         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")
+
         
         # load StableDiffusionInpaintPipeline pipeline
         self.pipe = AutoPipelineForInpainting.from_pretrained(
@@ -76,6 +82,11 @@ class EndpointHandler():
     
             return image
 
+        if(method == "smooth"):
+            image = self.smooth_pipe(prompt, image=image).images[0]
+
+            return image
+
         #pipe = AutoPipelineForInpainting.from_pretrained("diffusers/stable-diffusion-xl-1.0-inpainting-0.1", torch_dtype=torch.float16, variant="fp16").to("cuda")
 
         self.pipe.enable_xformers_memory_efficient_attention()