add multi contronet

app_haircolor_inpaint_15.py

@@ -15,6 +15,10 @@ from diffusers import (
     DDIMScheduler,
 )
 
+from controlnet_aux import (
+    CannyDetector,
+)
+
 BASE_MODEL = "stable-diffusion-v1-5/stable-diffusion-v1-5"
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
@@ -24,10 +28,18 @@ DEFAULT_NEGATIVE_PROMPT = "worst quality, normal quality, low quality, low res,
 
 DEFAULT_CATEGORY = "hair"
 
-controlnet = ControlNetModel.from_pretrained(
-    "lllyasviel/control_v11p_sd15_inpaint", 
-    torch_dtype=torch.float16,
-)
+canny_detector = CannyDetector()
+
+controlnet = [
+    ControlNetModel.from_pretrained(
+        "lllyasviel/control_v11p_sd15_inpaint", 
+        torch_dtype=torch.float16,
+    ),
+    ControlNetModel.from_pretrained(
+        "diffusers/controlnet-canny-sdxl-1.0",
+        torch_dtype=torch.float16,
+    ),
+]
 
 basepipeline = StableDiffusionControlNetInpaintPipeline.from_pretrained(
     BASE_MODEL,
@@ -50,11 +62,15 @@ def image_to_image(
     num_steps: int,
     guidance_scale: float,
     generate_size: int,
+    inpaint_scale: float = 1.0,
+    canny_scale: float = 0.5,
 ):
     run_task_time = 0
     time_cost_str = ''
     run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
-    control_image = make_inpaint_condition(input_image, mask_image)
+    canny_image = canny_detector(input_image, int(generate_size*0.375), generate_size)
+    inpaint_image = make_inpaint_condition(input_image, mask_image)
+    control_image = [inpaint_image, canny_image]
 
     generator = torch.Generator(device=DEVICE).manual_seed(seed)
     generated_image = basepipeline(
@@ -66,6 +82,7 @@ def image_to_image(
         control_image=control_image,
         guidance_scale=guidance_scale,
         num_inference_steps=num_steps,
+        controlnet_conditioning_scale=[inpaint_scale, canny_scale]
     ).images[0]
     
     run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
@@ -100,14 +117,17 @@ def create_demo() -> gr.Blocks:
             with gr.Column():
                 edit_prompt = gr.Textbox(lines=1, label="Edit Prompt", value=DEFAULT_EDIT_PROMPT)
                 generate_size = gr.Number(label="Generate Size", value=512)
-                seed = gr.Number(label="Seed", value=8)
-                category = gr.Textbox(label="Category", value=DEFAULT_CATEGORY, visible=False)
             with gr.Column():
                 num_steps = gr.Slider(minimum=1, maximum=100, value=15, step=1, label="Num Steps")
                 guidance_scale = gr.Slider(minimum=0, maximum=30, value=5, step=0.5, label="Guidance Scale")
-                mask_expansion = gr.Number(label="Mask Expansion", value=50, visible=True)
             with gr.Column():
-                mask_dilation = gr.Slider(minimum=0, maximum=10, value=2, step=1, label="Mask Dilation")
+                with gr.Accordion("Advanced Options", open=False):
+                    inpaint_scale = gr.Slider(minimum=0, maximum=3, value=1, step=0.1, label="Inpaint Scale")
+                    canny_scale = gr.Slider(minimum=0, maximum=3, value=0.8, step=0.1, label="Canny Scale")
+                    mask_expansion = gr.Number(label="Mask Expansion", value=50, visible=True)
+                    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)
                 g_btn = gr.Button("Edit Image")
                 
         with gr.Row():
@@ -127,7 +147,7 @@ def create_demo() -> gr.Blocks:
             outputs=[origin_area_image, mask_image, croper],
         ).success(
             fn=image_to_image,
-            inputs=[origin_area_image, mask_image, edit_prompt,seed, num_steps, guidance_scale, generate_size],
+            inputs=[origin_area_image, mask_image, edit_prompt,seed, num_steps, guidance_scale, generate_size, inpaint_scale, canny_scale],
             outputs=[generated_image, generated_cost],
         ).success(
             fn=restore_result,