Fixed Textual Inversion

1 files changed, 8 insertions, 2 deletions

diff --git a/train_ti.py b/train_ti.py
index a12b889..5f37d54 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -25,7 +25,7 @@ from pipelines.stable_diffusion.vlpn_stable_diffusion import VlpnStableDiffusion
 from data.csv import CSVDataModule, CSVDataItem
 from training.optimization import get_one_cycle_schedule
 from training.ti import patch_trainable_embeddings
-from training.util import AverageMeter, CheckpointerBase, save_args
+from training.util import AverageMeter, CheckpointerBase, save_args, freeze_params
 from models.clip.prompt import PromptProcessor
 
 logger = get_logger(__name__)
@@ -515,10 +515,16 @@ def main():
 
     vae.requires_grad_(False)
     unet.requires_grad_(False)
-    text_encoder.requires_grad_(False)
 
     patch_trainable_embeddings(text_encoder, placeholder_token_id)
 
+    freeze_params(itertools.chain(
+        text_encoder.text_model.encoder.parameters(),
+        text_encoder.text_model.final_layer_norm.parameters(),
+        text_encoder.text_model.embeddings.position_embedding.parameters(),
+        text_encoder.text_model.embeddings.token_embedding.parameters(),
+    ))
+
     prompt_processor = PromptProcessor(tokenizer, text_encoder)
 
     if args.scale_lr: