add control net

app_haircolor_inpaint_15.py

@@ -10,10 +10,20 @@ from segment_utils import(
     restore_result,
 )
 from diffusers import (
-    StableDiffusionInpaintPipeline,
+    StableDiffusionControlNetInpaintPipeline,
+    ControlNetModel,
+    DDIMScheduler,
+    DPMSolverMultistepScheduler,
     EulerAncestralDiscreteScheduler,
 )
 
+from controlnet_aux import (
+    CannyDetector,
+    LineartDetector,
+    PidiNetDetector,
+    HEDdetector,
+)
+
 # BASE_MODEL = "stable-diffusion-v1-5/stable-diffusion-v1-5"
 BASE_MODEL = "stable-diffusion-v1-5/stable-diffusion-inpainting"
 # BASE_MODEL = "SG161222/Realistic_Vision_V2.0"
@@ -25,12 +35,34 @@ DEFAULT_NEGATIVE_PROMPT = "worst quality, normal quality, low quality, low res,
 
 DEFAULT_CATEGORY = "hair"
 
-basepipeline = StableDiffusionInpaintPipeline.from_pretrained(
+canny_detector = CannyDetector()
+lineart_detector = LineartDetector.from_pretrained("lllyasviel/Annotators")
+lineart_detector = lineart_detector.to(DEVICE)
+
+pidiNet_detector = PidiNetDetector.from_pretrained('lllyasviel/Annotators')
+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_v11p_sd15_lineart", 
+        torch_dtype=torch.float16,
+    ),
+    ControlNetModel.from_pretrained(
+        "lllyasviel/control_v11p_sd15_softedge", 
+        torch_dtype=torch.float16,
+    ),
+]
+
+basepipeline = StableDiffusionControlNetInpaintPipeline.from_pretrained(
     BASE_MODEL,
     torch_dtype=torch.float16,
     # use_safetensors=True,
+    controlnet=controlnet,
 )
-
+# basepipeline.scheduler = DDIMScheduler.from_config(basepipeline.scheduler.config)
 basepipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(basepipeline.scheduler.config)
 
 basepipeline = basepipeline.to(DEVICE)
@@ -52,6 +84,11 @@ def image_to_image(
     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*1), generate_size)
+    lineart_image = lineart_detector(input_image, 384, generate_size)
+    run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
+    pidiNet_image = pidiNet_detector(input_image, 512, generate_size)
+    control_image = [lineart_image, pidiNet_image]
 
     generator = torch.Generator(device=DEVICE).manual_seed(seed)
     generated_image = basepipeline(
@@ -60,10 +97,12 @@ def image_to_image(
         negative_prompt=DEFAULT_NEGATIVE_PROMPT,
         image=input_image,
         mask_image=mask_image,
+        control_image=control_image,
         height=generate_size,
         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)
@@ -103,8 +142,8 @@ def create_demo() -> gr.Blocks:
                 guidance_scale = gr.Slider(minimum=0, maximum=30, value=5, step=0.5, label="Guidance Scale")
             with gr.Column():
                 with gr.Accordion("Advanced Options", open=False):
-                    cond_scale1 = gr.Slider(minimum=0, maximum=3, value=1, step=0.1, label="Cond Scale1")
-                    cond_scale2 = gr.Slider(minimum=0, maximum=3, value=0.6, step=0.1, label="Cond Scale2")
+                    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")
                     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)