Optimized Textual Inversion training by filtering dataset by existence of added tokens

1 files changed, 6 insertions, 0 deletions

diff --git a/dreambooth.py b/dreambooth.py
index 31416e9..5521b21 100644
--- a/dreambooth.py
+++ b/dreambooth.py
@@ -57,6 +57,11 @@ def parse_args():
         help="A folder containing the training data."
     )
     parser.add_argument(
+        "--train_data_template",
+        type=str,
+        default="template",
+    )
+    parser.add_argument(
         "--instance_identifier",
         type=str,
         default=None,
@@ -768,6 +773,7 @@ def main():
         repeats=args.repeats,
         dropout=args.tag_dropout,
         center_crop=args.center_crop,
+        template_key=args.train_data_template,
         valid_set_size=args.valid_set_size,
         num_workers=args.dataloader_num_workers,
         collate_fn=collate_fn