Added Lion optimizer

1 files changed, 5 insertions, 2 deletions

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)