From 7c02c2fe68da2411623f0a11c1187ccf0f7743d8 Mon Sep 17 00:00:00 2001 From: Volpeon Date: Thu, 1 Dec 2022 22:01:47 +0100 Subject: Update --- dreambooth.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) (limited to 'dreambooth.py') diff --git a/dreambooth.py b/dreambooth.py index 31dbea2..1ead6dd 100644 --- a/dreambooth.py +++ b/dreambooth.py @@ -550,11 +550,11 @@ class Checkpointer: def main(): args = parse_args() - if args.train_text_encoder and args.gradient_accumulation_steps > 1 and accelerator.num_processes > 1: - raise ValueError( - "Gradient accumulation is not supported when training the text encoder in distributed training. " - "Please set gradient_accumulation_steps to 1. This feature will be supported in the future." - ) + # if args.train_text_encoder and args.gradient_accumulation_steps > 1 and accelerator.num_processes > 1: + # raise ValueError( + # "Gradient accumulation is not supported when training the text encoder in distributed training. " + # "Please set gradient_accumulation_steps to 1. This feature will be supported in the future." + # ) instance_identifier = args.instance_identifier @@ -899,6 +899,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}") @@ -910,7 +913,7 @@ def main(): sample_checkpoint = False for step, batch in enumerate(train_dataloader): - with accelerator.accumulate(unet): + with accelerator.accumulate(itertools.chain(unet, text_encoder)): # Convert images to latent space latents = vae.encode(batch["pixel_values"]).latent_dist.sample() latents = latents * 0.18215 @@ -967,8 +970,6 @@ def main(): else: 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, :] if accelerator.sync_gradients: -- cgit v1.2.3-54-g00ecf