1 files changed, 27 insertions, 11 deletions

diff --git a/train_dreambooth.py b/train_dreambooth.py
index 5a7911c..8f0c6ea 100644
--- a/train_dreambooth.py
+++ b/train_dreambooth.py
@@ -285,9 +285,10 @@ def parse_args():
         default=0.9999
     )
     parser.add_argument(
-        "--use_8bit_adam",
-        action="store_true",
-        help="Whether or not to use 8-bit Adam from bitsandbytes."
+        "--optimizer",
+        type=str,
+        default="lion",
+        help='Optimizer to use ["adam", "adam8bit", "lion"]'
     )
     parser.add_argument(
         "--adam_beta1",
@@ -491,15 +492,34 @@ def main():
         args.learning_rate = 1e-6
         args.lr_scheduler = "exponential_growth"
 
-    if args.use_8bit_adam:
+    if args.optimizer == 'adam8bit':
         try:
             import bitsandbytes as bnb
         except ImportError:
             raise ImportError("To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`.")
 
-        optimizer_class = bnb.optim.AdamW8bit
+        create_optimizer = partial(
+            bnb.optim.AdamW8bit,
+            betas=(args.adam_beta1, args.adam_beta2),
+            weight_decay=args.adam_weight_decay,
+            eps=args.adam_epsilon,
+            amsgrad=args.adam_amsgrad,
+        )
+    elif args.optimizer == 'adam':
+        create_optimizer = partial(
+            torch.optim.AdamW,
+            betas=(args.adam_beta1, args.adam_beta2),
+            weight_decay=args.adam_weight_decay,
+            eps=args.adam_epsilon,
+            amsgrad=args.adam_amsgrad,
+        )
     else:
-        optimizer_class = torch.optim.AdamW
+        try:
+            from lion_pytorch import Lion
+        except ImportError:
+            raise ImportError("To use Lion, please install the lion_pytorch library: `pip install lion_pytorch`.")
+
+        create_optimizer = partial(Lion, use_triton=True)
 
     trainer = partial(
         train,
@@ -540,17 +560,13 @@ def main():
     )
     datamodule.setup()
 
-    optimizer = optimizer_class(
+    optimizer = create_optimizer(
         itertools.chain(
             unet.parameters(),
             text_encoder.text_model.encoder.parameters(),
             text_encoder.text_model.final_layer_norm.parameters(),
         ),
         lr=args.learning_rate,
-        betas=(args.adam_beta1, args.adam_beta2),
-        weight_decay=args.adam_weight_decay,
-        eps=args.adam_epsilon,
-        amsgrad=args.adam_amsgrad,
     )
 
     lr_scheduler = get_scheduler(