Added CLI arg to set dataloader worker num; improved text encoder handling with Dreambooth

1 files changed, 10 insertions, 0 deletions

diff --git a/textual_inversion.py b/textual_inversion.py
index dd7c3bd..115f3aa 100644
--- a/textual_inversion.py
+++ b/textual_inversion.py
@@ -117,6 +117,15 @@ def parse_args():
         help="Whether to center crop images before resizing to resolution"
     )
     parser.add_argument(
+        "--dataloader_num_workers",
+        type=int,
+        default=0,
+        help=(
+            "The number of subprocesses to use for data loading. 0 means that the data will be loaded in the main"
+            " process."
+        ),
+    )
+    parser.add_argument(
         "--num_train_epochs",
         type=int,
         default=100
@@ -626,6 +635,7 @@ def main():
         repeats=args.repeats,
         center_crop=args.center_crop,
         valid_set_size=args.sample_batch_size*args.sample_batches,
+        num_workers=args.dataloader_num_workers,
         collate_fn=collate_fn
     )