1 files changed, 9 insertions, 8 deletions

diff --git a/train_dreambooth.py b/train_dreambooth.py
index be7d6fe..4c36ae4 100644
--- a/train_dreambooth.py
+++ b/train_dreambooth.py
@@ -130,7 +130,8 @@ def parse_args():
         "--vector_shuffle",
         type=str,
         default="auto",
-        help='Vector shuffling algorithm. Choose between ["all", "trailing", "leading", "between", "auto", "off"]',
+        choices=["all", "trailing", "leading", "between", "auto", "off"],
+        help='Vector shuffling algorithm.',
     )
     parser.add_argument(
         "--guidance_scale",
@@ -229,10 +230,9 @@ def parse_args():
         "--lr_scheduler",
         type=str,
         default="one_cycle",
-        help=(
-            'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
-            ' "constant", "constant_with_warmup", "one_cycle"]'
-        ),
+        choices=["linear", "cosine", "cosine_with_restarts", "polynomial",
+                 "constant", "constant_with_warmup", "one_cycle"],
+        help='The scheduler type to use.',
     )
     parser.add_argument(
         "--lr_warmup_epochs",
@@ -256,7 +256,7 @@ def parse_args():
         "--lr_warmup_func",
         type=str,
         default="cos",
-        help='Choose between ["linear", "cos"]'
+        choices=["linear", "cos"],
     )
     parser.add_argument(
         "--lr_warmup_exp",
@@ -268,7 +268,7 @@ def parse_args():
         "--lr_annealing_func",
         type=str,
         default="cos",
-        help='Choose between ["linear", "half_cos", "cos"]'
+        choices=["linear", "half_cos", "cos"],
     )
     parser.add_argument(
         "--lr_annealing_exp",
@@ -306,7 +306,8 @@ def parse_args():
         "--optimizer",
         type=str,
         default="dadan",
-        help='Optimizer to use ["adam", "adam8bit", "lion", "dadam", "dadan", "adafactor"]'
+        choices=["adam", "adam8bit", "lion", "dadam", "dadan", "adafactor"],
+        help='Optimizer to use'
     )
     parser.add_argument(
         "--dadaptation_d0",