> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.flow_matching.adapters.ltx2

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

| Name                                                                                | Description                                                       |
| ----------------------------------------------------------------------------------- | ----------------------------------------------------------------- |
| [`LTX2Adapter`](#nemo_automodel-components-flow_matching-adapters-ltx2-LTX2Adapter) | Model adapter for LTX-2 dual-stream (video + audio) transformers. |

### Functions

| Name                                                                                                        | Description                                                                    |
| ----------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------ |
| [`_get_forward_parameters`](#nemo_automodel-components-flow_matching-adapters-ltx2-_get_forward_parameters) | Cache parameters by the underlying `forward` callable.                         |
| [`_pack_audio_latents`](#nemo_automodel-components-flow_matching-adapters-ltx2-_pack_audio_latents)         | Pack audio latents \[B, C, L, M] -> tokens \[B, L, C\*M].                      |
| [`_pack_video_latents`](#nemo_automodel-components-flow_matching-adapters-ltx2-_pack_video_latents)         | Pack video latents \[B, C, F, H, W] -> tokens \[B, F*H*W, C] (patch size 1).   |
| [`_unpack_audio_latents`](#nemo_automodel-components-flow_matching-adapters-ltx2-_unpack_audio_latents)     | Unpack audio tokens \[B, L, C\*M] -> latents \[B, C, L, M].                    |
| [`_unpack_video_latents`](#nemo_automodel-components-flow_matching-adapters-ltx2-_unpack_video_latents)     | Unpack video tokens \[B, F*H*W, C] -> latents \[B, C, F, H, W] (patch size 1). |

### Data

[`_AUDIO_LATENT_MEL_BINS`](#nemo_automodel-components-flow_matching-adapters-ltx2-_AUDIO_LATENT_MEL_BINS)

[`_MISSING_KEY_HINT`](#nemo_automodel-components-flow_matching-adapters-ltx2-_MISSING_KEY_HINT)

[`logger`](#nemo_automodel-components-flow_matching-adapters-ltx2-logger)

### API

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

**Bases:** [ModelAdapter](/nemo-automodel/nemo_automodel/components/flow_matching/adapters/base#nemo_automodel-components-flow_matching-adapters-base-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`** `float` — default: 1.0

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

---

**`fps`** `float` — default: 24.0

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

---

```python
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(...)`.

```python
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

```python
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].

```python
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.

```python
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.

```python
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].

```python
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, F*H*W, C] (patch size 1).

```python
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].

```python
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, F*H*W, C] -> latents \[B, C, F, H, W] (patch size 1).

**Parameters:**

**`tokens`** `torch.Tensor`

Video token sequence \[B, F*H*W, 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].

```python
nemo_automodel.components.flow_matching.adapters.ltx2._AUDIO_LATENT_MEL_BINS = 16
```

```python
nemo_automodel.components.flow_matching.adapters.ltx2._MISSING_KEY_HINT = "LTX2Adapter requires '{key}' in the batch. Preprocess your dataset with the 'lt...
```

```python
nemo_automodel.components.flow_matching.adapters.ltx2.logger = logging.getLogger(__name__)
```