diff --git a/extensions_built_in/diffusion_models/ltx2/ltx2.py b/extensions_built_in/diffusion_models/ltx2/ltx2.py index 2af79da..c643a0f 100644 --- a/extensions_built_in/diffusion_models/ltx2/ltx2.py +++ b/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 - # enabled. Prior/cfg/guidance passes run under no_grad, and the - # preservation pass (diff_output_preservation, blank_prompt_preservation) - # runs with grad but after the loss, so it must not restate the audio - # prediction the primary pass stored on the batch. + # a grad-enabled prediction is the primary (loss carrying) one unless + # the trainer declared a secondary slot on the batch (prior / + # guidance-unconditional / preservation passes). Trainers that make + # several grad predictions per step (e.g. turbo rollouts) get one + # 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, diff --git a/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py b/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py index 70e0fd8..2ca4d84 100644 --- a/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py +++ b/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 - # enabled. Prior/cfg/guidance passes run under no_grad, and the - # preservation pass (diff_output_preservation, blank_prompt_preservation) - # runs with grad but after the loss, so it must not restate the audio - # prediction the primary pass stored on the batch. + # a grad-enabled prediction is the primary (loss carrying) one unless + # the trainer declared a secondary slot on the batch (prior / + # guidance-unconditional / preservation passes). Trainers that make + # several grad predictions per step (e.g. turbo rollouts) get one + # 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) diff --git a/extensions_built_in/sd_trainer/SDTrainer.py b/extensions_built_in/sd_trainer/SDTrainer.py index a7e58eb..85829e5 100644 --- a/extensions_built_in/sd_trainer/SDTrainer.py +++ b/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 diff --git a/toolkit/data_transfer_object/data_loader.py b/toolkit/data_transfer_object/data_loader.py index 5edd815..9f8e2ee 100644 --- a/toolkit/data_transfer_object/data_loader.py +++ b/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 - # separate so it cannot stomp the primary pred we backprop through. + # audio predictions from the non primary passes. Kept separate so + # 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