From f894dfecfaa3ec17903b2ac37ac4f071408613db Mon Sep 17 00:00:00 2001
From: Volpeon <git@volpeon.ink>
Date: Fri, 17 Feb 2023 21:06:11 +0100
Subject: Added Lion optimizer

---
 infer.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

(limited to 'infer.py')

diff --git a/infer.py b/infer.py
index 51cf3a7..8910e68 100644
--- a/infer.py
+++ b/infer.py
@@ -23,6 +23,7 @@ from diffusers import (
     EulerAncestralDiscreteScheduler,
     KDPM2DiscreteScheduler,
     KDPM2AncestralDiscreteScheduler,
+    DEISMultistepScheduler,
     UniPCMultistepScheduler
 )
 from transformers import CLIPTextModel
@@ -126,13 +127,13 @@ def create_cmd_parser():
     parser.add_argument(
         "--scheduler",
         type=str,
-        choices=["plms", "ddim", "klms", "dpmsm", "dpmss", "euler_a", "kdpm2", "kdpm2_a", "unipc"],
+        choices=["plms", "ddim", "klms", "dpmsm", "dpmss", "euler_a", "kdpm2", "kdpm2_a", "deis", "unipc"],
     )
     parser.add_argument(
         "--subscheduler",
         type=str,
         default=None,
-        choices=["plms", "ddim", "klms", "dpmsm", "dpmss", "euler_a", "kdpm2", "kdpm2_a"],
+        choices=["plms", "ddim", "klms", "dpmsm", "dpmss", "euler_a", "kdpm2", "kdpm2_a", "deis"],
     )
     parser.add_argument(
         "--template",
@@ -252,6 +253,8 @@ def create_scheduler(config, scheduler: str, subscheduler: Optional[str] = None)
         return KDPM2DiscreteScheduler.from_config(config)
     elif scheduler == "kdpm2_a":
         return KDPM2AncestralDiscreteScheduler.from_config(config)
+    elif scheduler == "deis":
+        return DEISMultistepScheduler.from_config(config)
     elif scheduler == "unipc":
         if subscheduler is None:
             return UniPCMultistepScheduler.from_config(config)
-- 
cgit v1.2.3-54-g00ecf