add canny

app_haircolor_inpaint_15.py

@@ -53,6 +53,10 @@ controlnet = [
         "lllyasviel/control_v11p_sd15_softedge",
         torch_dtype=torch.float16,
     ),
+    ControlNetModel.from_pretrained(
+        "lllyasviel/control_v11p_sd15_canny",
+        torch_dtype=torch.float16,
+    ),
 ]
 
 basepipeline = StableDiffusionControlNetInpaintPipeline.from_pretrained(
@@ -79,15 +83,16 @@ def image_to_image(
     generate_size: int,
     cond_scale1: float = 1.2,
     cond_scale2: float = 1.2,
+    cond_scale3: float = 1.2,
 ):
     run_task_time = 0
     time_cost_str = ''
     run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
-    # canny_image = canny_detector(input_image, int(generate_size*0.375), generate_size)
-    lineart_image = lineart_detector(input_image, int(generate_size*0.375), generate_size)
+    canny_image = canny_detector(input_image, int(generate_size*1), generate_size)
+    lineart_image = lineart_detector(input_image, int(generate_size*1), generate_size)
     run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
-    pidiNet_image = pidiNet_detector(input_image, int(generate_size*0.5), generate_size)
-    control_image = [lineart_image, pidiNet_image]
+    pidiNet_image = pidiNet_detector(input_image, int(generate_size*1), generate_size)
+    control_image = [lineart_image, pidiNet_image, canny_image]
 
     generator = torch.Generator(device=DEVICE).manual_seed(seed)
     generated_image = basepipeline(
@@ -101,7 +106,7 @@ def image_to_image(
         width=generate_size,
         guidance_scale=guidance_scale,
         num_inference_steps=num_steps,
-        controlnet_conditioning_scale=[cond_scale1, cond_scale2]
+        controlnet_conditioning_scale=[cond_scale1, cond_scale2, cond_scale3],
     ).images[0]
     
     run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)