Misc improvements

1 files changed, 18 insertions, 29 deletions

diff --git a/train_dreambooth.py b/train_dreambooth.py
index 072150b..8fd78f1 100644
--- a/train_dreambooth.py
+++ b/train_dreambooth.py
@@ -83,16 +83,10 @@ def parse_args():
         help="A token to use as initializer word."
     )
     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(
         "--train_text_encoder",
@@ -142,10 +136,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",
@@ -391,11 +385,11 @@ def parse_args():
     if len(args.placeholder_token) != len(args.initializer_token):
         raise ValueError("Number of items in --placeholder_token and --initializer_token must match")
 
-    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")
@@ -655,17 +649,12 @@ def main():
         weight_dtype = torch.bfloat16
 
     def keyword_filter(item: CSVDataItem):
-        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 cond2 and cond3 and cond4
+        return cond3 and cond4
 
     def collate_fn(examples):
         prompts = [example["prompts"] for example in examples]
@@ -813,10 +802,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("dreambooth", config=config)
 
     # Train!