1 files changed, 14 insertions, 1 deletions

diff --git a/train_dreambooth.py b/train_dreambooth.py
index ab3ed16..7745d27 100644
--- a/train_dreambooth.py
+++ b/train_dreambooth.py
@@ -189,6 +189,12 @@ def parse_args():
         help="Tag dropout probability.",
     )
     parser.add_argument(
+        "--prompt_dropout",
+        type=float,
+        default=0,
+        help="Prompt dropout probability.",
+    )
+    parser.add_argument(
         "--no_tag_shuffle",
         action="store_true",
         help="Shuffle tags.",
@@ -255,6 +261,11 @@ def parse_args():
         help="Number of epochs the text encoder will be trained.",
     )
     parser.add_argument(
+        "--text_encoder_unfreeze_last_n_layers",
+        default=2,
+        help="Number of text encoder layers to train.",
+    )
+    parser.add_argument(
         "--find_lr",
         action="store_true",
         help="Automatically find a learning rate (no training).",
@@ -908,7 +919,8 @@ def main():
     dreambooth_datamodule = create_datamodule(
         valid_set_size=args.valid_set_size,
         batch_size=args.train_batch_size,
-        dropout=args.tag_dropout,
+        tag_dropout=args.tag_dropout,
+        prompt_dropout=args.prompt_dropout,
         filter=partial(keyword_filter, None, args.collection, args.exclude_collections),
     )
     dreambooth_datamodule.setup()
@@ -1051,6 +1063,7 @@ def main():
             checkpoint_output_dir=dreambooth_checkpoint_output_dir,
             sample_frequency=dreambooth_sample_frequency,
             input_pertubation=args.input_pertubation,
+            text_encoder_unfreeze_last_n_layers=args.text_encoder_unfreeze_last_n_layers,
             no_val=args.valid_set_size == 0,
             avg_loss=avg_loss,
             avg_acc=avg_acc,