change scheduler

app_haircolor_pix2pix.py

@@ -12,13 +12,14 @@ from enhance_utils import enhance_image
 from diffusers import (
     StableDiffusionInstructPix2PixPipeline,
     EulerAncestralDiscreteScheduler,
+    DDIMScheduler,
 )
 
 BASE_MODEL = "timbrooks/instruct-pix2pix"
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
-DEFAULT_EDIT_PROMPT = "change hair to blue"
+DEFAULT_EDIT_PROMPT = "hair to linen-blonde-hair"
 
 DEFAULT_CATEGORY = "hair"
 
@@ -28,7 +29,7 @@ basepipeline = StableDiffusionInstructPix2PixPipeline.from_pretrained(
     use_safetensors=True,
 )
 
-basepipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(basepipeline.scheduler.config)
+basepipeline.scheduler = DDIMScheduler.from_config(basepipeline.scheduler.config)
 
 basepipeline = basepipeline.to(DEVICE)
 
@@ -46,14 +47,12 @@ 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)
-    gray_image = input_image.convert("L")
-    p2p_image = Image.merge("RGB", [gray_image, gray_image, gray_image])
 
     generator = torch.Generator(device=DEVICE).manual_seed(seed)
     generated_image = basepipeline(
         generator=generator,
         prompt=edit_prompt,
-        image=p2p_image,
+        image=input_image,
         guidance_scale=guidance_scale,
         image_guidance_scale=image_guidance_scale,
         num_inference_steps=num_steps,