feat: make additional schedulers available in UI

ssube · Jan 8, 2023 · 93e53f6 · 93e53f6
1 parent 16c8b54
commit 93e53f6
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 1 deletion.
diff --git a/api/onnx_web/serve.py b/api/onnx_web/serve.py
@@ -3,8 +3,13 @@
     DDIMScheduler,
     DDPMScheduler,
     DPMSolverMultistepScheduler,
+    DPMSolverSinglestepScheduler,
     EulerDiscreteScheduler,
     EulerAncestralDiscreteScheduler,
+    HeunDiscreteScheduler,
+    KDPM2AncestralDiscreteScheduler,
+    KDPM2DiscreteScheduler,
+    KarrasVeScheduler,
     LMSDiscreteScheduler,
     PNDMScheduler,
     # onnx
@@ -54,8 +59,13 @@
     'ddim': DDIMScheduler,
     'ddpm': DDPMScheduler,
     'dpm-multi': DPMSolverMultistepScheduler,
+    'dpm-single': DPMSolverSinglestepScheduler,
     'euler': EulerDiscreteScheduler,
     'euler-a': EulerAncestralDiscreteScheduler,
+    'heun': HeunDiscreteScheduler,
+    'k-dpm-2-a': KDPM2AncestralDiscreteScheduler,
+    'k-dpm-2': KDPM2DiscreteScheduler,
+    'karras-ve': KarrasVeScheduler,
     'lms-discrete': LMSDiscreteScheduler,
     'pndm': PNDMScheduler,
 }

diff --git a/gui/src/strings.ts b/gui/src/strings.ts
@@ -12,8 +12,13 @@ export const SCHEDULER_LABELS: Record<string, string> = {
   'ddim': 'DDIM',
   'ddpm': 'DDPM',
   'dpm-multi': 'DPM Multistep',
+  'dpm-single': 'DPM Singlestep',
   'euler': 'Euler',
   'euler-a': 'Euler Ancestral',
-  'lms-discrete': 'LMS Discrete',
+  'heun': 'Heun',
+  'k-dpm-2-a': 'KDPM2 Ancestral',
+  'k-dpm-2': 'KDPM2',
+  'karras-ve': 'Karras Ve',
+  'lms-discrete': 'LMS',
   'pndm': 'PNDM',
 };