> This page is for version 0.5.0 · 26.06.
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> 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.base

Base classes and data structures for model adapters.

This module defines the abstract ModelAdapter class and the FlowMatchingContext
dataclass used to pass data between the pipeline and adapters.

## Module Contents

### Classes

| Name                                                                                                | Description                                                            |
| --------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------- |
| [`FlowMatchingContext`](#nemo_automodel-components-flow_matching-adapters-base-FlowMatchingContext) | Context object passed to model adapters containing all necessary data. |
| [`ModelAdapter`](#nemo_automodel-components-flow_matching-adapters-base-ModelAdapter)               | Abstract base class for model-specific forward pass logic.             |

### API

```python
class nemo_automodel.components.flow_matching.adapters.base.FlowMatchingContext(
    noisy_latents: torch.Tensor,
    latents: torch.Tensor,
    timesteps: torch.Tensor,
    sigma: torch.Tensor,
    task_type: str,
    data_type: str,
    device: torch.device,
    dtype: torch.dtype,
    batch: typing.Dict[str, typing.Any],
    cfg_dropout_prob: float = 0.0
)
```

Dataclass

Context object passed to model adapters containing all necessary data.

This provides a clean interface for adapters to access the data they need
without coupling to the batch dictionary structure.

**`batch`** `Dict[str, Any]`

Original batch dictionary (for model-specific data)

---

**`cfg_dropout_prob`** `float = 0.0`

Probability of dropping text embeddings (setting to 0) during
training for classifier-free guidance (CFG). Defaults to 0.0 for backward compatibility.

---

**`data_type`** `str`

"video" or "image"

---

**`device`** `device`

Device for tensor operations

---

**`dtype`** `dtype`

Data type for tensor operations

---

**`latents`** `Tensor`

\[B, C, F, H, W] for video or \[B, C, H, W] for image - Original clean latents
(also accessible via deprecated 'video\_latents' property for backward compatibility)

---

**`noisy_latents`** `Tensor`

\[B, C, F, H, W] or \[B, C, H, W] - Noisy latents after interpolation

---

**`sigma`** `Tensor`

\[B] - Sigma values

---

**`task_type`** `str`

"t2v" or "i2v"

---

**`timesteps`** `Tensor`

\[B] - Sampled timesteps

---

**`video_latents`** `Tensor`

Backward compatibility alias for 'latents' field.

---

```python
class nemo_automodel.components.flow_matching.adapters.base.ModelAdapter()
```

Abstract

Abstract base class for model-specific forward pass logic.

Implement this class to add support for new model architectures
without modifying the FlowMatchingPipeline.

The adapter pattern decouples the flow matching logic from model-specific
details like input preparation and forward pass conventions.

Example:
class MyCustomAdapter(ModelAdapter):
def prepare\_inputs(self, context: FlowMatchingContext) -> Dict\[str, Any]:
return \{
"x": context.noisy\_latents,
"t": context.timesteps,
"cond": context.batch\["my\_conditioning"],
}

def forward(self, model: nn.Module, inputs: Dict\[str, Any]) -> torch.Tensor:
return model(\*\*inputs)

pipeline = FlowMatchingPipelineV2(model\_adapter=MyCustomAdapter())

```python
nemo_automodel.components.flow_matching.adapters.base.ModelAdapter.auxiliary_losses(
    inputs: typing.Dict[str, typing.Any]
) -> typing.Dict[str, torch.Tensor] | None
```

Return additional named scalar losses computed from adapter-stashed state.

Called by `FlowMatchingPipeline.step()` after the primary loss is
computed. Override for models that train extra prediction streams
(e.g. a second modality) alongside the primary latents: stash the
needed tensors in the `inputs` dict during `prepare_inputs()` /
`forward()` and compute the losses here.

**Parameters:**

**`inputs`** `Dict[str, Any]`

The dict returned by `prepare_inputs()`, after
`forward()` has run (adapters may stash prediction tensors
in it during `forward()`).

---

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

Mapping of metric name to scalar loss tensor (each connected to

```python
nemo_automodel.components.flow_matching.adapters.base.ModelAdapter.compute_loss(
    model_pred: torch.Tensor,
    target: torch.Tensor
) -> torch.Tensor
```

Compute the unreduced flow-matching loss.

**Parameters:**

**`model_pred`** `torch.Tensor`

Tensor of shape \[batch, ...], with arbitrary trailing latent dimensions.

---

**`target`** `torch.Tensor`

Tensor of shape \[batch, ...] matching `model_pred`, containing the flow velocity target.

---

**Returns:** `torch.Tensor`

Float32 tensor of shape \[batch, ...] matching `model_pred`, containing the per-element loss.

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

abstract

Execute the model forward pass.

**Parameters:**

**`model`** `nn.Module`

The model to call

---

**`inputs`** `Dict[str, Any]`

Dictionary of inputs from prepare\_inputs()

---

**Returns:** `torch.Tensor`

Model prediction tensor

```python
nemo_automodel.components.flow_matching.adapters.base.ModelAdapter.post_process_prediction(
    model_pred: torch.Tensor
) -> torch.Tensor
```

Post-process model prediction if needed.

Override this for models that return extra outputs or need transformation.

**Parameters:**

**`model_pred`** `torch.Tensor`

Raw model output

---

**Returns:** `torch.Tensor`

Processed prediction tensor

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

abstract

Prepare model-specific inputs from the context.

**Parameters:**

**`context`** `FlowMatchingContext`

FlowMatchingContext containing all necessary data

---

**Returns:** `Dict[str, Any]`

Dictionary of inputs to pass to the model's forward method