ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsFlow MatchingAdapters

nemo_automodel.components.flow_matching.adapters.ltx2

View as Markdown

LTX-2 dual-stream (video + audio) model adapter for the FlowMatching pipeline.

LTX-2’s transformer jointly denoises a video token stream and an audio token stream with cross-modal attention. The pipeline owns noising and loss for the primary (video) latents; this adapter additionally noises the audio latents with the SAME per-sample sigma (independent noise draw), runs the dual-stream forward, and reports the audio flow-matching loss through the auxiliary_losses hook. Both streams therefore train from a single forward and a single backward pass.

Expected batch keys (produced by tools/diffusion/processors/ltx2.py):

  • video_latents: [B, 128, F, H, W] (consumed by the pipeline)
  • audio_latents: [B, 8, L, 16]
  • text_embeddings: [B, T, D_v] video-stream connector output
  • audio_text_embeddings: [B, T, D_a] audio-stream connector output
  • text_mask: [B, T] post-connector attention mask

Module Contents

Classes

NameDescription
LTX2AdapterModel adapter for LTX-2 dual-stream (video + audio) transformers.

Functions

NameDescription
_get_forward_parametersCache parameters by the underlying forward callable.
_pack_audio_latentsPack audio latents [B, C, L, M] -> tokens [B, L, C*M].
_pack_video_latentsPack video latents [B, C, F, H, W] -> tokens [B, FHW, C] (patch size 1).
_unpack_audio_latentsUnpack audio tokens [B, L, C*M] -> latents [B, C, L, M].
_unpack_video_latentsUnpack video tokens [B, FHW, C] -> latents [B, C, F, H, W] (patch size 1).

Data

_AUDIO_LATENT_MEL_BINS

_MISSING_KEY_HINT

logger

API

class nemo_automodel.components.flow_matching.adapters.ltx2.LTX2Adapter(
    audio_loss_weight: float = 1.0,
    fps: float = 24.0
)

Bases: ModelAdapter

Model adapter for LTX-2 dual-stream (video + audio) transformers.

The adapter is stateless: all per-step tensors live in the inputs dict created fresh by each prepare_inputs() call, so gradient accumulation over micro-batches is safe.

Parameters:

audio_loss_weight
floatDefaults to 1.0

Multiplier on the audio flow-matching MSE loss added to the video loss (LTX-2 reference training uses 1.0).

fps
floatDefaults to 24.0

Video frame rate the model was trained at (LTX-2 uses 24).

nemo_automodel.components.flow_matching.adapters.ltx2.LTX2Adapter._filter_model_kwargs(
    model: torch.nn.Module,
    inputs: typing.Dict[str, typing.Any]
) -> typing.Dict[str, typing.Any]
staticmethod

Drop private stash keys and kwargs the model’s forward doesn’t accept.

Parameters:

model
nn.Module

The transformer (possibly wrapped; unwrapped via .module for signature inspection only - the wrapped module is called).

inputs
Dict[str, Any]

Full inputs dict from prepare_inputs().

Returns: Dict[str, Any]

Kwargs safe to splat into model(...).

nemo_automodel.components.flow_matching.adapters.ltx2.LTX2Adapter.auxiliary_losses(
    inputs: typing.Dict[str, typing.Any]
) -> typing.Dict[str, torch.Tensor] | None

Compute the audio flow-matching loss from stashed tensors.

Parameters:

inputs
Dict[str, Any]

Dictionary from prepare_inputs() after forward() has run; must contain _audio_pred [B, 8, L, 16] and _audio_target (float32 [B, 8, L, 16]).

Returns: Dict[str, torch.Tensor] | None

{“audio_loss”: scalar} - unweighted MSE in float32 scaled by

nemo_automodel.components.flow_matching.adapters.ltx2.LTX2Adapter.forward(
    model: torch.nn.Module,
    inputs: typing.Dict[str, typing.Any]
) -> torch.Tensor

Execute the dual-stream forward pass.

Calls the LTX-2 transformer with the kwargs from prepare_inputs() (filtered to the model’s forward signature so minor diffusers-version differences in optional kwargs don’t break the call). The raw transformer outputs ARE the flow predictions (no conversion).

Parameters:

model
nn.Module

The LTX-2 transformer (possibly FSDP-wrapped).

inputs
Dict[str, Any]

Dictionary from prepare_inputs(). Mutated: the unpacked audio prediction is stashed under _audio_pred for auxiliary_losses().

Returns: torch.Tensor

Video flow prediction [B, 128, F, H, W].

nemo_automodel.components.flow_matching.adapters.ltx2.LTX2Adapter.prepare_inputs(
    context: nemo_automodel.components.flow_matching.adapters.base.FlowMatchingContext
) -> typing.Dict[str, typing.Any]

Prepare dual-stream transformer inputs; noise the audio latents.

The video latents are already noised by the pipeline (context.noisy_latents). Audio latents are noised here with the same per-sample sigma and an independent Gaussian draw, in float32 (matching the pipeline’s video noising), then cast to the model dtype.

Parameters:

context
FlowMatchingContext

FlowMatchingContext with batch data. context.timesteps is the rescaled timestep (sigma * num_train_timesteps) [B]; context.sigma is the raw sigma in [0, 1] [B].

Returns: Dict[str, Any]

Dictionary containing the transformer kwargs:

Raises:

  • KeyError: If the batch lacks LTX-2 audio/text cache keys.
nemo_automodel.components.flow_matching.adapters.ltx2._get_forward_parameters(
    forward_callable: typing.Callable[..., typing.Any]
) -> typing.Mapping[str, inspect.Parameter] | None

Cache parameters by the underlying forward callable.

nemo_automodel.components.flow_matching.adapters.ltx2._pack_audio_latents(
    latents: torch.Tensor
) -> torch.Tensor

Pack audio latents [B, C, L, M] -> tokens [B, L, C*M].

nemo_automodel.components.flow_matching.adapters.ltx2._pack_video_latents(
    latents: torch.Tensor
) -> torch.Tensor

Pack video latents [B, C, F, H, W] -> tokens [B, FHW, C] (patch size 1).

nemo_automodel.components.flow_matching.adapters.ltx2._unpack_audio_latents(
    tokens: torch.Tensor,
    num_mel_bins: int
) -> torch.Tensor

Unpack audio tokens [B, L, C*M] -> latents [B, C, L, M].

Parameters:

tokens
torch.Tensor

Audio token sequence [B, L, C*M].

num_mel_bins
int

Latent mel bin count M (16 for LTX-2).

Returns: torch.Tensor

Audio latents [B, C, L, M].

nemo_automodel.components.flow_matching.adapters.ltx2._unpack_video_latents(
    tokens: torch.Tensor,
    num_frames: int,
    height: int,
    width: int
) -> torch.Tensor

Unpack video tokens [B, FHW, C] -> latents [B, C, F, H, W] (patch size 1).

Parameters:

tokens
torch.Tensor

Video token sequence [B, FHW, C].

num_frames
int

Latent frame count F.

height
int

Latent height H.

width
int

Latent width W.

Returns: torch.Tensor

Video latents [B, C, F, H, W].

nemo_automodel.components.flow_matching.adapters.ltx2._AUDIO_LATENT_MEL_BINS = 16
 
nemo_automodel.components.flow_matching.adapters.ltx2._MISSING_KEY_HINT = "LTX2Adapter requires '{key}' in the batch. Preprocess your dataset with the 'lt...
 
nemo_automodel.components.flow_matching.adapters.ltx2.logger = logging.getLogger(__name__)