From 026d3260b4128fe4f2fc603ae9cfbd7ca18b7fae Mon Sep 17 00:00:00 2001
From: blessedcoolant <54517381+blessedcoolant@users.noreply.github.com>
Date: Sat, 13 May 2023 02:43:32 +1200
Subject: [PATCH] Add Heun Karras Scheduler

---
 invokeai/backend/args.py                                   | 1 +
 invokeai/backend/stable_diffusion/schedulers/schedulers.py | 3 ++-
 invokeai/backend/web/modules/parameters.py                 | 1 +
 invokeai/frontend/web/src/app/constants.ts                 | 1 +
 invokeai/frontend/web/src/app/types/invokeai.ts            | 1 +
 5 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/invokeai/backend/args.py b/invokeai/backend/args.py
index db6fbe08df..6a29847c6e 100644
--- a/invokeai/backend/args.py
+++ b/invokeai/backend/args.py
@@ -113,6 +113,7 @@ SAMPLER_CHOICES = [
     "lms",
     "pndm",
     "heun",
+    "heun_k",
     "euler",
     "euler_k",
     "euler_a",
diff --git a/invokeai/backend/stable_diffusion/schedulers/schedulers.py b/invokeai/backend/stable_diffusion/schedulers/schedulers.py
index fab28aca8c..08f85cf559 100644
--- a/invokeai/backend/stable_diffusion/schedulers/schedulers.py
+++ b/invokeai/backend/stable_diffusion/schedulers/schedulers.py
@@ -9,7 +9,8 @@ SCHEDULER_MAP = dict(
     deis=(DEISMultistepScheduler, dict()),
     lms=(LMSDiscreteScheduler, dict()),
     pndm=(PNDMScheduler, dict()),
-    heun=(HeunDiscreteScheduler, dict()),
+    heun=(HeunDiscreteScheduler, dict(use_karras_sigmas=False)),
+    heun_k=(HeunDiscreteScheduler, dict(use_karras_sigmas=True)),
     euler=(EulerDiscreteScheduler, dict(use_karras_sigmas=False)),
     euler_k=(EulerDiscreteScheduler, dict(use_karras_sigmas=True)),
     euler_a=(EulerAncestralDiscreteScheduler, dict()),
diff --git a/invokeai/backend/web/modules/parameters.py b/invokeai/backend/web/modules/parameters.py
index 72211857a3..9a4bc0aec3 100644
--- a/invokeai/backend/web/modules/parameters.py
+++ b/invokeai/backend/web/modules/parameters.py
@@ -9,6 +9,7 @@ SAMPLER_CHOICES = [
     "lms",
     "pndm",
     "heun",
+    'heun_k',
     "euler",
     "euler_k",
     "euler_a",
diff --git a/invokeai/frontend/web/src/app/constants.ts b/invokeai/frontend/web/src/app/constants.ts
index 189fbc9dd4..6ecd20d7fe 100644
--- a/invokeai/frontend/web/src/app/constants.ts
+++ b/invokeai/frontend/web/src/app/constants.ts
@@ -7,6 +7,7 @@ export const DIFFUSERS_SCHEDULERS: Array<string> = [
   'lms',
   'pndm',
   'heun',
+  'heun_k',
   'euler',
   'euler_k',
   'euler_a',
diff --git a/invokeai/frontend/web/src/app/types/invokeai.ts b/invokeai/frontend/web/src/app/types/invokeai.ts
index d0e5437d36..f684dc1ccf 100644
--- a/invokeai/frontend/web/src/app/types/invokeai.ts
+++ b/invokeai/frontend/web/src/app/types/invokeai.ts
@@ -52,6 +52,7 @@ export type CommonGeneratedImageMetadata = {
     | 'lms'
     | 'pndm'
     | 'heun'
+    | 'heun_k'
     | 'euler'
     | 'euler_k'
     | 'euler_a'