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`	Context object passed to model adapters containing all necessary data.
`ModelAdapter`	Abstract base class for model-specific forward pass logic.

API

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 for tensor operations

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.

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

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

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.

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

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

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