Update

1 files changed, 5 insertions, 4 deletions

diff --git a/textual_inversion.py b/textual_inversion.py
index d6be522..80f1d7d 100644
--- a/textual_inversion.py
+++ b/textual_inversion.py
@@ -545,6 +545,7 @@ def main():
     checkpoint_scheduler = DPMSolverMultistepScheduler.from_pretrained(
         args.pretrained_model_name_or_path, subfolder='scheduler')
 
+    vae.enable_slicing()
     unet.set_use_memory_efficient_attention_xformers(True)
 
     if args.gradient_checkpointing:
@@ -814,6 +815,9 @@ def main():
     )
     global_progress_bar.set_description("Total progress")
 
+    index_fixed_tokens = torch.arange(len(tokenizer))
+    index_fixed_tokens = index_fixed_tokens[~torch.isin(index_fixed_tokens, torch.tensor(placeholder_token_id))]
+
     try:
         for epoch in range(num_epochs):
             local_progress_bar.set_description(f"Epoch {epoch + 1} / {num_epochs}")
@@ -827,7 +831,7 @@ def main():
             for step, batch in enumerate(train_dataloader):
                 with accelerator.accumulate(text_encoder):
                     # Convert images to latent space
-                    latents = vae.encode(batch["pixel_values"]).latent_dist.sample()
+                    latents = vae.encode(batch["pixel_values"]).latent_dist.sample().detach()
                     latents = latents * 0.18215
 
                     # Sample noise that we'll add to the latents
@@ -883,7 +887,6 @@ def main():
                         token_embeds = text_encoder.get_input_embeddings().weight
 
                     # Get the index for tokens that we want to freeze
-                    index_fixed_tokens = torch.arange(len(tokenizer)) != placeholder_token_id
                     token_embeds.data[index_fixed_tokens, :] = original_token_embeds[index_fixed_tokens, :]
 
                     optimizer.step()
@@ -927,8 +930,6 @@ def main():
 
             accelerator.wait_for_everyone()
 
-            print(token_embeds[placeholder_token_id])
-
             text_encoder.eval()
             val_loss = 0.0