From 7c02c2fe68da2411623f0a11c1187ccf0f7743d8 Mon Sep 17 00:00:00 2001 From: Volpeon Date: Thu, 1 Dec 2022 22:01:47 +0100 Subject: Update --- textual_inversion.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) (limited to 'textual_inversion.py') 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 -- cgit v1.2.3-54-g00ecf