Skip to navigation
DevelopmentAPI ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsFlow Matchingnemo_automodel.components.flow_matching.adapters

nemo_automodel.components.flow_matching.adapters

View as Markdown

Model adapters for FlowMatching Pipeline.

This module provides model-specific adapters that decouple the flow matching logic from model-specific implementation details.

Available Adapters:

  • ModelAdapter: Abstract base class for all adapters
  • HunyuanAdapter: For HunyuanVideo 1.5 style models
  • SimpleAdapter: For simple transformer models (e.g., Wan)
  • FluxAdapter: For FLUX.1 text-to-image models
  • QwenImageAdapter: For Qwen-Image text-to-image models
  • QwenImage21Adapter: For Qwen-Image-2.1 text-to-image models
  • LTX2Adapter: For LTX-2 dual-stream video+audio models

Usage: from automodel.flow_matching.adapters import HunyuanAdapter, SimpleAdapter, FluxAdapter

Or import the base class to create custom adapters

from automodel.flow_matching.adapters import ModelAdapter

Submodules

Package Contents

Classes

NameDescription
FlowMatchingContextContext object passed to model adapters containing all necessary data.
Flux2AdapterModel adapter for FLUX.2-dev image generation models.
FluxAdapterModel adapter for FLUX.1 image generation models.
HunyuanAdapterModel adapter for HunyuanVideo 1.5 style models.
LTX2AdapterModel adapter for LTX-2 dual-stream (video + audio) transformers.
ModelAdapterAbstract base class for model-specific forward pass logic.
QwenImage21AdapterModel adapter for Qwen-Image-2.1 text-to-image models.
QwenImageAdapterModel adapter for Qwen-Image text-to-image models.
SimpleAdapterModel adapter for simple transformer models (e.g., Wan).

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

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.

class nemo_automodel.components.flow_matching.adapters.flux2.Flux2Adapter(
    guidance_scale: float = 3.5,
    use_guidance_embeds: bool = True
)

Bases: ModelAdapter

Model adapter for FLUX.2-dev image generation models.

The preprocessor stores latents already patchified (2x2 spatial patch → channel) and BN-normalized using vae.bn running statistics. The FlowMatchingPipeline adds noise in that space, so this adapter only needs to flatten the spatial dims (pack) before the transformer call and reshape back (unpack) after.

Batch format expected from multiresolution dataloader:

  • image_latents: patchified + BN-normalized latents [B, 128, H_p, W_p]
  • text_embeddings: Mistral3 stacked embeddings [B, seq_len, 15360]

FLUX.2 transformer forward interface:

  • hidden_states: packed latents [B, H_p*W_p, 128]
  • encoder_hidden_states: text embeddings [B, seq_len, 15360]
  • timestep: normalized [0, 1]
  • img_ids: 4D spatial coords [B, H_p*W_p, 4]
  • txt_ids: 4D text coords [B, seq_len, 4]
  • guidance: guidance scale [B]
nemo_automodel.components.flow_matching.adapters.flux2.Flux2Adapter._pack_latents(
    latents: torch.Tensor
) -> torch.Tensor

Flatten spatial dims: [B, C, H, W] -> [B, H*W, C].

nemo_automodel.components.flow_matching.adapters.flux2.Flux2Adapter._prepare_latent_ids(
    h_p: int,
    w_p: int,
    batch_size: int,
    device: torch.device
) -> torch.Tensor

Build 4D image positional IDs with coords (T=0, h, w, L=0).

Returns [B, H_p*W_p, 4] long tensor.

nemo_automodel.components.flow_matching.adapters.flux2.Flux2Adapter._prepare_text_ids(
    seq_len: int,
    batch_size: int,
    device: torch.device
) -> torch.Tensor

Build 4D text positional IDs with coords (T=0, H=0, W=0, position).

Returns [B, seq_len, 4] long tensor.

nemo_automodel.components.flow_matching.adapters.flux2.Flux2Adapter._unpack_latents(
    latents: torch.Tensor,
    h: int,
    w: int
) -> torch.Tensor

Restore spatial dims: [B, seq, C] -> [B, C, H, W].

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

Execute FLUX.2 transformer forward pass.

Returns prediction in [B, 128, H_p, W_p] — same space as the stored patchified latents, so MSE loss against (noise - latents) is correct.

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

Prepare inputs for FLUX.2 transformer from FlowMatchingContext.

Expects 4D patchified+BN-normalized latents [B, 128, H_p, W_p].

