From b73469706091c8aaf3f028de96ab017f5a845639 Mon Sep 17 00:00:00 2001
From: Volpeon <git@volpeon.ink>
Date: Tue, 13 Dec 2022 20:49:57 +0100
Subject: Optimized Textual Inversion training by filtering dataset by
 existence of added tokens

---
 dreambooth.py | 6 ++++++
 1 file changed, 6 insertions(+)

(limited to 'dreambooth.py')

diff --git a/dreambooth.py b/dreambooth.py
index 31416e9..5521b21 100644
--- a/dreambooth.py
+++ b/dreambooth.py
@@ -56,6 +56,11 @@ def parse_args():
         default=None,
         help="A folder containing the training data."
     )
+    parser.add_argument(
+        "--train_data_template",
+        type=str,
+        default="template",
+    )
     parser.add_argument(
         "--instance_identifier",
         type=str,
@@ -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
-- 
cgit v1.2.3-54-g00ecf