add t2i-adapter-sketch

app_haircolor.py

@@ -17,9 +17,11 @@ from diffusers import (
 from controlnet_aux import (
     LineartDetector,
     CannyDetector,
+    PidiNetDetector,
+    MidasDetector,
 )
 
-BASE_MODEL = "SG161222/RealVisXL_V5.0_Lightning"
+BASE_MODEL = "stabilityai/sdxl-turbo"
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 DEFAULT_EDIT_PROMPT = "a woman, blue hair, high detailed"
@@ -30,8 +32,16 @@ DEFAULT_CATEGORY = "hair"
 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)
+
 canndy_detector = CannyDetector()
 
+midas_detector = MidasDetector.from_pretrained(
+    "valhalla/t2iadapter-aux-models", filename="dpt_large_384.pt", model_type="dpt_large"
+)
+midas_detector = midas_detector.to(DEVICE)
+
 adapters = MultiAdapter(
     [
         T2IAdapter.from_pretrained(
@@ -44,6 +54,11 @@ adapters = MultiAdapter(
             torch_dtype=torch.float16,
             varient="fp16",
         ),
+        T2IAdapter.from_pretrained(
+            "TencentARC/t2i-adapter-sketch-sdxl-1.0",
+            torch_dtype=torch.float16,
+            varient="fp16",
+        ),
     ]
 )
 adapters = adapters.to(torch.float16)
@@ -61,7 +76,7 @@ basepipeline = basepipeline.to(DEVICE)
 
 basepipeline.enable_model_cpu_offload()
 
-@spaces.GPU(duration=30)
+@spaces.GPU(duration=15)
 def image_to_image(
     input_image: Image,
     edit_prompt: str,
@@ -71,6 +86,7 @@ def image_to_image(
     generate_size: int,
     lineart_scale: float = 1.0,
     canny_scale: float = 0.5,
+    sketch_scale:float = 0.5,
 ):
     run_task_time = 0
     time_cost_str = ''
@@ -79,9 +95,11 @@ def image_to_image(
     run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
     canny_image = canndy_detector(input_image, 384, generate_size)
     run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
+    sketch_image = pidinet_detector(input_image, 512, generate_size)
+    run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
 
-    cond_image = [lineart_image, canny_image]
-    cond_scale = [lineart_scale, canny_scale]
+    cond_image = [lineart_image, canny_image, sketch_image]
+    cond_scale = [lineart_scale, canny_scale, sketch_scale]
 
     generator = torch.Generator(device=DEVICE).manual_seed(seed)
     generated_image = basepipeline(
@@ -127,8 +145,9 @@ def create_demo() -> gr.Blocks:
                 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")
-                lineart_scale = gr.Slider(minimum=0, maximum=2, value=0.3, step=0.1, label="Lineart Scale")
-                canny_scale = gr.Slider(minimum=0, maximum=2, value=0.7, step=0.1, label="Canny Scale")
+                lineart_scale = gr.Slider(minimum=0, maximum=5, value=1, step=0.1, label="Lineart Scale")
+                canny_scale = gr.Slider(minimum=0, maximum=5, value=0.7, step=0.1, label="Canny Scale")
+                sketch_scale = gr.Slider(minimum=0, maximum=5, value=1, step=0.1, label="Sketch Scale")
                 g_btn = gr.Button("Edit Image")
                 
         with gr.Row():
@@ -147,7 +166,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, lineart_scale, canny_scale],
+            inputs=[origin_area_image, edit_prompt,seed, num_steps, guidance_scale, generate_size, lineart_scale, canny_scale, sketch_scale],
             outputs=[generated_image, generated_cost],
         ).success(
             fn=restore_result,