1 files changed, 9 insertions, 2 deletions

diff --git a/dreambooth.py b/dreambooth.py
index 3dd0920..31dbea2 100644
--- a/dreambooth.py
+++ b/dreambooth.py
@@ -32,6 +32,7 @@ logger = get_logger(__name__)
 
 
 torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.benchmark = True
 
 
 def parse_args():
@@ -474,7 +475,6 @@ class Checkpointer:
             scheduler=self.scheduler,
         ).to(self.accelerator.device)
         pipeline.set_progress_bar_config(dynamic_ncols=True)
-        pipeline.enable_vae_slicing()
 
         train_data = self.datamodule.train_dataloader()
         val_data = self.datamodule.val_dataloader()
@@ -550,6 +550,12 @@ 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."
+        )
+
     instance_identifier = args.instance_identifier
 
     if len(args.placeholder_token) != 0:
@@ -587,6 +593,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:
@@ -903,7 +910,7 @@ def main():
             sample_checkpoint = False
 
             for step, batch in enumerate(train_dataloader):
-                with accelerator.accumulate(itertools.chain(unet, text_encoder)):
+                with accelerator.accumulate(unet):
                     # Convert images to latent space
                     latents = vae.encode(batch["pixel_values"]).latent_dist.sample()
                     latents = latents * 0.18215