4 files changed, 22 insertions, 29 deletions

diff --git a/training/functional.py b/training/functional.py
index be39776..ed8ae3a 100644
--- a/training/functional.py
+++ b/training/functional.py
@@ -168,8 +168,7 @@ def save_samples(
             image_grid = pipeline.numpy_to_pil(image_grid.unsqueeze(0).permute(0, 2, 3, 1).numpy())[0]
             image_grid.save(file_path, quality=85)
 
-    del generator
-    del pipeline
+    del generator, pipeline
 
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
@@ -398,31 +397,32 @@ def loss_step(
     else:
         raise ValueError(f"Unknown prediction type {noise_scheduler.config.prediction_type}")
 
-    if disc is None:
-        if guidance_scale == 0 and prior_loss_weight != 0:
-            # Chunk the noise and model_pred into two parts and compute the loss on each part separately.
-            model_pred, model_pred_prior = torch.chunk(model_pred, 2, dim=0)
-            target, target_prior = torch.chunk(target, 2, dim=0)
+    acc = (model_pred == target).float().mean()
 
-            # Compute instance loss
-            loss = F.mse_loss(model_pred.float(), target.float(), reduction="none")
+    if guidance_scale == 0 and prior_loss_weight != 0:
+        # Chunk the noise and model_pred into two parts and compute the loss on each part separately.
+        model_pred, model_pred_prior = torch.chunk(model_pred, 2, dim=0)
+        target, target_prior = torch.chunk(target, 2, dim=0)
 
-            # Compute prior loss
-            prior_loss = F.mse_loss(model_pred_prior.float(), target_prior.float(), reduction="none")
+        # Compute instance loss
+        loss = F.mse_loss(model_pred.float(), target.float(), reduction="none")
 
-            # Add the prior loss to the instance loss.
-            loss = loss + prior_loss_weight * prior_loss
-        else:
-            loss = F.mse_loss(model_pred.float(), target.float(), reduction="none")
+        # Compute prior loss
+        prior_loss = F.mse_loss(model_pred_prior.float(), target_prior.float(), reduction="none")
 
-        loss = loss.mean([1, 2, 3])
+        # Add the prior loss to the instance loss.
+        loss = loss + prior_loss_weight * prior_loss
     else:
+        loss = F.mse_loss(model_pred.float(), target.float(), reduction="none")
+
+    loss = loss.mean([1, 2, 3])
+
+    if disc is not None:
         rec_latent = get_original(noise_scheduler, model_pred, noisy_latents, timesteps)
         rec_latent /= vae.config.scaling_factor
         rec_latent = rec_latent.to(dtype=vae.dtype)
         rec = vae.decode(rec_latent).sample
         loss = 1 - disc.get_score(rec)
-        del rec_latent, rec
 
     if min_snr_gamma != 0:
         snr = compute_snr(timesteps, noise_scheduler)
@@ -432,7 +432,6 @@ def loss_step(
         loss *= mse_loss_weights
 
     loss = loss.mean()
-    acc = (model_pred == target).float().mean()
 
     return loss, acc, bsz
 
diff --git a/training/strategy/dreambooth.py b/training/strategy/dreambooth.py
index fa51bc7..4ae28b7 100644
--- a/training/strategy/dreambooth.py
+++ b/training/strategy/dreambooth.py
@@ -142,9 +142,7 @@ def dreambooth_strategy_callbacks(
             )
             pipeline.save_pretrained(checkpoint_output_dir)
 
-        del unet_
-        del text_encoder_
-        del pipeline
+        del unet_, text_encoder_, pipeline
 
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
@@ -165,8 +163,7 @@ def dreambooth_strategy_callbacks(
         unet_.to(dtype=orig_unet_dtype)
         text_encoder_.to(dtype=orig_text_encoder_dtype)
 
-        del unet_
-        del text_encoder_
+        del unet_, text_encoder_
 
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
diff --git a/training/strategy/lora.py b/training/strategy/lora.py
index 73ec8f2..1517ee8 100644
--- a/training/strategy/lora.py
+++ b/training/strategy/lora.py
@@ -140,8 +140,7 @@ def lora_strategy_callbacks(
             with open(checkpoint_output_dir / "lora_config.json", "w") as f:
                 json.dump(lora_config, f)
 
-        del unet_
-        del text_encoder_
+        del unet_, text_encoder_
 
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
@@ -153,8 +152,7 @@ def lora_strategy_callbacks(
 
         save_samples_(step=step, unet=unet_, text_encoder=text_encoder_)
 
-        del unet_
-        del text_encoder_
+        del unet_, text_encoder_
 
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
diff --git a/training/strategy/ti.py b/training/strategy/ti.py
index 08af89d..ca7cc3d 100644
--- a/training/strategy/ti.py
+++ b/training/strategy/ti.py
@@ -158,8 +158,7 @@ def textual_inversion_strategy_callbacks(
         unet_.to(dtype=orig_unet_dtype)
         text_encoder_.to(dtype=orig_text_encoder_dtype)
 
-        del unet_
-        del text_encoder_
+        del unet_, text_encoder_
 
         if torch.cuda.is_available():
             torch.cuda.empty_cache()