add unwrap_model for accelerator, torch.no_grad context for validation, fix accelerator.accumulate call in advanced script

Author: linoytsaban

examples/advanced_diffusion_training/train_dreambooth_lora_flux_advanced.py

@@ -231,9 +231,10 @@ def log_validation(
     autocast_ctx = torch.autocast(accelerator.device.type)
 
     # pre-calculate  prompt embeds, pooled prompt embeds, text ids because t5 does not support autocast
-    prompt_embeds, pooled_prompt_embeds, text_ids = pipeline.encode_prompt(
-        pipeline_args["prompt"], prompt_2=pipeline_args["prompt"]
-    )
+    with torch.no_grad():
+        prompt_embeds, pooled_prompt_embeds, text_ids = pipeline.encode_prompt(
+            pipeline_args["prompt"], prompt_2=pipeline_args["prompt"]
+        )
     images = []
     for _ in range(args.num_validation_images):
         with autocast_ctx:
@@ -2044,6 +2045,11 @@ def get_sigmas(timesteps, n_dim=4, dtype=torch.float32):
                 pivoted_tr = True
 
         for step, batch in enumerate(train_dataloader):
+            models_to_accumulate = [transformer]
+            if not freeze_text_encoder:
+                models_to_accumulate.extend([text_encoder_one])
+                if args.enable_t5_ti:
+                    models_to_accumulate.extend([text_encoder_two])
             if pivoted_te:
                 # stopping optimization of text_encoder params
                 optimizer.param_groups[te_idx]["lr"] = 0.0
@@ -2052,7 +2058,7 @@ def get_sigmas(timesteps, n_dim=4, dtype=torch.float32):
                 logger.info(f"PIVOT TRANSFORMER {epoch}")
                 optimizer.param_groups[0]["lr"] = 0.0
 
-            with accelerator.accumulate(transformer):
+            with accelerator.accumulate(models_to_accumulate):
                 prompts = batch["prompts"]
 
                 # encode batch prompts when custom prompts are provided for each image -

examples/dreambooth/train_dreambooth_lora_flux.py

@@ -185,9 +185,10 @@ def log_validation(
     autocast_ctx = torch.autocast(accelerator.device.type)
 
     # pre-calculate  prompt embeds, pooled prompt embeds, text ids because t5 does not support autocast
-    prompt_embeds, pooled_prompt_embeds, text_ids = pipeline.encode_prompt(
-        pipeline_args["prompt"], prompt_2=pipeline_args["prompt"]
-    )
+    with torch.no_grad():
+        prompt_embeds, pooled_prompt_embeds, text_ids = pipeline.encode_prompt(
+            pipeline_args["prompt"], prompt_2=pipeline_args["prompt"]
+        )
     images = []
     for _ in range(args.num_validation_images):
         with autocast_ctx:
@@ -940,7 +941,7 @@ def _encode_prompt_with_t5(
 
     prompt_embeds = text_encoder(text_input_ids.to(device))[0]
 
-    dtype = text_encoder.dtype
+    dtype = unwrap_model(text_encoder).dtype
     prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
 
     _, seq_len, _ = prompt_embeds.shape
@@ -983,7 +984,7 @@ def _encode_prompt_with_clip(
 
     # Use pooled output of CLIPTextModel
     prompt_embeds = prompt_embeds.pooler_output
-    prompt_embeds = prompt_embeds.to(dtype=text_encoder.dtype, device=device)
+    prompt_embeds = prompt_embeds.to(dtype=unwrap_model(text_encoder).dtype, device=device)
 
     # duplicate text embeddings for each generation per prompt, using mps friendly method
     prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
@@ -1002,7 +1003,7 @@ def encode_prompt(
     text_input_ids_list=None,
 ):
     prompt = [prompt] if isinstance(prompt, str) else prompt
-    dtype = text_encoders[0].dtype
+    dtype = unwrap_model(text_encoders[0]).dtype
 
     pooled_prompt_embeds = _encode_prompt_with_clip(
         text_encoder=text_encoders[0],