Fix audio losses for DOP and other preservation losses

extensions_built_in/diffusion_models/ltx2/ltx2.py

@@ -852,13 +852,15 @@ class LTX2Model(BaseModel):
         batch: "DataLoaderBatchDTO" = None,
         **kwargs,
     ):
-        # the primary (loss carrying) prediction is the first one made with grad
+        # a grad-enabled prediction is the primary (loss carrying) one unless
-        # enabled. Prior/cfg/guidance passes run under no_grad, and the
+        # the trainer declared a secondary slot on the batch (prior /
-        # preservation pass (diff_output_preservation, blank_prompt_preservation)
+        # guidance-unconditional / preservation passes). Trainers that make
-        # runs with grad but after the loss, so it must not restate the audio
+        # several grad predictions per step (e.g. turbo rollouts) get one
-        # prediction the primary pass stored on the batch.
+        # primary per prediction, last writer wins.
         is_primary_pred = (
-            torch.is_grad_enabled() and batch is not None and batch.audio_pred is None
+            torch.is_grad_enabled()
+            and batch is not None
+            and batch.audio_pred_slot is None
         )
         with torch.no_grad():
             if self.model.device == torch.device("cpu"):
@@ -1064,7 +1066,7 @@ class LTX2Model(BaseModel):
             if is_primary_pred:
                 batch.audio_pred = noise_pred_audio
             else:
-                batch.audio_pred_uncond = noise_pred_audio
+                batch.set_secondary_audio_pred(noise_pred_audio)
 
         unpacked_output = self.pipeline._unpack_latents(
             latents=noise_pred_video,

extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py

@@ -719,13 +719,15 @@ class MinimaxH3Model(BaseModel):
         if self.model.device == torch.device("cpu"):
             self.model.to(device)
 
-        # the primary (loss carrying) prediction is the first one made with grad
+        # a grad-enabled prediction is the primary (loss carrying) one unless
-        # enabled. Prior/cfg/guidance passes run under no_grad, and the
+        # the trainer declared a secondary slot on the batch (prior /
-        # preservation pass (diff_output_preservation, blank_prompt_preservation)
+        # guidance-unconditional / preservation passes). Trainers that make
-        # runs with grad but after the loss, so it must not restate the audio
+        # several grad predictions per step (e.g. turbo rollouts) get one
-        # prediction the primary pass stored on the batch.
+        # primary per prediction, last writer wins.
         is_primary_pred = (
-            torch.is_grad_enabled() and batch is not None and batch.audio_pred is None
+            torch.is_grad_enabled()
+            and batch is not None
+            and batch.audio_pred_slot is None
         )
 
         batch_size, _, t_lat, h_lat, w_lat = latent_model_input.shape
@@ -809,10 +811,16 @@ class MinimaxH3Model(BaseModel):
                 batch.audio_latents = raw_audio
                 if batch.audio_target is None:
                     # model predicts clean - noise; audio_pred is negated below
-                    # so the stored target follows ai-toolkit's noise - clean
+                    # so the stored target follows ai-toolkit's noise - clean.
+                    # With the shared noise this is the same value on every
+                    # pass, so first writer is fine (and it keeps a guidance
+                    # extrapolated target from being overwritten).
                     batch.audio_target = (audio_noise - raw_audio).detach()
+                if is_primary_pred:
                     # expose what audio perceptual losses need to rebuild the
-                    # clean estimate (x0 = noisy - sigma_a * pred)
+                    # clean estimate (x0 = noisy - sigma_a * pred). Tied to the
+                    # primary pass so they always match audio_pred, even when a
+                    # trainer makes primary predictions at several sigmas.
                     batch.audio_noisy = audio_rows
                     batch.audio_sigma = sigma_a
             else:
@@ -893,7 +901,7 @@ class MinimaxH3Model(BaseModel):
             if is_primary_pred:
                 batch.audio_pred = -audio_pred
             else:
-                batch.audio_pred_uncond = -audio_pred
+                batch.set_secondary_audio_pred(-audio_pred)
 
         video_pred = video_pred[:, num_cond:]
         noise_pred = unpatchify_video_tokens(video_pred, t_lat, h_lat, w_lat)

extensions_built_in/sd_trainer/SDTrainer.py

@@ -700,6 +700,9 @@ class SDTrainer(BaseSDTrainProcess):
                 unconditional_embeds = concat_prompt_embeds(
                     [self.unconditional_embeds] * noisy_latents.shape[0],
                 )
+                # joint audio models route this pass's audio pred to its own
+                # slot so it cannot stomp the primary pred on the batch
+                batch.audio_pred_slot = 'audio_pred_uncond'
                 unconditional_target = self.predict_noise(
                     noisy_latents=noisy_latents,
                     timesteps=timesteps,
@@ -707,6 +710,7 @@ class SDTrainer(BaseSDTrainProcess):
                     unconditional_embeds=None,
                     batch=batch,
                 )
+                batch.audio_pred_slot = None
                 is_video = len(target.shape) == 5
                 
                 if self.train_config.do_guidance_loss_cfg_zero:
@@ -1910,6 +1914,10 @@ class SDTrainer(BaseSDTrainProcess):
                                 [blank_embeds] * noisy_latents.shape[0]
                             )
                         
