Add color jitter

1 files changed, 12 insertions, 3 deletions

diff --git a/train_dreambooth.py b/train_dreambooth.py
index 4c36ae4..48921d4 100644
--- a/train_dreambooth.py
+++ b/train_dreambooth.py
@@ -306,7 +306,7 @@ def parse_args():
         "--optimizer",
         type=str,
         default="dadan",
-        choices=["adam", "adam8bit", "lion", "dadam", "dadan", "adafactor"],
+        choices=["adam", "adam8bit", "adan", "lion", "dadam", "dadan", "adafactor"],
         help='Optimizer to use'
     )
     parser.add_argument(
@@ -513,8 +513,6 @@ def main():
         if not embeddings_dir.exists() or not embeddings_dir.is_dir():
             raise ValueError("--embeddings_dir must point to an existing directory")
 
-        embeddings.persist()
-
         added_tokens, added_ids = load_embeddings_from_dir(tokenizer, embeddings, embeddings_dir)
         print(f"Added {len(added_tokens)} tokens from embeddings dir: {list(zip(added_tokens, added_ids))}")
 
@@ -549,6 +547,17 @@ def main():
             eps=args.adam_epsilon,
             amsgrad=args.adam_amsgrad,
         )
+    elif args.optimizer == 'adan':
+        try:
+            import timm.optim
+        except ImportError:
+            raise ImportError("To use Adan, please install the PyTorch Image Models library: `pip install timm`.")
+
+        create_optimizer = partial(
+            timm.optim.Adan,
+            weight_decay=args.adam_weight_decay,
+            eps=args.adam_epsilon,
+        )
     elif args.optimizer == 'lion':
         try:
             import lion_pytorch