1 files changed, 12 insertions, 21 deletions

diff --git a/train_dreambooth.py b/train_dreambooth.py
index a1802a0..c180170 100644
--- a/train_dreambooth.py
+++ b/train_dreambooth.py
@@ -195,15 +195,6 @@ def parse_args():
         ),
     )
     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
@@ -577,24 +568,24 @@ def main():
         )
 
     now = datetime.datetime.now().strftime("%Y-%m-%dT%H-%M-%S")
-    basepath = Path(args.output_dir).joinpath(slugify(args.project), now)
-    basepath.mkdir(parents=True, exist_ok=True)
+    output_dir = Path(args.output_dir).joinpath(slugify(args.project), now)
+    output_dir.mkdir(parents=True, exist_ok=True)
 
     accelerator = Accelerator(
         log_with=LoggerType.TENSORBOARD,
-        logging_dir=f"{basepath}",
+        logging_dir=f"{output_dir}",
         gradient_accumulation_steps=args.gradient_accumulation_steps,
         mixed_precision=args.mixed_precision
     )
 
-    logging.basicConfig(filename=basepath.joinpath("log.txt"), level=logging.DEBUG)
+    logging.basicConfig(filename=output_dir.joinpath("log.txt"), level=logging.DEBUG)
 
     if args.seed is None:
         args.seed = torch.random.seed() >> 32
 
     set_seed(args.seed)
 
-    save_args(basepath, args)
+    save_args(output_dir, args)
 
     tokenizer, text_encoder, vae, unet, noise_scheduler, sample_scheduler, embeddings = get_models(
         args.pretrained_model_name_or_path)
@@ -618,7 +609,7 @@ def main():
         added_tokens, added_ids = load_embeddings_from_dir(tokenizer, embeddings, embeddings_dir)
         print(f"Added {len(added_tokens)} tokens from embeddings dir: {list(zip(added_tokens, added_ids))}")
 
-    placeholder_token_ids = add_placeholder_tokens(
+    placeholder_token_ids, initializer_token_ids = add_placeholder_tokens(
         tokenizer=tokenizer,
         embeddings=embeddings,
         placeholder_tokens=args.placeholder_tokens,
@@ -627,7 +618,9 @@ def main():
     )
 
     if len(placeholder_token_ids) != 0:
-        print(f"Added {len(placeholder_token_ids)} new tokens: {list(zip(args.placeholder_tokens, placeholder_token_ids))}")
+        initializer_token_id_lens = [len(id) for id in initializer_token_ids]
+        placeholder_token_stats = list(zip(args.placeholder_tokens, placeholder_token_ids, initializer_token_id_lens))
+        print(f"Added {len(placeholder_token_ids)} new tokens: {placeholder_token_stats}")
 
     if args.use_ema:
         ema_unet = EMAModel(
@@ -726,7 +719,6 @@ def main():
         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,
         dtype=weight_dtype
@@ -830,7 +822,6 @@ def main():
         noise_scheduler,
         unet,
         text_encoder,
-        args.num_class_images,
         args.prior_loss_weight,
         args.seed,
     )
@@ -848,7 +839,8 @@ def main():
         scheduler=sample_scheduler,
         placeholder_tokens=args.placeholder_tokens,
         placeholder_token_ids=placeholder_token_ids,
-        output_dir=basepath,
+        output_dir=output_dir,
+        sample_steps=args.sample_steps,
         sample_image_size=args.sample_image_size,
         sample_batch_size=args.sample_batch_size,
         sample_batches=args.sample_batches,
@@ -873,7 +865,7 @@ def main():
         )
         lr_finder.run(num_epochs=100, end_lr=1e2)
 
-        plt.savefig(basepath.joinpath("lr.png"), dpi=300)
+        plt.savefig(output_dir.joinpath("lr.png"), dpi=300)
         plt.close()
     else:
         train_loop(
@@ -886,7 +878,6 @@ def main():
             val_dataloader=val_dataloader,
             loss_step=loss_step_,
             sample_frequency=args.sample_frequency,
-            sample_steps=args.sample_steps,
             checkpoint_frequency=args.checkpoint_frequency,
             global_step_offset=0,
             num_epochs=args.num_train_epochs,