change to img2img

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 
-from app_haircolor_inpaint_15 import create_demo as create_demo_haircolor
+from app_haircolor_img2img import create_demo as create_demo_haircolor
 
 with gr.Blocks(css="style.css") as demo:
     with gr.Tabs():

app_haircolor_img2img.py

@@ -10,22 +10,19 @@ from segment_utils import(
     restore_result,
 )
 from diffusers import (
-    StableDiffusionControlNetImg2ImgPipeline,
-    ControlNetModel,
-    DDIMScheduler,
-    DPMSolverMultistepScheduler,
+    DiffusionPipeline,
+    T2IAdapter,
+    MultiAdapter,
+    AutoencoderKL,
     EulerAncestralDiscreteScheduler,
-    UniPCMultistepScheduler,
 )
 
 from controlnet_aux import (
     CannyDetector,
     LineartDetector,
-    PidiNetDetector,
-    HEDdetector,
 )
 
-BASE_MODEL = "stable-diffusion-v1-5/stable-diffusion-v1-5"
+BASE_MODEL = "stabilityai/stable-diffusion-xl-base-1.0"
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -37,32 +34,34 @@ 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,
-    ),
-]
+adapters = MultiAdapter(
+    [
+        T2IAdapter.from_pretrained(
+            "TencentARC/t2i-adapter-lineart-sdxl-1.0",
+            torch_dtype=torch.float16,
+            varient="fp16",
+        ),
+        T2IAdapter.from_pretrained(
+            "TencentARC/t2i-adapter-canny-sdxl-1.0",
+            torch_dtype=torch.float16,
+            varient="fp16",
+        ),
+    ]
+)
+adapters = adapters.to(torch.float16)
 
-basepipeline = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
+basepipeline = DiffusionPipeline.from_pretrained(
     BASE_MODEL,
-    torch_dtype=torch.float16,
+    torch_dtype=torch.float16, 
+    variant="fp16", 
     use_safetensors=True,
-    controlnet=controlnet,
+    vae=AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16),
+    scheduler=EulerAncestralDiscreteScheduler.from_pretrained(BASE_MODEL, subfolder="scheduler"),
+    adapter=adapters,
+    custom_pipeline="./pipelines/pipeline_sdxl_adapter_img2img.py",
 )
 
-basepipeline.scheduler = UniPCMultistepScheduler.from_config(basepipeline.scheduler.config)
-
 basepipeline = basepipeline.to(DEVICE)
 
 basepipeline.enable_model_cpu_offload()
@@ -78,15 +77,17 @@ def image_to_image(
     generate_size: int,
     cond_scale1: float = 1.2,
     cond_scale2: float = 1.2,
+    lineart_detect:float = 0.375,
+    canny_detect:float = 0.375,
 ):
     run_task_time = 0
     time_cost_str = ''
     run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
-    lineart_image = lineart_detector(input_image, 768, generate_size)
+    lineart_image = lineart_detector(input_image, int(generate_size * lineart_detect), generate_size)
+    canny_image = canny_detector(input_image, int(generate_size * canny_detect), generate_size)
 
-    pidinet_image = pidiNet_detector(input_image, 768, generate_size)
-
-    cond_image = [lineart_image, pidinet_image]
+    cond_image = [lineart_image, canny_image]
+    cond_scale = [cond_scale1, cond_scale2]
 
     generator = torch.Generator(device=DEVICE).manual_seed(seed)
     generated_image = basepipeline(
@@ -99,7 +100,7 @@ def image_to_image(
         guidance_scale=guidance_scale,
         strength=strength,
         num_inference_steps=num_steps,
-        controlnet_conditioning_scale=[cond_scale1, cond_scale2],
+        controlnet_conditioning_scale=cond_scale,
     ).images[0]
     
     run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
@@ -137,15 +138,17 @@ def create_demo() -> gr.Blocks:
             with gr.Column():
                 num_steps = gr.Slider(minimum=1, maximum=100, value=20, step=1, label="Num Steps")
                 guidance_scale = gr.Slider(minimum=0, maximum=30, value=5, step=0.5, label="Guidance Scale")
-                strength = gr.Slider(minimum=0, maximum=3, value=0.2, step=0.1, label="Strength")
             with gr.Column():
+                strength = gr.Slider(minimum=0, maximum=3, value=0.2, step=0.1, label="Strength")
                 with gr.Accordion("Advanced Options", open=False):
                     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)
-                    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")
+                    cond_scale1 = gr.Slider(minimum=0, maximum=3, value=0.8, step=0.1, label="Cond_scale1")
+                    cond_scale2 = gr.Slider(minimum=0, maximum=3, value=0.3, step=0.1, label="Cond_scale2")
+                    lineart_detect = gr.Slider(minimum=0, maximum=1, value=0.375, step=0.01, label="Lineart Detect")
+                    canny_detect = gr.Slider(minimum=0, maximum=1, value=0.75, step=0.01, label="Canny Detect")
                 g_btn = gr.Button("Edit Image")
                 
         with gr.Row():
@@ -164,7 +167,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, strength, generate_size, cond_scale1, cond_scale2],
+            inputs=[origin_area_image, edit_prompt,seed, num_steps, guidance_scale, strength, generate_size, cond_scale1, cond_scale2, lineart_detect, canny_detect],
             outputs=[generated_image, generated_cost],
         ).success(
             fn=restore_result,