github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Fix audio losses for DOP and other preservation losses

Browse Source
This commit is contained in:
Jaret Burkett
2026-08-04 22:08:48 -06:00
parent 0f9094db95
commit 3afa270ab5
4 changed files with 70 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
extensions_built_in/diffusion_models/ltx2/ltx2.py
View File

@@ -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,

26
extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py
View File

@@ -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)

23
extensions_built_in/sd_trainer/SDTrainer.py
View File

@@ -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

23
toolkit/data_transfer_object/data_loader.py
View File

@@ -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