Update handler.py

handler.py

@@ -1,10 +1,10 @@
 from typing import  Dict, List, Any
 import torch
-from diffusers import DPMSolverMultistepScheduler, StableDiffusionInpaintPipeline, AutoPipelineForInpainting, AutoPipelineForImage2Image, StableDiffusionXLImg2ImgPipeline
+from diffusers import DPMSolverMultistepScheduler, StableDiffusionInpaintPipeline, AutoPipelineForInpainting, AutoPipelineForImage2Image, StableDiffusionXLImg2ImgPipeline, StableDiffusionControlNetInpaintPipeline, ControlNetModel
 from PIL import Image
 import base64
 from io import BytesIO
-
+import numpy as np
 
 # set device
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@@ -12,6 +12,16 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 if device.type != 'cuda':
     raise ValueError("need to run on GPU")
 
+def make_inpaint_condition(image, image_mask):
+    image = np.array(image.convert("RGB")).astype(np.float32) / 255.0
+    image_mask = np.array(image_mask.convert("L")).astype(np.float32) / 255.0
+
+    assert image.shape[0:1] == image_mask.shape[0:1], "image and image_mask must have the same image size"
+    image[image_mask > 0.5] = -1.0  # set as masked pixel
+    image = np.expand_dims(image, 0).transpose(0, 3, 1, 2)
+    image = torch.from_numpy(image)
+    return image
+
 class EndpointHandler():
     def __init__(self, path=""):
 
@@ -24,7 +34,19 @@ class EndpointHandler():
         # )
         # self.smooth_pipe.to("cuda")
 
-        
+        self.controlnet = ControlNetModel.from_pretrained(
+            "lllyasviel/control_v11p_sd15_inpaint", torch_dtype=torch.float16
+        )
+
+        self.pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
+            "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16
+        )
+
+        self.pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+        self.pipe.enable_model_cpu_offload()
+        self.pipe.enable_xformers_memory_efficient_attention()
+
+        """
         # load StableDiffusionInpaintPipeline pipeline
         self.pipe = AutoPipelineForInpainting.from_pretrained(
             "runwayml/stable-diffusion-inpainting",
@@ -50,7 +72,7 @@ class EndpointHandler():
         self.pipe3 = AutoPipelineForImage2Image.from_pipe(self.pipe2)
         #self.pipe3.enable_model_cpu_offload()
         self.pipe3.enable_xformers_memory_efficient_attention()
-        
+        """
 
 
     def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
@@ -102,7 +124,7 @@ class EndpointHandler():
         """
 
         #pipe = AutoPipelineForInpainting.from_pretrained("diffusers/stable-diffusion-xl-1.0-inpainting-0.1", torch_dtype=torch.float16, variant="fp16").to("cuda")
-        
+        """
         # run inference pipeline
         out = self.pipe(prompt=prompt, negative_prompt=negative_prompt, image=image, mask_image=mask_image)
 
@@ -137,6 +159,21 @@ class EndpointHandler():
             
         # return first generate PIL image
         return image2
+        """
+
+        control_image = make_inpaint_condition(image, mask_image)
+
+        # generate image
+        image = pipe(
+            prompt,
+            num_inference_steps=num_inference_steps,
+            eta=1.0,
+            image=image,
+            mask_image=mask_image,
+            control_image=control_image,
+        ).images[0]
+
+        return image
     
     # helper to decode input image
     def decode_base64_image(self, image_string):