change base model

app_haircolor.py

@@ -11,13 +11,15 @@ from segment_utils import(
 from diffusers import (
     DiffusionPipeline,
     T2IAdapter,
+    MultiAdapter,
 )
 
 from controlnet_aux import (
     LineartDetector,
+    CannyDetector,
 )
 
-BASE_MODEL = "stabilityai/sdxl-turbo"
+BASE_MODEL = "stabilityai/stable-diffusion-xl-base-1.0"
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 DEFAULT_EDIT_PROMPT = "a woman, blue hair, high detailed"
@@ -28,18 +30,30 @@ DEFAULT_CATEGORY = "hair"
 lineart_detector = LineartDetector.from_pretrained("lllyasviel/Annotators")
 lineart_detector = lineart_detector.to(DEVICE)
 
-adapter = T2IAdapter.from_pretrained(
-    "TencentARC/t2i-adapter-lineart-sdxl-1.0",
-    torch_dtype=torch.float16,
-    varient="fp16",
+canndy_detector = CannyDetector()
+
+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 = DiffusionPipeline.from_pretrained(
     BASE_MODEL,
     torch_dtype=torch.float16,
     use_safetensors=True,
-    adapter=adapter,
+    adapter=adapters,
     custom_pipeline="./pipelines/pipeline_sdxl_adapter_img2img.py",
 )
 
@@ -55,13 +69,20 @@ def image_to_image(
     num_steps: int,
     guidance_scale: float,
     generate_size: int,
-    adapter_weight: float = 1.0,
+    lineart_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)
     lineart_image = lineart_detector(input_image, int(generate_size*0.375), generate_size)
     run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
+    canny_image = canndy_detector(input_image, int(generate_size*0.375), 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]
+
     generator = torch.Generator(device=DEVICE).manual_seed(seed)
     generated_image = basepipeline(
         generator=generator,
@@ -72,8 +93,8 @@ def image_to_image(
         width=generate_size,
         guidance_scale=guidance_scale,
         num_inference_steps=num_steps,
-        adapter_image=lineart_image,
-        adapter_conditioning_scale=adapter_weight,
+        adapter_image=cond_image,
+        adapter_conditioning_scale=cond_scale,
     ).images[0]
     
     run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)