nemo_automodel.components.flow_matching.adapters.simple#

Simple transformer model adapter for FlowMatching Pipeline.

This adapter supports simple transformer models with a basic interface, such as Wan-style models.

Module Contents#

Classes#

SimpleAdapter

Model adapter for simple transformer models (e.g., Wan).

API#

class nemo_automodel.components.flow_matching.adapters.simple.SimpleAdapter#

Bases: nemo_automodel.components.flow_matching.adapters.base.ModelAdapter

Model adapter for simple transformer models (e.g., Wan).

These models use a simple interface with:

  • hidden_states: noisy latents

  • timestep: timestep values

  • encoder_hidden_states: text embeddings

Expected batch keys:

  • text_embeddings: Text encoder output [B, seq_len, dim]

.. rubric:: Example

adapter = SimpleAdapter() pipeline = FlowMatchingPipelineV2(model_adapter=adapter)

prepare_inputs(
context: nemo_automodel.components.flow_matching.adapters.base.FlowMatchingContext,
) → Dict[str, Any]#

Prepare inputs for simple transformer model.

Parameters:

context – FlowMatchingContext with batch data

Returns:

  • hidden_states: Noisy latents

  • timestep: Timestep values

  • encoder_hidden_states: Text embeddings

Return type:

Dictionary containing

forward(
model: torch.nn.Module,
inputs: Dict[str, Any],
) → torch.Tensor#

Execute forward pass for simple transformer model.

Parameters:

  • model – Transformer model

  • inputs – Dictionary from prepare_inputs()

Returns:

Model prediction tensor