Update handler.py

handler.py

@@ -1,10 +1,11 @@
 from typing import  Dict, List, Any
 import torch
-from diffusers import DPMSolverMultistepScheduler, EulerDiscreteScheduler, EulerAncestralDiscreteScheduler, DDIMScheduler, StableDiffusionInpaintPipeline, AutoPipelineForInpainting, AutoPipelineForImage2Image, DiffusionPipeline, StableDiffusionXLImg2ImgPipeline, StableDiffusionControlNetInpaintPipeline, ControlNetModel
+from diffusers import DPMSolverMultistepScheduler, EulerDiscreteScheduler, EulerAncestralDiscreteScheduler, DDIMScheduler, StableDiffusionInpaintPipeline, AutoPipelineForInpainting, AutoPipelineForImage2Image, DiffusionPipeline, StableDiffusionXLImg2ImgPipeline, StableDiffusionControlNetInpaintPipeline, ControlNetModel, StableDiffusionPipeline
 from PIL import Image
 import base64
 from io import BytesIO
 import numpy as np
+import cv2
 
 # set device
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@@ -24,13 +25,27 @@ 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.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.controlnet, torch_dtype=torch.float16
+            "runwayml/stable-diffusion-v1-5", controlnet=self.controlnets, torch_dtype=torch.float16
         )
 
         self.pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(self.pipe.scheduler.config)
@@ -98,6 +113,8 @@ class EndpointHandler():
         else:
             image = None
             mask_image = None 
+
+        
         """
         if(method == "fast"):
             image = self.fast_pipe(
@@ -156,17 +173,31 @@ class EndpointHandler():
         
         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)
+        
+        help_image = cv2.Canny(image, low_threshold, high_threshold)
+        help_image = help_image[:, :, None]
+        help_image = np.concatenate([help_image, help_image, help_image], axis=2)
+        canny_image = Image.fromarray(image)
+
+        input_images = [canny_image.resize((1024, 1024)), image.resize((1024, 1024))]
+
         # generate image
         image = self.pipe(
             prompt=prompt,
             negative_prompt=negative_prompt,
             num_inference_steps=num_inference_steps,
             eta=1.0,
-            image=image,
+            image=input_images,
             mask_image=mask_image,
             control_image=control_image,
-            guidance_scale=guidance_scale,
-            strength=strength
+            # guidance_scale=guidance_scale,
+            strength=strength,
+            controlnet_conditioning_scale=[0.8, 1.0]
         ).images[0]
 
         return image