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

Migrate to a new DTO for latents to carry more information that a normal tensor such as audio.

Browse Source
This commit is contained in:
Jaret Burkett
2026-08-30 10:30:21 -06:00
parent 2a69c1e7de
commit 764b5064fb
8 changed files with 369 additions and 215 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -9,6 +9,7 @@ from transformers import Gemma3Config
import yaml
from toolkit.config_modules import GenerateImageConfig, ModelConfig
from toolkit.data_transfer_object.data_loader import DataLoaderBatchDTO
from toolkit.dto import DTO
from toolkit.models.base_model import BaseModel
from toolkit.basic import flush
from toolkit.prompt_utils import PromptEmbeds
@@ -825,16 +826,7 @@ class LTX2Model(BaseModel):
batch: "DataLoaderBatchDTO" = None,
**kwargs,
):
# 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_slot is None
)
audio_target = None
with torch.no_grad():
if self.model.device == torch.device("cpu"):
self.model.to(self.device_torch)
@@ -937,22 +929,27 @@ class LTX2Model(BaseModel):
raw_audio_latents = self.encode_audio(batch.audio_data)
audio_num_frames = raw_audio_latents.shape[1]
# add the audio targets to the batch for loss calculation later
# the audio noise is drawn once per step and shared by every
# pass (prior, primary, cfg/guidance, preservation) so they all
# see the same soundtrack and the stored target keeps matching
# see the same soundtrack and every pass's target matches. It
# rides on the latents DTO.
audio_noise = (
batch.latents.get("audio_noise")
if isinstance(batch.latents, DTO)
else None
)
if (
batch.audio_noise is not None
and batch.audio_noise.shape == raw_audio_latents.shape
audio_noise is not None
and audio_noise.shape == raw_audio_latents.shape
):
audio_noise = batch.audio_noise.to(
audio_noise = audio_noise.to(
raw_audio_latents.device, dtype=raw_audio_latents.dtype
)
else:
audio_noise = torch.randn_like(raw_audio_latents)
batch.audio_noise = audio_noise
if batch.audio_target is None:
batch.audio_target = (audio_noise - raw_audio_latents).detach()
if batch.latents is not None:
batch.latents = DTO(batch.latents, audio_noise=audio_noise)
audio_target = (audio_noise - raw_audio_latents).detach()
audio_latents = self.add_noise(
raw_audio_latents,
audio_noise,
@@ -1044,13 +1041,6 @@ class LTX2Model(BaseModel):
return_dict=False,
)
# add audio latent to batch if we had audio
if batch.audio_target is not None:
if is_primary_pred:
batch.audio_pred = noise_pred_audio
else:
batch.set_secondary_audio_pred(noise_pred_audio)
unpacked_output = self.pipeline._unpack_latents(
latents=noise_pred_video,
num_frames=latent_num_frames,
@@ -1060,6 +1050,13 @@ class LTX2Model(BaseModel):
patch_size_t=self.pipeline.transformer_temporal_patch_size,
)
if audio_target is not None:
# every pass's DTO carries its own audio stream and target
return DTO(
unpacked_output,
audio=noise_pred_audio,
audio_target=audio_target,
)
return unpacked_output
def get_prompt_embeds(self, prompt: str) -> PromptEmbeds:

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

@@ -49,6 +49,7 @@ from toolkit.accelerator import unwrap_model
from toolkit.advanced_prompt_embeds import AdvancedPromptEmbeds
from toolkit.basic import flush
from toolkit.config_modules import GenerateImageConfig, ModelConfig
from toolkit.dto import DTO
from toolkit.metadata import get_meta_for_safetensors
from toolkit.models.base_model import BaseModel
from toolkit.models.v2.text_encoders.qwen3_vl import Qwen3VLTextEncoder
@@ -779,17 +780,6 @@ class MinimaxH3Model(BaseModel):
if self.model.device == torch.device("cpu"):
self.model.to(device)
# 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_slot is None
)
batch_size, _, t_lat, h_lat, w_lat = latent_model_input.shape
with torch.no_grad():
@@ -836,6 +826,8 @@ class MinimaxH3Model(BaseModel):
)
sa = sigma_a.view(-1, 1, 1)
audio_target = None
noisy_audio_rows = None
if raw_audio is not None:
expected_rows = a_lat * packing.AUDIO_CHANNELS
if raw_audio.shape[1] > expected_rows:
@@ -846,31 +838,29 @@ class MinimaxH3Model(BaseModel):
)
# the audio noise is drawn once per step and shared by every
# pass (prior, primary, cfg/guidance, preservation) so they all
# see the same soundtrack and the stored target keeps matching
if (
batch.audio_noise is not None
and batch.audio_noise.shape == raw_audio.shape
):
audio_noise = batch.audio_noise.to(device, torch.float32)
# see the same soundtrack and every pass's target matches. It
# rides on the latents DTO along with the trimmed audio so
# on-the-fly encodes aren't repeated per pass.
audio_noise = (
batch.latents.get("audio_noise")
if isinstance(batch.latents, DTO)
else None
)
if audio_noise is not None and audio_noise.shape == raw_audio.shape:
audio_noise = audio_noise.to(device, torch.float32)
else:
audio_noise = torch.randn_like(raw_audio)
batch.audio_noise = audio_noise
if batch.latents is not None:
batch.latents = DTO(
batch.latents, audio=raw_audio, audio_noise=audio_noise
)
audio_rows = (1.0 - sa) * raw_audio + sa * audio_noise
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.
# 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). 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
# model predicts clean - noise; audio_pred is negated below so
# the target follows ai-toolkit's noise - clean convention
audio_target = (audio_noise - raw_audio).detach()
# what audio perceptual losses need to rebuild the clean
# estimate (x0 = noisy - sigma_a * pred); rides the pred DTO
noisy_audio_rows = audio_rows
else:
# no soundtrack: silence (zeros) noised at the audio sigma
# rides along without contributing to the loss
@@ -966,15 +956,19 @@ class MinimaxH3Model(BaseModel):
if num_cond_audio > 0:
# reference soundtrack rows are conditioning, not targets
audio_pred = audio_pred[:, num_cond_audio:]
if batch is not None and batch.audio_target is not None:
# flip to ai-toolkit's noise - clean convention
if is_primary_pred:
batch.audio_pred = -audio_pred
else:
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)
if audio_target is not None:
# every pass's DTO carries its own audio stream; preds flipped to
# ai-toolkit's noise - clean convention
return DTO(
-noise_pred,
audio=-audio_pred,
audio_target=audio_target,
audio_noisy=noisy_audio_rows,
audio_sigma=sigma_a,
)
return -noise_pred
def get_loss_target(self, *args, **kwargs):