class nemo_automodel.components.flow_matching.adapters.flux.FluxAdapter(
    guidance_scale: float = 3.5,
    use_guidance_embeds: bool = True
)

Bases: ModelAdapter

Model adapter for FLUX.1 image generation models.

Supports batch format from multiresolution dataloader:

  • image_latents: [B, C, H, W] for images
  • text_embeddings: T5 embeddings [B, seq_len, 4096]
  • pooled_prompt_embeds: CLIP pooled [B, 768]

FLUX model forward interface:

  • hidden_states: Packed latents
  • encoder_hidden_states: T5 text embeddings
  • pooled_projections: CLIP pooled embeddings
  • timestep: Normalized timesteps [0, 1]
  • img_ids / txt_ids: Positional embeddings
nemo_automodel.components.flow_matching.adapters.flux.FluxAdapter._pack_latents(
    latents: torch.Tensor
) -> torch.Tensor

Pack latents from [B, C, H, W] to Flux format [B, (H//2)(W//2), C4].

Flux uses a 2x2 patch embedding, so latents are reshaped accordingly.

nemo_automodel.components.flow_matching.adapters.flux.FluxAdapter._prepare_latent_image_ids(
    batch_size: int,
    height: int,
    width: int,
    device: torch.device,
    dtype: torch.dtype
) -> torch.Tensor

Prepare positional IDs for image latents.

Returns tensor of shape [B, (H//2)*(W//2), 3] containing (batch_idx, y, x).

nemo_automodel.components.flow_matching.adapters.flux.FluxAdapter._unpack_latents(
    latents: torch.Tensor,
    height: int,
    width: int,
    vae_scale_factor: int = 8
) -> torch.Tensor
staticmethod

Unpack latents from Flux format back to [B, C, H, W].

Parameters:

latents
torch.Tensor

Packed latents of shape [B, num_patches, channels]

height
int

Original image height in pixels

width
int

Original image width in pixels

vae_scale_factor
intDefaults to 8

VAE compression factor (default: 8)

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

Execute forward pass for Flux model.

Returns unpacked prediction in [B, C, H, W] format.

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

Prepare inputs for Flux model from FlowMatchingContext.

Expects 4D image latents: [B, C, H, W]

class nemo_automodel.components.flow_matching.adapters.hunyuan.HunyuanAdapter(
    default_image_embed_shape: typing.Tuple[int, int] = (729, 1152),
    use_condition_latents: bool = True
)

Bases: ModelAdapter

Model adapter for HunyuanVideo 1.5 style models.

These models use:

  • Condition latents concatenated with noisy latents
  • Dual text encoders with attention masks
  • Image embeddings for i2v

Expected batch keys:

  • text_embeddings: Primary text encoder output [B, seq_len, dim]
  • text_mask: Attention mask for primary encoder [B, seq_len] (optional)
  • text_embeddings_2: Secondary text encoder output [B, seq_len, dim] (optional)
  • text_mask_2: Attention mask for secondary encoder [B, seq_len] (optional)
  • image_embeds: Image embeddings for i2v [B, seq_len, dim] (optional)

Example: adapter = HunyuanAdapter() pipeline = FlowMatchingPipelineV2(model_adapter=adapter)

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

Execute forward pass for HunyuanVideo model.

Parameters:

model
nn.Module

HunyuanVideo model

inputs
Dict[str, Any]

Dictionary from prepare_inputs()

Returns: torch.Tensor

Model prediction tensor

nemo_automodel.components.flow_matching.adapters.hunyuan.HunyuanAdapter.get_condition_latents(
    latents: torch.Tensor,
    task_type: str
) -> torch.Tensor

Generate conditional latents based on task type.

Parameters:

latents
torch.Tensor

Input latents [B, C, F, H, W]

task_type
str

Task type (“t2v” or “i2v”)

Returns: torch.Tensor

Conditional latents [B, C+1, F, H, W]

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

Prepare inputs for HunyuanVideo model.

Parameters:

context
FlowMatchingContext

FlowMatchingContext with batch data

Returns: Dict[str, Any]

Dictionary containing:

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

class nemo_automodel.components.flow_matching.adapters.qwen_image_21.QwenImage21Adapter()

Bases: ModelAdapter

Model adapter for Qwen-Image-2.1 text-to-image models.

Supports batch format from multiresolution dataloader:

  • image_latents: [B, 64, H, W]
  • text_embeddings: Qwen3-VL embeddings [B, seq_len, 4096], right-padded
  • text_attention_mask: optional [B, seq_len] bool marking valid text tokens

Qwen-Image-2.1 transformer forward interface:

  • hidden_states: Flattened latents [B, H*W, 64]
  • encoder_hidden_states: Text embeddings [B, text_len, 4096]
  • timestep: Normalized timesteps [0, 1]
  • img_shapes: [[(1, H, W)]] per sample
  • img_mask: [B, text_len + H*W/4] bool, True at target-image slots

The transformer lays out RoPE from row 0 of img_mask for the whole batch, and every text position (padding included) advances the target image’s frame index. Each sample is therefore run as its own unpadded transformer call, so it sees exactly the positions of single-prompt inference.

nemo_automodel.components.flow_matching.adapters.qwen_image_21.QwenImage21Adapter._build_img_mask(
    batch_size: int,
    text_len: int,
    image_tokens: int,
    device: torch.device
) -> torch.Tensor
staticmethod

Text positions first, then one slot per 2x2 group of target latent tokens.

nemo_automodel.components.flow_matching.adapters.qwen_image_21.QwenImage21Adapter._pack_latents(
    latents: torch.Tensor
) -> torch.Tensor
staticmethod

Flatten latents from [B, C, H, W] to [B, H*W, C] (no patching in 2.1).

nemo_automodel.components.flow_matching.adapters.qwen_image_21.QwenImage21Adapter._unpack_latents(
    latents: torch.Tensor,
    height: int,
    width: int
) -> torch.Tensor
staticmethod

Restore [B, H*W, C] token predictions to [B, C, H, W].

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

Execute forward pass for Qwen-Image-2.1 model.

Runs one transformer call per sample, trimmed to that sample’s prompt length, and returns the target-image prediction in [B, C, H, W] format.

The call count is always the local batch size, never the number of distinct prompt lengths: under FSDP every call issues collectives, so all ranks must make the same number of calls, and the diffusion sampler gives every rank the same local batch size.

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

Prepare inputs for Qwen-Image-2.1 model from FlowMatchingContext.

Expects 4D image latents: [B, C, H, W] with even H and W.

class nemo_automodel.components.flow_matching.adapters.qwen_image.QwenImageAdapter(
    guidance_scale: float = 3.5,
    use_guidance_embeds: bool = False
)

Bases: ModelAdapter

Model adapter for Qwen-Image text-to-image models.

Supports batch format from multiresolution dataloader:

  • image_latents: [B, C, H, W]
  • text_embeddings: Qwen2 embeddings [B, seq_len, hidden_dim]

Qwen-Image transformer forward interface:

  • hidden_states: Packed latents [B, num_patches, C*4]
  • encoder_hidden_states: Qwen2 text embeddings [B, seq_len, hidden_dim]
  • encoder_hidden_states_mask: Attention mask (None for flash attention)
  • timestep: Normalized timesteps [0, 1]
  • img_shapes: List of image shape tuples [[(1, h//2, w//2)]] per sample
  • guidance: Optional guidance scale embedding [B]
nemo_automodel.components.flow_matching.adapters.qwen_image.QwenImageAdapter._pack_latents(
    latents: torch.Tensor
) -> torch.Tensor

Pack latents from [B, C, H, W] to [B, (H//2)(W//2), C4].

Uses 2x2 patch grouping to match the transformer’s patch embedding.

nemo_automodel.components.flow_matching.adapters.qwen_image.QwenImageAdapter._unpack_latents(
    latents: torch.Tensor,
    height: int,
    width: int,
    vae_scale_factor: int = 8
) -> torch.Tensor
staticmethod

Unpack latents from [B, num_patches, channels] back to [B, C, H, W].

Parameters:

latents
torch.Tensor

Packed latents of shape [B, num_patches, channels]

height
int

Original image height in pixels

width
int

Original image width in pixels

vae_scale_factor
intDefaults to 8

VAE compression factor (default: 8)

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

Execute forward pass for Qwen-Image model.

Returns unpacked prediction in [B, C, H, W] format.

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

Prepare inputs for Qwen-Image model from FlowMatchingContext.

Expects 4D image latents: [B, C, H, W]

class nemo_automodel.components.flow_matching.adapters.simple.SimpleAdapter()

Bases: 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]

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

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

Execute forward pass for simple transformer model.

Parameters:

model
nn.Module

Transformer model

inputs
Dict[str, Any]

Dictionary from prepare_inputs()

Returns: torch.Tensor

Model prediction tensor

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

Prepare inputs for simple transformer model.

Parameters:

context
FlowMatchingContext

FlowMatchingContext with batch data

Returns: Dict[str, Any]

Dictionary containing: