2 files changed, 9 insertions, 3 deletions

diff --git a/train_dreambooth.py b/train_dreambooth.py
index da3a075..0fe590f 100644
--- a/train_dreambooth.py
+++ b/train_dreambooth.py
@@ -907,7 +907,7 @@ def main():
             on_before_optimize=on_before_optimize,
             on_after_optimize=on_after_optimize,
         )
-        lr_finder.run(num_epochs=100, end_lr=1e3)
+        lr_finder.run(num_epochs=100, end_lr=1e2)
 
         plt.savefig(basepath.joinpath("lr.png"), dpi=300)
         plt.close()
diff --git a/train_ti.py b/train_ti.py
index 3b7e3b1..e18ee38 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -415,11 +415,17 @@ def parse_args():
     )
     parser.add_argument(
         "--decay_factor",
-        default=100,
+        default=1,
         type=float,
         help="Embedding decay factor."
     )
     parser.add_argument(
+        "--decay_start",
+        default=1e-4,
+        type=float,
+        help="Embedding decay start offset."
+    )
+    parser.add_argument(
         "--noise_timesteps",
         type=int,
         default=1000,
@@ -833,7 +839,7 @@ def main():
     def on_after_optimize(lr: float):
         text_encoder.text_model.embeddings.normalize(
             args.decay_target,
-            min(1.0, args.decay_factor * lr)
+            min(1.0, max(0.0, args.decay_factor * ((lr - args.decay_start) / (args.learning_rate - args.decay_start))))
         )
 
     loop = partial(