> This page is for version Nightly (default).
> 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.flux3

Flux3 model adapter for FlowMatching Pipeline.

Supports FLUX.3 finetuning for two image workloads, dispatched from the batch:

* Text-to-image: 4D patchified + BN-normalized image latents \[B, 128, H_p, W_p]
* Image-to-image: t2i plus a clean conditioning image stream
  (`batch["context_latents"]`, position ids with l = -1, timestep 0)

Common FLUX.3 conventions (vs Flux1/2):

* Qwen3-VL text embeddings (text_embeddings, \[B, seq, 20480])
* No CLIP pooled projections and no guidance embedding (FLUX.3 is not
  guidance-distilled)
* Per-sample timesteps (\[B]) and unbatched float32 4D position ids
  (\[seq, 4], axes (t, h, w, l)); image H/W axes use the 0.5 RoPE scale

## Module Contents

### Classes

| Name                                                                                                         | Description                                                               |
| ------------------------------------------------------------------------------------------------------------ | ------------------------------------------------------------------------- |
| [`Flux3Adapter`](#nemo_automodel-components-flow_matching-adapters-flux3-Flux3Adapter)                       | Model adapter for FLUX.3 text-to-image and image-to-image.                |
| [`Flux3ConditioningInputs`](#nemo_automodel-components-flow_matching-adapters-flux3-Flux3ConditioningInputs) | Optional reference-image tokens and their position IDs.                   |
| [`Flux3Inputs`](#nemo_automodel-components-flow_matching-adapters-flux3-Flux3Inputs)                         | Image transformer inputs and the latent grid needed to unpack its output. |

### Data

[`_IMAGE_ROPE_SCALE`](#nemo_automodel-components-flow_matching-adapters-flux3-_IMAGE_ROPE_SCALE)

### API

```python
class nemo_automodel.components.flow_matching.adapters.flux3.Flux3Adapter()
```

**Bases:** [ModelAdapter](/nemo-automodel/nemo_automodel/components/flow_matching/adapters/base#nemo_automodel-components-flow_matching-adapters-base-ModelAdapter)

Model adapter for FLUX.3 text-to-image and image-to-image.

The preprocessor stores latents already in the transformer's expected space
(2x2-patchified + BN-normalized). The FlowMatchingPipeline adds noise in
that space, so this adapter only packs tokens for the transformer call and
unpacks the prediction back so MSE loss against (noise - latents) is
elementwise aligned.

Mode dispatch in `prepare_inputs`:

* 4D with `batch["context_latents"]` -> image-to-image (adds the clean
  `image_cond_*` stream: same latent space as the target, ids with l = -1,
  timestep 0; the transformer returns only generated tokens, so the loss
  covers only the generated stream)
* 4D otherwise -> text-to-image

```python
nemo_automodel.components.flow_matching.adapters.flux3.Flux3Adapter._pack_latents(
    latents: torch.Tensor
) -> torch.Tensor
```

Flatten spatial dims.

**Parameters:**

**`latents`** `torch.Tensor`

Tensor of shape \[batch, channels, height, width].

---

**Returns:** `torch.Tensor`

Tensor of shape \[batch, height\*width, channels].

```python
nemo_automodel.components.flow_matching.adapters.flux3.Flux3Adapter._prepare_image_cond_ids(
    h_p: int,
    w_p: int,
    device: torch.device
) -> torch.Tensor
```

Build 4D conditioning-image position ids: like image ids but l = -1.

**Returns:** `torch.Tensor`

Float32 tensor of shape \[h_p\*w_p, 4].

```python
nemo_automodel.components.flow_matching.adapters.flux3.Flux3Adapter._prepare_latent_ids(
    h_p: int,
    w_p: int,
    device: torch.device
) -> torch.Tensor
```

Build 4D image position ids with axes (t=0, h, w, l=0).

The H/W coordinates are scaled by the FLUX.3 image RoPE scale (0.5).
FLUX.3 position ids are unbatched and shared across the batch.

**Returns:** `torch.Tensor`

Float32 tensor of shape \[h_p\*w_p, 4].

```python
nemo_automodel.components.flow_matching.adapters.flux3.Flux3Adapter._prepare_text_ids(
    seq_len: int,
    device: torch.device
) -> torch.Tensor
```

Build 4D text position ids with axes (t=0, h=0, w=0, l=position).

**Returns:** `torch.Tensor`

Float32 tensor of shape \[seq_len, 4].

```python
nemo_automodel.components.flow_matching.adapters.flux3.Flux3Adapter._unpack_latents(
    latents: torch.Tensor,
    h: int,
    w: int
) -> torch.Tensor
```

Restore spatial dims.

**Parameters:**

**`latents`** `torch.Tensor`

Tensor of shape \[batch, height\*width, channels].

---

**`h`** `int`

Latent height.

---

**`w`** `int`

Latent width.

---

**Returns:** `torch.Tensor`

Tensor of shape \[batch, channels, height, width].

```python
nemo_automodel.components.flow_matching.adapters.flux3.Flux3Adapter.forward(
    model: torch.nn.Module,
    inputs: nemo_automodel.components.flow_matching.adapters.flux3.Flux3Inputs
) -> torch.Tensor
```

Execute the FLUX.3 transformer forward pass.

**Parameters:**

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

The Flux3Transformer2DModel to call.

---

**`inputs`** `Flux3Inputs`

Dict from `prepare_inputs()`; tensor layouts are documented
there.

---

**Returns:** `torch.Tensor`

Prediction tensor of shape \[batch, 128, H_p, W_p], the same unpacked

```python
nemo_automodel.components.flow_matching.adapters.flux3.Flux3Adapter.prepare_inputs(
    context: nemo_automodel.components.flow_matching.adapters.base.FlowMatchingContext
) -> nemo_automodel.components.flow_matching.adapters.flux3.Flux3Inputs
```

Prepare FLUX.3 transformer inputs from FlowMatchingContext.

Expects `context.noisy_latents` of shape \[batch, 128, H_p, W_p],
text embeddings in `context.batch["text_embeddings"]` of shape
\[batch, seq_len, 20480], and, for image-to-image,
`context.batch["context_latents"]` as a single-entry list of clean
conditioning latents \[batch, 128, H_c, W_c].

**Returns:** `Flux3Inputs`

Dict of transformer forward kwargs plus the private latent grid

```python
class nemo_automodel.components.flow_matching.adapters.flux3.Flux3ConditioningInputs
```

**Bases:** `typing.TypedDict`

Optional reference-image tokens and their position IDs.

**`image_cond_hidden_states`** `Tensor`

---

**`image_cond_ids`** `Tensor`

---

```python
class nemo_automodel.components.flow_matching.adapters.flux3.Flux3Inputs()
```

**Bases:** [Flux3ConditioningInputs](#nemo_automodel-components-flow_matching-adapters-flux3-Flux3ConditioningInputs)

Image transformer inputs and the latent grid needed to unpack its output.

**`_h_p`** `int`

---

**`_w_p`** `int`

---

**`encoder_hidden_states`** `Tensor`

---

**`hidden_states`** `Tensor`

---

**`img_ids`** `Tensor`

---

**`timestep`** `Tensor`

---

**`txt_ids`** `Tensor`

---

```python
nemo_automodel.components.flow_matching.adapters.flux3._IMAGE_ROPE_SCALE = 0.5
```