Update

1 files changed, 9 insertions, 8 deletions

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: