1 files changed, 17 insertions, 18 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