From e09aaedd0e74f2fc6e2a53f233914803c65e127c Mon Sep 17 00:00:00 2001 From: Volpeon Date: Sat, 24 Dec 2022 10:25:58 +0100 Subject: Training update --- train_ti.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) (limited to 'train_ti.py') 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, -- cgit v1.2.3-54-g00ecf