nemo_automodel.components.models.qwen_image_edit.adapter

Flow-matching adapter for cached Qwen-Image-Edit-2511 conditioning.

Module Contents

Classes

Name	Description
`QwenImageEditAdapter`	Adapt cached image-edit batches to the upstream Diffusers transformer.

API

class nemo_automodel.components.models.qwen_image_edit.adapter.QwenImageEditAdapter()

Bases: ModelAdapter

Adapt cached image-edit batches to the upstream Diffusers transformer.

Qwen-Image-Edit concatenates the packed noisy target tokens followed by the packed context-image tokens. The transformer predicts every image token, but flow matching supervises only the leading target-token span.

nemo_automodel.components.models.qwen_image_edit.adapter.QwenImageEditAdapter._pack_latents(
    latents: torch.Tensor
) -> torch.Tensor

staticmethod

Pack 2-by-2 latent patches in the upstream Diffusers layout.

Parameters:

latents

torch.Tensor

Tensor of shape [batch, channels, latent_height, latent_width]. Both spatial dimensions must be even.

Returns: torch.Tensor

Tensor of shape [batch, packed_tokens, 4 * channels], where

nemo_automodel.components.models.qwen_image_edit.adapter.QwenImageEditAdapter._unpack_latents(
    packed_latents: torch.Tensor,
    height: int,
    width: int
) -> torch.Tensor

staticmethod

Restore packed predictions to a channels-first latent grid.

Parameters:

packed_latents

torch.Tensor

Tensor of shape [batch, packed_tokens, 4 * channels] in the upstream Diffusers 2-by-2 patch layout.

height

int

Output latent height. It must be even.

width

int

Output latent width. It must be even.

Returns: torch.Tensor

Tensor of shape [batch, channels, height, width]. The returned

nemo_automodel.components.models.qwen_image_edit.adapter.QwenImageEditAdapter._validate_context_latents(
    context_latents: typing.Any,
    batch_size: int,
    channels: int
) -> list[torch.Tensor]

staticmethod

Validate the ordered cached context-latent sequence.

Parameters:

context_latents

Any

Ordered list of tensors, each with shape [batch, channels, context_height, context_width].

batch_size

int

Required leading size for every context tensor.

channels

int

Required latent channel count for every context tensor.

Returns: list[torch.Tensor]

The validated list of context tensors in its original order. The

nemo_automodel.components.models.qwen_image_edit.adapter.QwenImageEditAdapter.compute_loss(
    model_pred: torch.Tensor,
    target: torch.Tensor
) -> torch.Tensor

Compute unreduced flow-matching velocity MSE in float32.

Parameters:

model_pred

torch.Tensor

Predicted target velocity tensor of shape [batch, channels, target_height, target_width].

target

torch.Tensor

Reference noise - clean_latent velocity tensor with the same shape as model_pred.

Returns: torch.Tensor

Per-element float32 squared-error tensor with the same shape as

nemo_automodel.components.models.qwen_image_edit.adapter.QwenImageEditAdapter.forward(
    model: torch.nn.Module,
    inputs: dict[str, typing.Any]
) -> torch.Tensor

Run the upstream transformer and return only target predictions.

Parameters:

model

nn.Module

Upstream Diffusers Qwen image transformer. Its packed output has shape [batch, total_image_tokens, 4 * channels].

inputs

dict[str, Any]

Mapping returned by prepare_inputs. Tensor-bearing fields have the layouts documented by that method.

Returns: torch.Tensor

Target velocity prediction tensor of shape [batch, channels,

nemo_automodel.components.models.qwen_image_edit.adapter.QwenImageEditAdapter.prepare_inputs(
    context: nemo_automodel.components.flow_matching.adapters.base.FlowMatchingContext
) -> dict[str, typing.Any]

Build Qwen edit conditioning from a cached flow-matching batch.

Parameters:

context

FlowMatchingContext

Flow context whose noisy_latents tensor has shape [batch, channels, target_height, target_width]. Its batch must contain context_latents as an ordered list of tensors with shape [batch, channels, context_height, context_width], text_embeddings with shape [batch, text_tokens, hidden], and text_attention_mask with shape [batch, text_tokens].

Returns: dict[str, Any]

Mapping containing hidden_states with shape [batch, total_image_tokens,