Misc improvements

1 files changed, 18 insertions, 29 deletions

diff --git a/train_ti.py b/train_ti.py
index 6aa4007..088c1a6 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -93,16 +93,10 @@ def parse_args():
         help="The directory where class images will be saved.",
     )
     parser.add_argument(
-        "--exclude_keywords",
+        "--exclude_collections",
         type=str,
         nargs='*',
-        help="Skip dataset items containing a listed keyword.",
-    )
-    parser.add_argument(
-        "--exclude_modes",
-        type=str,
-        nargs='*',
-        help="Exclude all items with a listed mode.",
+        help="Exclude all items with a listed collection.",
     )
     parser.add_argument(
         "--repeats",
@@ -123,10 +117,10 @@ def parse_args():
         help="The embeddings directory where Textual Inversion embeddings are stored.",
     )
     parser.add_argument(
-        "--mode",
+        "--collection",
         type=str,
-        default=None,
-        help="A mode to filter the dataset.",
+        nargs='*',
+        help="A collection to filter the dataset.",
     )
     parser.add_argument(
         "--seed",
@@ -369,11 +363,11 @@ def parse_args():
     if len(args.placeholder_token) != len(args.initializer_token):
         raise ValueError("You must specify --placeholder_token")
 
-    if isinstance(args.exclude_keywords, str):
-        args.exclude_keywords = [args.exclude_keywords]
+    if isinstance(args.collection, str):
+        args.collection = [args.collection]
 
-    if isinstance(args.exclude_modes, str):
-        args.exclude_modes = [args.exclude_modes]
+    if isinstance(args.exclude_collections, str):
+        args.exclude_collections = [args.exclude_collections]
 
     if args.output_dir is None:
         raise ValueError("You must specify --output_dir")
@@ -600,17 +594,12 @@ def main():
             for keyword in args.placeholder_token
             for part in item.prompt
         )
-        cond2 = args.exclude_keywords is None or not any(
-            keyword in part
-            for keyword in args.exclude_keywords
-            for part in item.prompt
-        )
-        cond3 = args.mode is None or args.mode in item.mode
-        cond4 = args.exclude_modes is None or not any(
-            mode in item.mode
-            for mode in args.exclude_modes
+        cond3 = args.collection is None or args.collection in item.collection
+        cond4 = args.exclude_collections is None or not any(
+            collection in item.collection
+            for collection in args.exclude_collections
         )
-        return cond1 and cond2 and cond3 and cond4
+        return cond1 and cond3 and cond4
 
     def collate_fn(examples):
         prompts = [example["prompts"] for example in examples]
@@ -827,10 +816,10 @@ def main():
         config = vars(args).copy()
         config["initializer_token"] = " ".join(config["initializer_token"])
         config["placeholder_token"] = " ".join(config["placeholder_token"])
-        if config["exclude_modes"] is not None:
-            config["exclude_modes"] = " ".join(config["exclude_modes"])
-        if config["exclude_keywords"] is not None:
-            config["exclude_keywords"] = " ".join(config["exclude_keywords"])
+        if config["collection"] is not None:
+            config["collection"] = " ".join(config["collection"])
+        if config["exclude_collections"] is not None:
+            config["exclude_collections"] = " ".join(config["exclude_collections"])
         accelerator.init_trackers("textual_inversion", config=config)
 
     # Train!