1 files changed, 11 insertions, 5 deletions

diff --git a/train_lora.py b/train_lora.py
index b16a99b..684d0cc 100644
--- a/train_lora.py
+++ b/train_lora.py
@@ -88,7 +88,7 @@ def parse_args():
     parser.add_argument(
         "--num_buckets",
         type=int,
-        default=0,
+        default=2,
         help="Number of aspect ratio buckets in either direction.",
     )
     parser.add_argument(
@@ -111,7 +111,7 @@ def parse_args():
     parser.add_argument(
         "--tag_dropout",
         type=float,
-        default=0.1,
+        default=0,
         help="Tag dropout probability.",
     )
     parser.add_argument(
@@ -120,6 +120,11 @@ def parse_args():
         help="Shuffle tags.",
     )
     parser.add_argument(
+        "--guidance_scale",
+        type=float,
+        default=0,
+    )
+    parser.add_argument(
         "--num_class_images",
         type=int,
         default=0,
@@ -167,7 +172,7 @@ def parse_args():
     parser.add_argument(
         "--offset_noise_strength",
         type=float,
-        default=0.15,
+        default=0,
         help="Perlin offset noise strength.",
     )
     parser.add_argument(
@@ -589,8 +594,8 @@ def main():
         vae=vae,
         noise_scheduler=noise_scheduler,
         dtype=weight_dtype,
-        with_prior_preservation=args.num_class_images != 0,
-        prior_loss_weight=args.prior_loss_weight,
+        guidance_scale=args.guidance_scale,
+        prior_loss_weight=args.prior_loss_weight if args.num_class_images != 0 else 0,
         no_val=args.valid_set_size == 0,
     )
 
@@ -602,6 +607,7 @@ def main():
         batch_size=args.train_batch_size,
         tokenizer=tokenizer,
         class_subdir=args.class_image_dir,
+        with_guidance=args.guidance_scale != 0,
         num_class_images=args.num_class_images,
         size=args.resolution,
         num_buckets=args.num_buckets,