Update handler.py

handler.py

@@ -25,29 +25,14 @@ class EndpointHandler():
         # )
         # self.smooth_pipe.to("cuda")
 
-        self.canny_pipe = StableDiffusionPipeline.from_pretrained(
-            "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16
-        )
-        self.canny_pipe = self.canny_pipe.to("cuda")
-        self.canny_pipe.enable_model_cpu_offload()
-        self.canny_pipe.enable_xformers_memory_efficient_attention()
-
-        self.controlnets = [
-            ControlNetModel.from_pretrained(
-                "diffusers/controlnet-canny-sdxl-1.0", torch_dtype=torch.float16, use_safetensors=True
-            ),
-            ControlNetModel.from_pretrained(
-                "lllyasviel/control_v11p_sd15_inpaint", torch_dtype=torch.float16
-            )
-        ]
-        """
+
         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=self.controlnets, torch_dtype=torch.float16
+            "diffusers/stable-diffusion-xl-1.0-inpainting-0.1", controlnet=self.controlnets, torch_dtype=torch.float16
         )
 
         self.pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(self.pipe.scheduler.config)
@@ -172,53 +157,20 @@ class EndpointHandler():
         return image
         """
         
-        
         control_image = self.make_inpaint_condition(image, mask_image)
 
-        low_threshold = 100
-        high_threshold = 200
-
-        # generate a first version of the prompt for the canny image
-        gen_canny_img = self.canny_pipe(prompt).images[0]
-
-        gen_canny_img = np.array(gen_canny_img)
-        
-        help_image = cv2.Canny(gen_canny_img, low_threshold, high_threshold)
-        
-        # get bounding box from selected area in mask image
-        # make help_image fit exactly into the bounding box
-        # create black image with canny edges only in the selected area
-        
-        # get bounding box from selected area in mask image
-        mask_image = np.array(mask_image)
-        mask_image = cv2.cvtColor(mask_image, cv2.COLOR_RGB2GRAY)
-        mask_image = cv2.bitwise_not(mask_image)
-        contours, _ = cv2.findContours(mask_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-        x, y, w, h = cv2.boundingRect(contours[0])
-        
-        # create a completely black image with the same size as the mask image
-        black_image = np.zeros_like(mask_image)
-        
-        # put the canny edges into the black image but resize it to the bounding box
-        help_image = cv2.resize(help_image, (w, h))
-        black_image[y:y+h, x:x+w] = help_image
-
-        canny_image = Image.fromarray(black_image)
-
-        input_images = [canny_image.resize((1024, 1024)), image.resize((1024, 1024))]
 
         # generate image
         image = self.pipe(
-            prompt,
-            input_images,
+            prompt=prompt,
+            image=image,
             negative_prompt=negative_prompt,
             num_inference_steps=num_inference_steps,
             eta=1.0,
-            # mask_image=mask_image,
-            # control_image=control_image,
-            # guidance_scale=guidance_scale,
+            mask_image=mask_image,
+            control_image=control_image,
+            guidance_scale=guidance_scale,
             strength=strength,
-            controlnet_conditioning_scale=[0.8, 1.0]
         ).images[0]
 
         return image