added canny

handler.py

@@ -91,15 +91,15 @@ class EndpointHandler():
 
 
         controlnet_pose_model = "thibaud/controlnet-openpose-sdxl-1.0"
-        # controlnet_canny_model = "diffusers/controlnet-canny-sdxl-1.0"
+        controlnet_canny_model = "diffusers/controlnet-canny-sdxl-1.0"
         # controlnet_depth_model = "diffusers/controlnet-depth-sdxl-1.0-small"
 
         controlnet_pose = ControlNetModel.from_pretrained(
             controlnet_pose_model, torch_dtype=dtype
         ).to(device)
-        # controlnet_canny = ControlNetModel.from_pretrained(
-        #     controlnet_canny_model, torch_dtype=dtype
-        # ).to(device)
+        controlnet_canny = ControlNetModel.from_pretrained(
+            controlnet_canny_model, torch_dtype=dtype
+        ).to(device)
         # controlnet_depth = ControlNetModel.from_pretrained(
         #     controlnet_depth_model, torch_dtype=dtype
         # ).to(device)
@@ -107,10 +107,10 @@ class EndpointHandler():
         openpose = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
         # depth_anything = DepthAnything.from_pretrained('LiheYoung/depth_anything_vitl14').to(device).eval()
 
-        # def get_canny_image(image, t1=100, t2=200):
-        #     image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
-        #     edges = cv2.Canny(image, t1, t2)
-        #     return Image.fromarray(edges, "L")
+        def get_canny_image(image, t1=100, t2=200):
+            image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+            edges = cv2.Canny(image, t1, t2)
+            return Image.fromarray(edges, "L")
         
         # def get_depth_map(image):
         
@@ -135,13 +135,13 @@ class EndpointHandler():
         
         self.controlnet_map = {
             "pose": controlnet_pose,
-            # "canny": controlnet_canny,
+            "canny": controlnet_canny,
             # "depth": controlnet_depth,
         }
 
         self.controlnet_map_fn = {
             "pose": openpose,
-            # "canny": get_canny_image,
+            "canny": get_canny_image,
             # "depth": get_depth_map,
         }
         
@@ -152,12 +152,13 @@ class EndpointHandler():
 
         identitynet_strength_ratio = 0.8
         pose_strength = 0.4
-        self.my_controlnet_selection = ["pose"]
+        canny_strength = 0.3
+        self.my_controlnet_selection = ["pose", "canny"]
 
 
         controlnet_scales = {
             "pose": pose_strength,
-            # "canny": canny_strength,
+            "canny": canny_strength,
             # "depth": depth_strength,
         }
         self.pipe.controlnet = MultiControlNetModel(
@@ -184,7 +185,6 @@ class EndpointHandler():
         pose_image_path = data.pop("pose_image_path", "https://i.ibb.co/9bP9tMb/pose-2-1.jpg")
         
         adapter_strength_ratio = 0.8
-        # canny_strength = 0.3
         # depth_strength = 0.5
         # controlnet_selection = ["pose", "canny", "depth"]