> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/flow_matching/adapters/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

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

* **[`nemo_automodel.components.flow_matching.adapters.base`](/nemo-automodel/nemo_automodel/components/flow_matching/adapters/base)**
* **[`nemo_automodel.components.flow_matching.adapters.flux`](/nemo-automodel/nemo_automodel/components/flow_matching/adapters/flux)**
* **[`nemo_automodel.components.flow_matching.adapters.flux2`](/nemo-automodel/nemo_automodel/components/flow_matching/adapters/flux2)**
* **[`nemo_automodel.components.flow_matching.adapters.hunyuan`](/nemo-automodel/nemo_automodel/components/flow_matching/adapters/hunyuan)**
* **[`nemo_automodel.components.flow_matching.adapters.ltx2`](/nemo-automodel/nemo_automodel/components/flow_matching/adapters/ltx2)**
* **[`nemo_automodel.components.flow_matching.adapters.qwen_image`](/nemo-automodel/nemo_automodel/components/flow_matching/adapters/qwen_image)**
* **[`nemo_automodel.components.flow_matching.adapters.qwen_image_21`](/nemo-automodel/nemo_automodel/components/flow_matching/adapters/qwen_image_21)**
* **[`nemo_automodel.components.flow_matching.adapters.simple`](/nemo-automodel/nemo_automodel/components/flow_matching/adapters/simple)**

## Package Contents

### Classes

| Name                                                                                                       | Description                                                            |
| ---------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------- |
| [`FlowMatchingContext`](#nemo_automodel-components-flow_matching-adapters-base-FlowMatchingContext)        | Context object passed to model adapters containing all necessary data. |
| [`Flux2Adapter`](#nemo_automodel-components-flow_matching-adapters-flux2-Flux2Adapter)                     | Model adapter for FLUX.2-dev image generation models.                  |
| [`FluxAdapter`](#nemo_automodel-components-flow_matching-adapters-flux-FluxAdapter)                        | Model adapter for FLUX.1 image generation models.                      |
| [`HunyuanAdapter`](#nemo_automodel-components-flow_matching-adapters-hunyuan-HunyuanAdapter)               | Model adapter for HunyuanVideo 1.5 style models.                       |
| [`LTX2Adapter`](#nemo_automodel-components-flow_matching-adapters-ltx2-LTX2Adapter)                        | Model adapter for LTX-2 dual-stream (video + audio) transformers.      |
| [`ModelAdapter`](#nemo_automodel-components-flow_matching-adapters-base-ModelAdapter)                      | Abstract base class for model-specific forward pass logic.             |
| [`QwenImage21Adapter`](#nemo_automodel-components-flow_matching-adapters-qwen_image_21-QwenImage21Adapter) | Model adapter for Qwen-Image-2.1 text-to-image models.                 |
| [`QwenImageAdapter`](#nemo_automodel-components-flow_matching-adapters-qwen_image-QwenImageAdapter)        | Model adapter for Qwen-Image text-to-image models.                     |
| [`SimpleAdapter`](#nemo_automodel-components-flow_matching-adapters-simple-SimpleAdapter)                  | Model adapter for simple transformer models (e.g., Wan).               |

### 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.flux2.Flux2Adapter(
    guidance_scale: float = 3.5,
    use_guidance_embeds: bool = True
)
```

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

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

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

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

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

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

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

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

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

```python
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), C*4].

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

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

```python
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`** `int` — default: 8

VAE compression factor (default: 8)

---

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

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

```python
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](/nemo-automodel/nemo_automodel/components/flow_matching/adapters/base#nemo_automodel-components-flow_matching-adapters-base-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)

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

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

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

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

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

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

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

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

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

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

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

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

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

```python
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), C*4].

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

```python
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`** `int` — default: 8

VAE compression factor (default: 8)

---

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

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

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

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

Example:
adapter = SimpleAdapter()
pipeline = FlowMatchingPipelineV2(model\_adapter=adapter)

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

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

## Docs

- [nemo_automodel.components.flow_matching.adapters.base](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/flow_matching/adapters/base.md)
- [nemo_automodel.components.flow_matching.adapters.flux](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/flow_matching/adapters/flux.md)
- [nemo_automodel.components.flow_matching.adapters.flux2](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/flow_matching/adapters/flux2.md)
- [nemo_automodel.components.flow_matching.adapters.hunyuan](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/flow_matching/adapters/hunyuan.md)
- [nemo_automodel.components.flow_matching.adapters.ltx2](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/flow_matching/adapters/ltx2.md)
- [nemo_automodel.components.flow_matching.adapters.qwen_image](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/flow_matching/adapters/qwen_image.md)
- [nemo_automodel.components.flow_matching.adapters.qwen_image_21](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/flow_matching/adapters/qwen_image_21.md)
- [nemo_automodel.components.flow_matching.adapters.simple](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/flow_matching/adapters/simple.md)