1 files changed, 10 insertions, 3 deletions

diff --git a/train_ti.py b/train_ti.py
index 082e9b7..94ddbb6 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -72,6 +72,12 @@ def parse_args():
         help="How many cycles to run automatically."
     )
     parser.add_argument(
+        "--cycle_decay",
+        type=float,
+        default=1.0,
+        help="Learning rate decay per cycle."
+    )
+    parser.add_argument(
         "--placeholder_tokens",
         type=str,
         nargs='*',
@@ -672,7 +678,6 @@ def main():
     convnext.to(accelerator.device, dtype=weight_dtype)
     convnext.requires_grad_(False)
     convnext.eval()
-    disc = ConvNeXtDiscriminator(convnext, input_size=384)
 
     if len(args.alias_tokens) != 0:
         alias_placeholder_tokens = args.alias_tokens[::2]
@@ -815,7 +820,6 @@ def main():
         milestone_checkpoints=not args.no_milestone_checkpoints,
         global_step_offset=global_step_offset,
         offset_noise_strength=args.offset_noise_strength,
-        disc=disc,
         # --
         use_emb_decay=args.use_emb_decay,
         emb_decay_target=args.emb_decay_target,
@@ -890,6 +894,7 @@ def main():
             sample_frequency = math.ceil(num_train_epochs / args.sample_num)
 
         training_iter = 0
+        learning_rate = args.learning_rate
 
         project = placeholder_tokens[0] if len(placeholder_tokens) == 1 else "ti"
 
@@ -908,7 +913,7 @@ def main():
 
             optimizer = create_optimizer(
                 text_encoder.text_model.embeddings.token_override_embedding.parameters(),
-                lr=args.learning_rate,
+                lr=learning_rate,
             )
 
             lr_scheduler = get_scheduler(
@@ -948,6 +953,8 @@ def main():
             )
 
             training_iter += 1
+            if args.learning_rate is not None:
+                learning_rate *= args.cycle_decay
 
         accelerator.end_training()