feat: add support for DEIS multistep scheduler

ssube · Feb 14, 2023 · 6b6f2f0 · 6b6f2f0
1 parent 05756b2
commit 6b6f2f0
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 0 deletions.
diff --git a/api/onnx_web/diffusion/load.py b/api/onnx_web/diffusion/load.py
@@ -18,6 +18,11 @@
     PNDMScheduler,
 )
 
+try:
+    from diffusers import DEISMultistepScheduler
+except:
+    from .stub_scheduler import StubScheduler as DEISMultistepScheduler
+
 from ..params import DeviceParams, Size
 from ..utils import ServerContext, run_gc
 
@@ -29,6 +34,7 @@
 pipeline_schedulers = {
     "ddim": DDIMScheduler,
     "ddpm": DDPMScheduler,
+    "deis-multi": DEISMultistepScheduler,
     "dpm-multi": DPMSolverMultistepScheduler,
     "dpm-single": DPMSolverSinglestepScheduler,
     "euler": EulerDiscreteScheduler,

diff --git a/api/onnx_web/diffusion/stub_scheduler.py b/api/onnx_web/diffusion/stub_scheduler.py
@@ -0,0 +1,6 @@
+from typing import Any
+
+
+class StubScheduler():
+  def step(self, model_output: Any, timestep: int, sample: Any, return_dict: bool = True) -> None:
+    raise NotImplemented("scheduler not available, try updating diffusers")
diff --git a/gui/src/strings.ts b/gui/src/strings.ts
@@ -44,6 +44,7 @@ export const PLATFORM_LABELS: Record<string, string> = {
 export const SCHEDULER_LABELS: Record<string, string> = {
   'ddim': 'DDIM',
   'ddpm': 'DDPM',
+  'deis-multi': 'DEIS Multistep',
   'dpm-multi': 'DPM Multistep',
   'dpm-single': 'DPM Singlestep',
   'euler': 'Euler',

diff --git a/onnx-web.code-workspace b/onnx-web.code-workspace
@@ -22,6 +22,7 @@
 			"CUDA",
 			"ddim",
 			"ddpm",
+			"deis",
 			"denoise",
 			"denoising",
 			"directml",