From 0bc909409648a3cae0061c3de2b39e486473ae39 Mon Sep 17 00:00:00 2001
From: Volpeon <git@volpeon.ink>
Date: Thu, 27 Oct 2022 17:57:05 +0200
Subject: Added CLI arg to set dataloader worker num; improved text encoder
 handling with Dreambooth

---
 textual_inversion.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

(limited to 'textual_inversion.py')

diff --git a/textual_inversion.py b/textual_inversion.py
index dd7c3bd..115f3aa 100644
--- a/textual_inversion.py
+++ b/textual_inversion.py
@@ -116,6 +116,15 @@ def parse_args():
         action="store_true",
         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,
@@ -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
     )
 
-- 
cgit v1.2.3-54-g00ecf