+                        # joint audio models stash their audio pred on the batch.
+                        # Give this pass its own slot so the preservation loss
+                        # can pair it with the preservation pass below.
+                        batch.audio_pred_slot = 'audio_pred_prior'
                         prior_pred = self.get_prior_prediction(
                             noisy_latents=noisy_latents,
                             conditional_embeds=prior_embeds_to_use,
@@ -1922,6 +1930,7 @@ class SDTrainer(BaseSDTrainProcess):
                             unconditional_embeds=unconditional_embeds,
                             conditioned_prompts=conditioned_prompts
                         )
+                        batch.audio_pred_slot = None
                         if prior_pred is not None:
                             prior_pred = prior_pred.detach()
 
@@ -2110,6 +2119,7 @@ class SDTrainer(BaseSDTrainProcess):
                                 preservation_embeds = concat_prompt_embeds(
                                     [blank_embeds] * noisy_latents.shape[0]
                                 )
+                        batch.audio_pred_slot = 'audio_pred_preservation'
                         preservation_pred = self.predict_noise(
                             noisy_latents=noisy_latents.to(self.device_torch, dtype=dtype),
                             timesteps=timesteps,
@@ -2118,10 +2128,23 @@ class SDTrainer(BaseSDTrainProcess):
                             batch=batch,
                             **pred_kwargs
                         )
+                        batch.audio_pred_slot = None
                         multiplier = self.train_config.diff_output_preservation_multiplier if self.train_config.diff_output_preservation else self.train_config.blank_prompt_preservation_multiplier
                         preservation_loss = torch.nn.functional.mse_loss(preservation_pred, prior_pred) * multiplier
                         self.additional_logs['loss/normal'] = loss.item()
                         self.additional_logs['loss/preservation'] = preservation_loss.item()
+
+                        # preserve the audio stream of joint audio models too.
+                        # Both passes ran on the same noisy audio, so this holds
+                        # the audio branch to its base model output.
+                        if batch.audio_pred_preservation is not None and batch.audio_pred_prior is not None:
+                            audio_preservation_loss = torch.nn.functional.mse_loss(
+                                batch.audio_pred_preservation.float(),
+                                batch.audio_pred_prior.float(),
+                            ) * multiplier * self.train_config.audio_loss_multiplier
+                            self.additional_logs['loss/preservation_audio'] = audio_preservation_loss.item()
+                            preservation_loss = preservation_loss + audio_preservation_loss
+
                         loss = loss + preservation_loss
 
                 # check if nan

toolkit/data_transfer_object/data_loader.py

@@ -225,9 +225,17 @@ class DataLoaderBatchDTO:
             # prediction. Secondary passes (cfg / guidance loss / prior preds)
             # reuse it so their noisy audio matches the stored audio_target.
             self.audio_noise: Union[torch.Tensor, None] = None
-            # audio prediction from an unconditional/secondary pass. Kept
+            # audio predictions from the non primary passes. Kept separate so
-            # separate so it cannot stomp the primary pred we backprop through.
+            # they cannot stomp the primary pred we backprop through.
             self.audio_pred_uncond: Union[torch.Tensor, None] = None
+            self.audio_pred_prior: Union[torch.Tensor, None] = None
+            self.audio_pred_preservation: Union[torch.Tensor, None] = None
+            # which of the above the current secondary pass writes to. None (the
+            # default) means no secondary pass is in flight: any grad-enabled
+            # prediction is a primary one and writes audio_pred (and the
+            # noisy/sigma bookkeeping) directly. The trainer sets this around
+            # its prior / guidance-unconditional / preservation passes.
+            self.audio_pred_slot: Union[str, None] = None
             # noisy audio rows and audio sigma of the primary pass, used to
             # rebuild the clean audio estimate for perceptual losses
             self.audio_noisy: Union[torch.Tensor, None] = None
@@ -491,6 +499,15 @@ class DataLoaderBatchDTO:
     ):
         return [x.caption_short for x in self.file_items]
 
+    def set_secondary_audio_pred(self, pred):
+        """Route an audio prediction from a non primary pass (prior,
+        unconditional/guidance, preservation) to its own slot so it cannot
+        stomp the primary prediction the loss backprops through. Passes that
+        did not declare a slot (e.g. a trainer's extra no_grad prediction)
+        are simply not stored."""
+        if self.audio_pred_slot is not None:
+            setattr(self, self.audio_pred_slot, pred)
+
     def cleanup(self):
         del self.latents
         del self.tensor
@@ -501,6 +518,8 @@ class DataLoaderBatchDTO:
         del self.audio_pred
         del self.audio_noise
         del self.audio_pred_uncond
+        del self.audio_pred_prior
+        del self.audio_pred_preservation
         del self.audio_noisy
         del self.audio_sigma
         del self.first_frame_latents