Add UniPC scheduler and remove DDIM

app.py

@@ -152,7 +152,7 @@ with gr.Blocks(
                 )
 
             # Model settings
-            gr.HTML("<h3>Settings</h3>")
+            gr.HTML("<h3>Model</h3>")
             with gr.Row():
                 model = gr.Dropdown(
                     choices=Config.MODELS,

lib/config.py

@@ -9,6 +9,7 @@ from diffusers import (
     EulerDiscreteScheduler,
     StableDiffusionXLImg2ImgPipeline,
     StableDiffusionXLPipeline,
+    UniPCMultistepScheduler,
 )
 from diffusers.utils import logging as diffusers_logging
 from transformers import logging as transformers_logging
@@ -79,13 +80,13 @@ Config = SimpleNamespace(
     ],
     VAE_MODEL="madebyollin/sdxl-vae-fp16-fix",
     REFINER_MODEL="stabilityai/stable-diffusion-xl-refiner-1.0",
-    SCHEDULER="Euler",
+    SCHEDULER="UniPC",
     SCHEDULERS={
-        "DDIM": DDIMScheduler,
-        "DEIS 2M": DEISMultistepScheduler,
+        "DEIS": DEISMultistepScheduler,
         "DPM++ 2M": DPMSolverMultistepScheduler,
         "Euler": EulerDiscreteScheduler,
         "Euler a": EulerAncestralDiscreteScheduler,
+        "UniPC": UniPCMultistepScheduler,
     },
     WIDTH=1024,
     HEIGHT=1024,

lib/loader.py

@@ -126,6 +126,9 @@ class Loader:
                 self.refiner.deepcache.enable()
 
     def load(self, kind, model, scheduler, deepcache_interval, scale, use_karras, use_refiner, progress=None):
+        Pipeline = Config.PIPELINES[kind]
+        Scheduler = Config.SCHEDULERS[scheduler]
+
         scheduler_kwargs = {
             "beta_start": 0.00085,
             "beta_end": 0.012,
@@ -134,32 +137,25 @@ class Loader:
             "steps_offset": 1,
         }
 
-        if scheduler not in ["DDIM", "Euler a"]:
-            scheduler_kwargs["use_karras_sigmas"] = use_karras
-
-        if scheduler == "DDIM":
-            scheduler_kwargs["clip_sample"] = False
-            scheduler_kwargs["set_alpha_to_one"] = False
-
-        if model not in Config.SINGLE_FILE_MODELS:
-            variant = "fp16"
-        else:
-            variant = None
-
-        dtype = torch.float16
         pipeline_kwargs = {
-            "variant": variant,
-            "torch_dtype": dtype,
+            "torch_dtype": torch.float16,
             "add_watermarker": False,
             "scheduler": Config.SCHEDULERS[scheduler](**scheduler_kwargs),
-            "vae": AutoencoderKL.from_pretrained(Config.VAE_MODEL, torch_dtype=dtype),
+            "vae": AutoencoderKL.from_pretrained(Config.VAE_MODEL, torch_dtype=torch.float16),
         }
 
-        self.unload(model, use_refiner, deepcache_interval, scale)
+        if scheduler not in ["Euler a"]:
+            scheduler_kwargs["use_karras_sigmas"] = use_karras
 
-        Pipeline = Config.PIPELINES[kind]
-        Scheduler = Config.SCHEDULERS[scheduler]
+        if model not in Config.SINGLE_FILE_MODELS:
+            pipeline_kwargs["variant"] = "fp16"
+        else:
+            pipeline_kwargs["variant"] = None
+
+        # Unload
+        self.unload(model, use_refiner, deepcache_interval, scale)
 
+        # Load
         try:
             with timer(f"Loading {model}", logger=self.log.info):
                 self.model = model