Update app.py

app.py

@@ -11,6 +11,13 @@ import PIL.Image
 import spaces
 import torch
 from diffusers import AutoencoderKL, StableDiffusionXLPipeline
+from diffusers import (
+    DDIMScheduler,
+    DPMSolverMultistepScheduler,
+    DPMSolverSinglestepScheduler,
+    EulerAncestralDiscreteScheduler,
+    EulerDiscreteScheduler,
+)
 
 DESCRIPTION = "# humblemikey/PixelWave10"
 if not torch.cuda.is_available():
@@ -46,6 +53,16 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
         seed = random.randint(0, MAX_SEED)
     return seed
 
+def get_scheduler(scheduler_config: Dict, name: str) -> Optional[Callable]:
+    scheduler_factory_map = {
+        "DPM++ 2M Karras": lambda: DPMSolverMultistepScheduler.from_config(scheduler_config, use_karras_sigmas=True),
+        "DPM++ SDE Karras": lambda: DPMSolverSinglestepScheduler.from_config(scheduler_config, use_karras_sigmas=True),
+        "DPM++ 2M SDE Karras": lambda: DPMSolverMultistepScheduler.from_config(scheduler_config, use_karras_sigmas=True, algorithm_type="sde-dpmsolver++"),
+        "Euler": lambda: EulerDiscreteScheduler.from_config(scheduler_config),
+        "Euler a": lambda: EulerAncestralDiscreteScheduler.from_config(scheduler_config),
+        "DDIM": lambda: DDIMScheduler.from_config(scheduler_config),
+    }
+    return scheduler_factory_map.get(name, lambda: None)()
 
 @spaces.GPU
 def generate(
@@ -57,9 +74,14 @@ def generate(
     height: int = 1024,
     guidance_scale_base: float = 4.0,
     num_inference_steps_base: int = 40,
+    sampler: str = "DPM++ 2M SDE Karras",
+    progress=gr.Progress(track_tqdm=True)
 ) -> PIL.Image.Image:
     generator = torch.Generator().manual_seed(seed)
 
+    #backup_scheduler = pipe.scheduler
+    pipe.scheduler = get_scheduler(pipe.scheduler.config, sampler)
+
     if not use_negative_prompt:
         negative_prompt = None  # type: ignore