p2p add canny

app_haircolor_img2img.py

@@ -44,10 +44,16 @@ pidiNet_detector = pidiNet_detector.to(DEVICE)
 hed_detector = HEDdetector.from_pretrained('lllyasviel/Annotators')
 hed_detector = hed_detector.to(DEVICE)
 
-controlnet = ControlNetModel.from_pretrained(
-    "lllyasviel/control_v11e_sd15_ip2p", 
-    torch_dtype=torch.float16,
-)
+controlnet = [
+    ControlNetModel.from_pretrained(
+        "lllyasviel/control_v11e_sd15_ip2p", 
+        torch_dtype=torch.float16,
+    ),
+    ControlNetModel.from_pretrained(
+        "lllyasviel/control_v11p_sd15_canny", 
+        torch_dtype=torch.float16,
+    ),
+]
 
 basepipeline = StableDiffusionControlNetPipeline.from_pretrained(
     BASE_MODEL,
@@ -70,12 +76,15 @@ def image_to_image(
     num_steps: int,
     guidance_scale: float,
     generate_size: int,
+    cond_scale1: float = 1.2,
+    cond_scale2: 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)
 
-    cond_image = input_image
+    cond_image = [input_image, canny_image]
 
     generator = torch.Generator(device=DEVICE).manual_seed(seed)
     generated_image = basepipeline(
@@ -87,6 +96,7 @@ def image_to_image(
         width=generate_size,
         guidance_scale=guidance_scale,
         num_inference_steps=num_steps,
+        controlnet_conditioning_scale=[cond_scale1, cond_scale2],
     ).images[0]
     
     run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
@@ -130,6 +140,8 @@ def create_demo() -> gr.Blocks:
                     mask_dilation = gr.Slider(minimum=0, maximum=10, value=2, step=1, label="Mask Dilation")
                     seed = gr.Number(label="Seed", value=8)
                     category = gr.Textbox(label="Category", value=DEFAULT_CATEGORY, visible=False)
+                    cond_scale1 = gr.Slider(minimum=0, maximum=3, value=1.2, step=0.1, label="Cond_scale1")
+                    cond_scale2 = gr.Slider(minimum=0, maximum=3, value=1.2, step=0.1, label="Cond_scale2")
                 g_btn = gr.Button("Edit Image")
                 
         with gr.Row():
@@ -148,7 +160,7 @@ def create_demo() -> gr.Blocks:
             outputs=[origin_area_image, croper],
         ).success(
             fn=image_to_image,
-            inputs=[origin_area_image, edit_prompt,seed, num_steps, guidance_scale, generate_size],
+            inputs=[origin_area_image, edit_prompt,seed, num_steps, guidance_scale, generate_size, cond_scale1, cond_scale2],
             outputs=[generated_image, generated_cost],
         ).success(
             fn=restore_result,