1 files changed, 10 insertions, 0 deletions

diff --git a/train_dreambooth.py b/train_dreambooth.py
index d396249..aa5ff01 100644
--- a/train_dreambooth.py
+++ b/train_dreambooth.py
@@ -384,6 +384,12 @@ def parse_args():
         help="Number of images in the validation dataset."
     )
     parser.add_argument(
+        "--valid_set_repeat",
+        type=int,
+        default=None,
+        help="Times the images in the validation dataset are repeated."
+    )
+    parser.add_argument(
         "--train_batch_size",
         type=int,
         default=1,
@@ -451,6 +457,9 @@ def parse_args():
     if isinstance(args.exclude_collections, str):
         args.exclude_collections = [args.exclude_collections]
 
+    if args.valid_set_repeat is None:
+        args.valid_set_repeat = args.train_batch_size
+
     if args.output_dir is None:
         raise ValueError("You must specify --output_dir")
 
@@ -764,6 +773,7 @@ def main():
         dropout=args.tag_dropout,
         template_key=args.train_data_template,
         valid_set_size=args.valid_set_size,
+        valid_set_repeat=args.valid_set_repeat,
         num_workers=args.dataloader_num_workers,
         seed=args.seed,
         filter=keyword_filter,