From 93e53f6dc3d312f348df156780d5f89712f1c087 Mon Sep 17 00:00:00 2001
From: Sean Sube <seansube@gmail.com>
Date: Sun, 8 Jan 2023 08:19:24 -0600
Subject: [PATCH] feat: make additional schedulers available in UI

---
 api/onnx_web/serve.py | 10 ++++++++++
 gui/src/strings.ts    |  7 ++++++-
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/api/onnx_web/serve.py b/api/onnx_web/serve.py
index ee5a8c17..1c5fe1ae 100644
--- a/api/onnx_web/serve.py
+++ b/api/onnx_web/serve.py
@@ -3,8 +3,13 @@ from diffusers import (
     DDIMScheduler,
     DDPMScheduler,
     DPMSolverMultistepScheduler,
+    DPMSolverSinglestepScheduler,
     EulerDiscreteScheduler,
     EulerAncestralDiscreteScheduler,
+    HeunDiscreteScheduler,
+    KDPM2AncestralDiscreteScheduler,
+    KDPM2DiscreteScheduler,
+    KarrasVeScheduler,
     LMSDiscreteScheduler,
     PNDMScheduler,
     # onnx
@@ -54,8 +59,13 @@ pipeline_schedulers = {
     '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
index dd1af4bf..ae213724 100644
--- 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',
 };