Fix TI

1 files changed, 10 insertions, 1 deletions

diff --git a/train_ti.py b/train_ti.py
index ef39c38..9ae8d1b 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -155,7 +155,7 @@ def parse_args():
     parser.add_argument(
         "--num_buckets",
         type=int,
-        default=2,
+        default=0,
         help="Number of aspect ratio buckets in either direction.",
     )
     parser.add_argument(
@@ -507,9 +507,18 @@ def parse_args():
     if isinstance(args.num_vectors, list) and len(args.placeholder_tokens) != len(args.num_vectors):
         raise ValueError("--placeholder_tokens and --num_vectors must have the same number of items")
 
+    if args.alias_tokens is None:
+        args.alias_tokens = []
+
     if not isinstance(args.alias_tokens, list) or len(args.alias_tokens) % 2 != 0:
         raise ValueError("--alias_tokens must be a list with an even number of items")
 
+    args.alias_tokens += [
+        item
+        for pair in zip(args.placeholder_tokens, args.initializer_tokens)
+        for item in pair
+    ]
+
     if args.sequential:
         if isinstance(args.train_data_template, str):
             args.train_data_template = [args.train_data_template] * len(args.placeholder_tokens)