Fixed Lora PTI

1 files changed, 20 insertions, 17 deletions

diff --git a/train_lora.py b/train_lora.py
index 476efcf..5b0a292 100644
--- a/train_lora.py
+++ b/train_lora.py
@@ -21,7 +21,6 @@ from data.csv import VlpnDataModule, keyword_filter
 from training.functional import train, add_placeholder_tokens, get_models
 from training.lr import plot_metrics
 from training.strategy.lora import lora_strategy
-from training.strategy.ti import textual_inversion_strategy
 from training.optimization import get_scheduler
 from training.util import save_args
 
@@ -829,6 +828,12 @@ def main():
         sample_num_batches=args.sample_batches,
         sample_num_steps=args.sample_steps,
         sample_image_size=args.sample_image_size,
+        placeholder_tokens=args.placeholder_tokens,
+        placeholder_token_ids=placeholder_token_ids,
+        use_emb_decay=args.use_emb_decay,
+        emb_decay_target=args.emb_decay_target,
+        emb_decay=args.emb_decay,
+        max_grad_norm=args.max_grad_norm,
     )
 
     create_datamodule = partial(
@@ -907,7 +912,8 @@ def main():
         )
 
         metrics = trainer(
-            strategy=textual_inversion_strategy,
+            strategy=lora_strategy,
+            pti_mode=True,
             project="pti",
             train_dataloader=pti_datamodule.train_dataloader,
             val_dataloader=pti_datamodule.val_dataloader,
@@ -919,11 +925,6 @@ def main():
             sample_output_dir=pti_sample_output_dir,
             checkpoint_output_dir=pti_checkpoint_output_dir,
             sample_frequency=math.inf,
-            placeholder_tokens=args.placeholder_tokens,
-            placeholder_token_ids=placeholder_token_ids,
-            use_emb_decay=args.use_emb_decay,
-            emb_decay_target=args.emb_decay_target,
-            emb_decay=args.emb_decay,
         )
 
         plot_metrics(metrics, pti_output_dir / "lr.png")
@@ -952,13 +953,21 @@ def main():
     lora_optimizer = create_optimizer(
         [
             {
-                "params": unet.parameters(),
+                "params": (
+                    param
+                    for param in unet.parameters()
+                    if param.requires_grad
+                ),
                 "lr": args.learning_rate_unet,
             },
             {
-                "params": itertools.chain(
-                    text_encoder.text_model.encoder.parameters(),
-                    text_encoder.text_model.final_layer_norm.parameters(),
+                "params": (
+                    param
+                    for param in itertools.chain(
+                        text_encoder.text_model.encoder.parameters(),
+                        text_encoder.text_model.final_layer_norm.parameters(),
+                    )
+                    if param.requires_grad
                 ),
                 "lr": args.learning_rate_text,
             },
@@ -990,12 +999,6 @@ def main():
         sample_output_dir=lora_sample_output_dir,
         checkpoint_output_dir=lora_checkpoint_output_dir,
         sample_frequency=lora_sample_frequency,
-        placeholder_tokens=args.placeholder_tokens,
-        placeholder_token_ids=placeholder_token_ids,
-        use_emb_decay=args.use_emb_decay,
-        emb_decay_target=args.emb_decay_target,
-        emb_decay=args.emb_decay,
-        max_grad_norm=args.max_grad_norm,
     )
 
     plot_metrics(metrics, lora_output_dir / "lr.png")