Skip to navigation
ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsFlow MatchingAdaptersnemo_automodel.components.flow_matching.adapters.flux3

nemo_automodel.components.flow_matching.adapters.flux3

View as Markdown

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

NameDescription
Flux3AdapterModel adapter for FLUX.3 text-to-image and image-to-image.
Flux3ConditioningInputsOptional reference-image tokens and their position IDs.
Flux3InputsImage transformer inputs and the latent grid needed to unpack its output.

Data

_IMAGE_ROPE_SCALE

API

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

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

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

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

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

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

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

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

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
class nemo_automodel.components.flow_matching.adapters.flux3.Flux3Inputs()

Bases: 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
nemo_automodel.components.flow_matching.adapters.flux3._IMAGE_ROPE_SCALE = 0.5