From 6b8a93f46f053668c8023520225a18445d48d8f1 Mon Sep 17 00:00:00 2001
From: Volpeon <git@volpeon.ink>
Date: Sat, 25 Mar 2023 16:34:48 +0100
Subject: Update

---
 train_ti.py | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

(limited to 'train_ti.py')

diff --git a/train_ti.py b/train_ti.py
index bbc5524..83ad46d 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -90,6 +90,11 @@ def parse_args():
         "--sequential",
         action="store_true",
     )
+    parser.add_argument(
+        "--guidance_scale",
+        type=float,
+        default=0,
+    )
     parser.add_argument(
         "--num_class_images",
         type=int,
@@ -167,7 +172,7 @@ def parse_args():
     parser.add_argument(
         "--tag_dropout",
         type=float,
-        default=0.1,
+        default=0,
         help="Tag dropout probability.",
     )
     parser.add_argument(
@@ -190,7 +195,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(
@@ -651,8 +656,8 @@ def main():
         noise_scheduler=noise_scheduler,
         dtype=weight_dtype,
         seed=args.seed,
-        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,
         strategy=textual_inversion_strategy,
         num_train_epochs=args.num_train_epochs,
@@ -705,6 +710,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,
-- 
cgit v1.2.3-54-g00ecf