nemo_automodel.components.flow_matching.adapters.qwen_image_21

Qwen-Image-2.1 model adapter for FlowMatching Pipeline.

Qwen-Image-2.1 is a single-stream, block-causal DiT. Unlike Qwen-Image it:

consumes 64-channel latents unpatched (one token per 16x16 pixel tile),
runs text and image tokens through one joint sequence whose layout is described by img_mask (one slot per 2x2 group of target latent tokens),
predicts every joint-sequence token; only the trailing target-image tokens are supervised.

Module Contents

Classes

Name	Description
`QwenImage21Adapter`	Model adapter for Qwen-Image-2.1 text-to-image models.

Data

_IMG_TOKENS_PER_SLOT

API

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.

nemo_automodel.components.flow_matching.adapters.qwen_image_21._IMG_TOKENS_PER_SLOT = 4