Fix choice args

3 files changed, 25 insertions, 23 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",
diff --git a/train_lora.py b/train_lora.py
index a0cd174..538a7f7 100644
--- a/train_lora.py
+++ b/train_lora.py
@@ -261,10 +261,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",
@@ -288,7 +287,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",
@@ -300,7 +299,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",
@@ -318,7 +317,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",
diff --git a/train_ti.py b/train_ti.py
index 4366c9e..6757bde 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -204,7 +204,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(
         "--offset_noise_strength",
@@ -253,10 +254,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",
@@ -280,7 +280,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",
@@ -292,7 +292,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",
@@ -330,7 +330,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",