Training update

1 files changed, 4 insertions, 6 deletions

diff --git a/train_ti.py b/train_ti.py
index 52bd675..a12b889 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -368,7 +368,6 @@ class Checkpointer(CheckpointerBase):
         tokenizer,
         text_encoder,
         scheduler,
-        text_embeddings,
         placeholder_token,
         placeholder_token_id,
         output_dir: Path,
@@ -394,7 +393,6 @@ class Checkpointer(CheckpointerBase):
         self.tokenizer = tokenizer
         self.text_encoder = text_encoder
         self.scheduler = scheduler
-        self.text_embeddings = text_embeddings
 
     @torch.no_grad()
     def checkpoint(self, step, postfix):
@@ -407,7 +405,7 @@ class Checkpointer(CheckpointerBase):
 
         for (placeholder_token, placeholder_token_id) in zip(self.placeholder_token, self.placeholder_token_id):
             # Save a checkpoint
-            learned_embeds = self.text_embeddings.trainable_embedding.weight[placeholder_token_id]
+            learned_embeds = text_encoder.text_model.embeddings.trainable_embedding.weight[placeholder_token_id]
             learned_embeds_dict = {placeholder_token: learned_embeds.detach().cpu()}
 
             filename = f"%s_%d_%s.bin" % (slugify(placeholder_token), step, postfix)
@@ -517,8 +515,9 @@ def main():
 
     vae.requires_grad_(False)
     unet.requires_grad_(False)
+    text_encoder.requires_grad_(False)
 
-    text_embeddings = patch_trainable_embeddings(text_encoder, placeholder_token_id)
+    patch_trainable_embeddings(text_encoder, placeholder_token_id)
 
     prompt_processor = PromptProcessor(tokenizer, text_encoder)
 
@@ -541,7 +540,7 @@ def main():
 
     # Initialize the optimizer
     optimizer = optimizer_class(
-        text_embeddings.trainable_embedding.parameters(),  # only optimize the embeddings
+        text_encoder.text_model.embeddings.trainable_embedding.parameters(),  # only optimize the embeddings
         lr=args.learning_rate,
         betas=(args.adam_beta1, args.adam_beta2),
         weight_decay=args.adam_weight_decay,
@@ -741,7 +740,6 @@ def main():
         tokenizer=tokenizer,
         text_encoder=text_encoder,
         scheduler=checkpoint_scheduler,
-        text_embeddings=text_embeddings,
         placeholder_token=args.placeholder_token,
         placeholder_token_id=placeholder_token_id,
         output_dir=basepath,