79
extensions_built_in/sd_trainer/SDTrainer.py
View File

@@ -16,6 +16,7 @@ from toolkit.clip_vision_adapter import ClipVisionAdapter
from toolkit.config_modules import GenerateImageConfig
from toolkit.data_loader import get_dataloader_datasets
from toolkit.data_transfer_object.data_loader import DataLoaderBatchDTO, FileItemDTO
from toolkit.dto import DTO
from toolkit.guidance import get_targeted_guidance_loss, get_guidance_loss, GuidanceType
from toolkit.image_utils import show_tensors, show_latents
from toolkit.ip_adapter import IPAdapter
@@ -541,6 +542,13 @@ class SDTrainer(BaseSDTrainProcess):
target_mask_multiplier = None
dtype = get_torch_dtype(self.train_config.dtype)
# joint audio models return a DTO: video pred as the tensor, the audio
# pred and its target riding as extras on each pass's own prediction.
# Grab them before any math rebinds noise_pred to a plain tensor.
audio_pred = noise_pred.get('audio') if isinstance(noise_pred, DTO) else None
audio_target = noise_pred.get('audio_target') if isinstance(noise_pred, DTO) else None
audio_sigma = noise_pred.get('audio_sigma') if isinstance(noise_pred, DTO) else None
has_mask = batch.mask_tensor is not None
with torch.no_grad():
@@ -625,6 +633,9 @@ class SDTrainer(BaseSDTrainProcess):
# matching adapter prediction
target = prior_pred
if getattr(self.sd, 'dopsd_self_ref', False):
if isinstance(prior_pred, DTO) and prior_pred.get('audio') is not None:
# the teacher's audio prediction is the audio target too
audio_target = prior_pred.get('audio').detach()
# D-OPSD bleed: also train against the normal (non-teacher) target
if hasattr(self.sd, 'get_loss_target'):
dopsd_normal_target = self.sd.get_loss_target(
@@ -749,9 +760,6 @@ 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,
@@ -759,7 +767,8 @@ class SDTrainer(BaseSDTrainProcess):
unconditional_embeds=None,
batch=batch,
)
batch.audio_pred_slot = None
# joint audio models: this pass's DTO carries its own audio pred
audio_uncond = unconditional_target.get('audio') if isinstance(unconditional_target, DTO) else None
is_video = len(target.shape) == 5
if self.train_config.do_guidance_loss_cfg_zero:
@@ -797,17 +806,17 @@ class SDTrainer(BaseSDTrainProcess):
unconditional_target = unconditional_target * alpha
target = unconditional_target + guidance_scale * (target - unconditional_target)
# joint audio models (ltx2, minimax_h3, flux3) carry their audio
# target/pred on the batch. Extrapolate the audio target the
# same way so the audio stream trains contrastively as well.
audio_uncond = getattr(batch, 'audio_pred_uncond', None)
if batch.audio_target is not None and audio_uncond is not None:
audio_target = batch.audio_target.float()
# joint audio models carry their audio pred/target on the pred
# DTOs. Extrapolate the audio target the same way so the audio
# stream trains contrastively as well.
if audio_target is not None and audio_uncond is not None:
a_dtype = audio_target.dtype
a_target = audio_target.float()
audio_uncond = audio_uncond.float()
audio_dims = [1] * (audio_target.dim() - 1)
audio_dims = [1] * (a_target.dim() - 1)
if self.train_config.do_guidance_loss_cfg_zero:
batch_size = audio_target.shape[0]
a_pos_flat = audio_target.view(batch_size, -1)
batch_size = a_target.shape[0]
a_pos_flat = a_target.view(batch_size, -1)
a_neg_flat = audio_uncond.view(batch_size, -1)
a_dot = torch.sum(a_pos_flat * a_neg_flat, dim=1, keepdim=True)
a_squared_norm = torch.sum(a_neg_flat ** 2, dim=1, keepdim=True) + 1e-8
@@ -816,22 +825,22 @@ class SDTrainer(BaseSDTrainProcess):
audio_guidance_scale = self._guidance_loss_target_batch
if isinstance(audio_guidance_scale, list):
audio_guidance_scale = torch.tensor(audio_guidance_scale).to(
audio_target.device, dtype=audio_target.dtype
a_target.device, dtype=a_target.dtype
).view(-1, *audio_dims)
if self.train_config.guidance_loss_schedule == 'sigma':
# audio streams can run on their own remapped sigma
audio_sigma = getattr(batch, 'audio_sigma', None)
if audio_sigma is None:
audio_sigma = timesteps / 1000.0
audio_sigma = audio_sigma.to(
audio_target.device, dtype=audio_target.dtype
a_sigma = audio_sigma
if a_sigma is None:
a_sigma = timesteps / 1000.0
a_sigma = a_sigma.to(
a_target.device, dtype=a_target.dtype
).view(-1, *audio_dims)
audio_guidance_scale = 1.0 + (audio_guidance_scale - 1.0) * audio_sigma
audio_guidance_scale = 1.0 + (audio_guidance_scale - 1.0) * a_sigma
batch.audio_target = (
audio_uncond + audio_guidance_scale * (audio_target - audio_uncond)
).to(batch.audio_target.dtype).detach()
audio_target = (
audio_uncond + audio_guidance_scale * (a_target - audio_uncond)
).to(a_dtype).detach()
if self.train_config.do_differential_guidance:
with torch.no_grad():
@@ -865,7 +874,6 @@ class SDTrainer(BaseSDTrainProcess):
# we also denoise as the unaugmented tensor is not a noisy diffirental
with torch.no_grad():
unaugmented_latents = self.sd.encode_images(batch.unaugmented_tensor).to(self.device_torch, dtype=dtype)
unaugmented_latents = unaugmented_latents * self.train_config.latent_multiplier
target = unaugmented_latents.detach()
# Get the target for loss depending on the prediction type
@@ -1040,8 +1048,8 @@ class SDTrainer(BaseSDTrainProcess):
loss = loss.mean()
# check for audio loss
if batch.audio_pred is not None and batch.audio_target is not None:
audio_loss = torch.nn.functional.mse_loss(batch.audio_pred.float(), batch.audio_target.float(), reduction="mean")
if audio_pred is not None and audio_target is not None:
audio_loss = torch.nn.functional.mse_loss(audio_pred.float(), audio_target.float(), reduction="mean")
audio_loss = audio_loss * self.train_config.audio_loss_multiplier
self.additional_logs['loss/img'] = loss.item()
self.additional_logs['loss/audio'] = audio_loss.item()
@@ -2035,10 +2043,6 @@ class SDTrainer(BaseSDTrainProcess):
)
batch.dopsd_teacher_pass = True
# 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,
@@ -2051,13 +2055,10 @@ class SDTrainer(BaseSDTrainProcess):
unconditional_embeds=unconditional_embeds,
conditioned_prompts=conditioned_prompts
)
batch.audio_pred_slot = None
if is_dopsd:
batch.dopsd_teacher_pass = False
if batch.audio_pred_prior is not None:
# the teacher's audio prediction is the audio target too
batch.audio_target = batch.audio_pred_prior.detach()
if prior_pred is not None:
# a DTO prior pred keeps its audio extras through detach
prior_pred = prior_pred.detach()
# do the custom adapter after the prior prediction
@@ -2235,7 +2236,6 @@ 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,
@@ -2244,7 +2244,6 @@ 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()
@@ -2253,10 +2252,12 @@ class SDTrainer(BaseSDTrainProcess):
# 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_pres = preservation_pred.get('audio') if isinstance(preservation_pred, DTO) else None
audio_prior = prior_pred.get('audio') if isinstance(prior_pred, DTO) else None
if audio_pres is not None and audio_prior is not None:
audio_preservation_loss = torch.nn.functional.mse_loss(
batch.audio_pred_preservation.float(),
batch.audio_pred_prior.float(),
audio_pres.float(),
audio_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