Update app.py

app.py

@@ -5,10 +5,21 @@ import PIL.Image
 from PIL import Image
 import random
 from diffusers import ControlNetModel, StableDiffusionXLPipeline, AutoencoderKL
-from diffusers import DDIMScheduler, EulerAncestralDiscreteScheduler
 import cv2
 import torch
 
+from diffusers import (
+    DDIMScheduler,
+    DPMSolverMultistepScheduler,
+    EulerDiscreteScheduler,
+    EulerAncestralDiscreteScheduler,
+    HeunDiscreteScheduler,
+    KDPM2DiscreteScheduler,
+    KDPM2AncestralDiscreteScheduler,
+    LMSDiscreteScheduler,
+    UniPCMultistepScheduler,
+)
+
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
@@ -36,8 +47,31 @@ MAX_IMAGE_SIZE = 1216
 
     
 @spaces.GPU
-def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
+def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, sampler_name):
+
+    # サンプラーの設定
+    if sampler_name == "DDIM":
+        pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+    elif sampler_name == "DPMSolverMultistep":
+        pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
+    elif sampler_name == "Euler":
+        pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config)
+    elif sampler_name == "EulerAncestral":
+        pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+    elif sampler_name == "Heun":
+        pipe.scheduler = HeunDiscreteScheduler.from_config(pipe.scheduler.config)
+    elif sampler_name == "KDPM2":
+        pipe.scheduler = KDPM2DiscreteScheduler.from_config(pipe.scheduler.config)
+    elif sampler_name == "KDPM2Ancestral":
+        pipe.scheduler = KDPM2AncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+    elif sampler_name == "LMS":
+        pipe.scheduler = LMSDiscreteScheduler.from_config(pipe.scheduler.config)
+    elif sampler_name == "UniPC":
+        pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
+    else:
+        pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
 
+    
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
@@ -102,6 +136,12 @@ with gr.Blocks(css=css) as demo:
                 value=0,
             )
 
+            sampler_name = gr.Dropdown(
+                label="Sampler",
+                choices=["DDIM", "DPMSolverMultistep", "Euler", "EulerAncestral", "Heun", "KDPM2", "KDPM2Ancestral", "LMS", "UniPC"],
+                value="EulerAncestral",
+            )
+
             randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
 
             with gr.Row():
@@ -138,9 +178,11 @@ with gr.Blocks(css=css) as demo:
                     value=28,
                 )
 
+                
+
     run_button.click(#lambda x: None, inputs=None, outputs=result).then(
         fn=infer,
-        inputs=[prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
+        inputs=[prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, sampler_name],
         outputs=[result]
     )