Merge pull request ddPn08#87 from Isotr0py/karras

Support karras samplers
camenduru · May 16, 2023 · 4ced4f1 · 4ced4f1
2 parents 0eb692e + e6b58cf
commit 4ced4f1
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 5 deletions.
diff --git a/lib/diffusers/scheduler.py b/lib/diffusers/scheduler.py
@@ -1,5 +1,7 @@
 import diffusers
 
+# TODO: add rest Auto111 samplers
+# Commented schedulers are still unavailable in diffusers 0.16.1
 SCHEDULERS = {
     "unipc": diffusers.schedulers.UniPCMultistepScheduler,
     "euler_a": diffusers.schedulers.EulerAncestralDiscreteScheduler,
@@ -9,8 +11,24 @@
     "deis": diffusers.schedulers.DEISMultistepScheduler,
     "dpm2": diffusers.schedulers.KDPM2DiscreteScheduler,
     "dpm2-a": diffusers.schedulers.KDPM2AncestralDiscreteScheduler,
-    "heun": diffusers.schedulers.DPMSolverMultistepScheduler,
-    "dpm++": diffusers.schedulers.DPMSolverMultistepScheduler,
-    "dpm": diffusers.schedulers.DPMSolverMultistepScheduler,
+    "dpm++_2s": diffusers.schedulers.DPMSolverSinglestepScheduler,
+    "dpm++_2m": diffusers.schedulers.DPMSolverMultistepScheduler,
+    "dpm++_2m_karras": diffusers.schedulers.DPMSolverMultistepScheduler,
+    # "dpm++_sde": diffusers.schedulers.DPMSolverSDEScheduler,
+    # "dpm++_sde_karras": diffusers.schedulers.DPMSolverSDEScheduler,
+    "heun": diffusers.schedulers.HeunDiscreteScheduler,
+    "heun_karras": diffusers.schedulers.HeunDiscreteScheduler,
+    "lms": diffusers.schedulers.LMSDiscreteScheduler,
+    # "lms_karras": diffusers.schedulers.LMSDiscreteScheduler,
     "pndm": diffusers.schedulers.PNDMScheduler,
 }
+
+
+def parser_schedulers_config(scheduler_id: str):
+    """
+    add extra config parameter to scheduler
+    """
+    kwargs = {}
+    if "karras" in scheduler_id:
+        kwargs["use_karras_sigmas"] = True
+    return kwargs
diff --git a/modules/model.py b/modules/model.py
@@ -8,7 +8,7 @@
 import torch
 
 from api.models.diffusion import ImageGenerationOptions
-from lib.diffusers.scheduler import SCHEDULERS
+from lib.diffusers.scheduler import SCHEDULERS, parser_schedulers_config
 
 from . import config, utils
 from .images import save_image
@@ -133,7 +133,7 @@ def swap_scheduler(self, scheduler_id: str):
         if not self.activated:
             raise RuntimeError("Model not activated")
         self.pipe.scheduler = SCHEDULERS[scheduler_id].from_config(
-            self.pipe.scheduler.config
+            self.pipe.scheduler.config, **parser_schedulers_config(scheduler_id)
         )
 
     def __call__(self, opts: ImageGenerationOptions, plugin_data: Dict[str, List] = {}):