1 files changed, 1 insertions, 8 deletions

diff --git a/train_dreambooth.py b/train_dreambooth.py
index 659b84c..0543a35 100644
--- a/train_dreambooth.py
+++ b/train_dreambooth.py
@@ -246,12 +246,6 @@ def parse_args():
         ),
     )
     parser.add_argument(
-        "--offset_noise_strength",
-        type=float,
-        default=0,
-        help="Perlin offset noise strength.",
-    )
-    parser.add_argument(
         "--input_pertubation",
         type=float,
         default=0,
@@ -496,7 +490,6 @@ def parse_args():
         default=1.0,
         help="The weight of prior preservation loss.",
     )
-    parser.add_argument("--run_pti", action="store_true", help="Whether to run PTI.")
     parser.add_argument("--emb_alpha", type=float, default=1.0, help="Embedding alpha")
     parser.add_argument(
         "--emb_dropout",
@@ -679,6 +672,7 @@ def main():
 
     if args.gradient_checkpointing:
         unet.enable_gradient_checkpointing()
+        text_encoder.gradient_checkpointing_enable()
 
     if len(args.alias_tokens) != 0:
         alias_placeholder_tokens = args.alias_tokens[::2]
@@ -1074,7 +1068,6 @@ def main():
             sample_output_dir=dreambooth_sample_output_dir,
             checkpoint_output_dir=dreambooth_checkpoint_output_dir,
             sample_frequency=dreambooth_sample_frequency,
-            offset_noise_strength=args.offset_noise_strength,
             input_pertubation=args.input_pertubation,
             no_val=args.valid_set_size == 0,
             avg_loss=avg_loss,