dew.nn.backbones.qwen_image

Qwen-Image 2.1’s transformer, as Diffusers’ QwenImage21Transformer2DModel runs it at commit 6256aa76.

One residual stream carries the image and the text together, image first. The text is the vision-language encoder’s hidden states, projected through a zero-centred RMS norm and a GELU MLP; the image is the latent, one token per latent position, projected by one linear. Every block modulates both with scales and tanh gates read from one projection of the time embedding that all blocks share, attends, and runs a SwiGLU feed-forward.

Two things set it apart from the MM-DiT families. Attention is block-causal: the text attends causally and the image attends to all of the text and to itself. And under causal_condition the text is modulated from time zero rather than from the sampled time, so its activations do not change across a walk. The rotary table spans three axes: a text token advances one shared position on all three, and the image sits at the frame position after the text on a height and width grid centred on zero.

The source takes the text padded to the longest prompt of the call and starts the image’s frame position after that padding. Here each row starts it after its own text, which is the source’s value for a prompt alone and for a batch of prompts of one length, and keeps a row’s output independent of what it is batched with.

Name	Summary
QwenImageTransformer	QwenImage21Transformer2DModel over Dew’s interface.

QwenImageTransformer

Flax module source

class QwenImageTransformer(
    in_channels: int = 64,
    out_channels: int = 64,
    num_layers: int = 32,
    heads: int = 32,
    head_dim: int = 128,
    context_in_dim: int = 4096,
    mlp_ratio: int = 3,
    axes_dims_rope: Sequence[int] = (16, 56, 56),
    eps: float = 1e-06,
    causal_condition: bool = True,
    dtype: Dtype | None = None,
    precision: PrecisionLike = None,
    attention_impl: str = 'auto',
)

QwenImage21Transformer2DModel over Dew’s interface.

__call__ takes NHWC latents, one token per position, the model time the schedule supplies - the sigma times the training count, which is the product the source reaches by dividing its timestep by a thousand and multiplying it back - and a DenoisingCondition whose context is the encoder’s text states after the system prompt, right-padded, with mask marking the real ones. Each row’s real tokens lead its text, the layout the rotary positions and the attention’s key lengths both read. It returns the prediction at the image’s tokens.