Diffusers documentation

QwenImage21Transformer2DModel

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.42.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

QwenImage21Transformer2DModel

Qwen-Image 2.1 使用的单流 Transformer。文本和图像潜变量共享同一个序列,且单个共享的 modulation 投影为每个块提供输入。

2.1 与早期 QwenImage 转换器在两种行为上有所不同:

  • 块因果注意力 — 注意力遵循 (q_idx >= kv_idx) or same_image_block,因此联合序列是因果的,而每个图像块内部保持双向。 QwenImage21AttnProcessor 将其实现为每个前缀段一次注意力调用,并且是默认选项。 QwenImage21FlexAttnProcessor 将其实现为一次由 BlockMask 驱动的 flex_attention 调用,在模型编译后速度更快。两者产生的结果相同。
  • causal_condition — 文本和条件图像 token 的调制从 t = 0 开始,而非从采样的时间步开始。 它们的激活值与去噪步骤无关,因此该前缀的键和值可以通过 kv_cache 参数在步骤之间进行缓存。

使用以下代码加载:

import torch
from diffusers import QwenImage21Transformer2DModel

transformer = QwenImage21Transformer2DModel.from_pretrained(
    "Qwen/Qwen-Image-2.1", subfolder="transformer", dtype=torch.bfloat16
)

QwenImage21Transformer2DModel

类 diffusers.QwenImage21Transformer2DModel

< >

( patch_size: int = 1in_channels: int = 64out_channels: int | None = 64num_layers: int = 32attention_head_dim: int = 128num_attention_heads: int = 32context_in_dim: int = 4096mlp_ratio: int = 3axes_dims_rope: tuple = (16, 56, 56)eps: float = 1e-06causal_condition: bool = True )

参数

  • patch_size (int,默认为 1) — 折叠到通道维的潜在块(latent patch)的边长。2.1 版本在未分块(unpatched)的情况下消费潜在表示。
  • in_channels (int,默认为 64) — 输入潜空间的通道数。
  • out_channels (int,可选,默认为 64) — 输出的潜在通道数。若未指定,则回退为 in_channels。
  • num_layers (int,默认为 32) — 单流块的数量。
  • attention_head_dim (int,默认为 128) — 每个注意力头的通道数。
  • num_attention_heads (int,默认为 32) — 注意力头的数量。
  • context_in_dim (int,默认为 4096) — encoder_hidden_states 的通道维度。
  • mlp_ratio (int,默认为 3) — 前馈扩展因子。
  • axes_dims_rope (tuple[int],默认为 (16, 56, 56)) — 帧、高度和宽度轴的旋转维度。
  • eps (float,默认为 1e-6) — 归一化层的 epsilon 值。
  • causal_condition (bool,默认为 True) — 从 t = 0 开始调制文本和条件图像 token。使用 KV 缓存时必需。

Qwen-Image 2.1 使用的单流 Transformer。

文本和图像潜在表示共享同一个序列:条件图像的 token 被替换到视觉语言编码器为其预留的文本流位置中,而目标图像的 token 则被追加到序列末尾。所有块都由同一个共享的 modulation 投影提供输入,因此各块自身不持有任何调制参数。

2.1 版本有两个显著的行为特征:

  • 块因果注意力 — 注意力遵循 (q_idx >= kv_idx) or same_image_block,因此序列是因果的,同时每个图像块内部保持双向。使用 flex 注意力后端可提供高效的单遍注意力;若未使用,模型将采用精确的多遍 SDPA 预填充,分别处理每个块。
  • causal_condition — 文本和条件图像 token 从 t = 0 开始进行调制,而不是从采样的时间步开始,这使得它们的激活值与时间步无关,从而可以在去噪步骤之间进行缓存。

build_token_metadata

< >

( image_pad_mask: Tensorimg_shapes: list ) → tuple[torch.Tensor, torch.Tensor]

参数

  • image_pad_mask (torch.Tensor) — (seq_len,) 布尔值,在图像标记位置为 True。
  • img_shapes (list[tuple[int, int, int]]) — 每张图像在潜在 token 中的 (frame, height, width),条件图像在前,目标图像在后。

返回值

tuple[torch.Tensor, torch.Tensor]

image_ids (seq_len,),其中文本位置为 -1,每个图像块具有唯一的 id,以及 target_token_mask (seq_len,),用于标记目标图像的 token。

为联合序列中的每个 token 标注其所属的图像块。

块边界由 img_shapes 中的 token 数量决定,而不是由 image_pad_mask 中 True 的连续段决定:如果两张条件图像恰好相邻且中间没有文本,它们会形成一个连续段,但必须保持为独立的块,否则它们会进行双向注意力计算。

forward

< >

( hidden_states: Tensorencoder_hidden_states: Tensortimestep: Tensorimg_shapes: listimg_mask: Tensorencoder_hidden_states_mask: typing.Optional[torch.Tensor] = Noneattention_kwargs: dict[str, typing.Any] | None = Nonekv_cache: diffusers.models.transformers.transformer_qwenimage21.QwenImage21KVCache | None = Nonekv_cache_mode: str | None = Nonereturn_dict: bool = True ) → Transformer2DModelOutput or tuple

参数

  • hidden_states(形状为 (batch_size, image_sequence_length, in_channels) 的 torch.Tensor)— 打包后的潜在变量,条件图像在前,目标图像在最后。
  • encoder_hidden_states (torch.Tensor,形状为 (batch_size, text_sequence_length, context_in_dim)) — 来自视觉-语言编码器的文本嵌入。
  • timestep (torch.Tensor) — 当前去噪步骤,已缩放到 [0, 1]。
  • img_shapes (list[list[tuple[int, int, int]]]) — 每个样本的 (frame, height, width) 列表,以潜在 token 为单位,条件图像在前,目标图像在后。所有样本必须具有相同的布局。
  • img_mask(形状为 (batch_size, vlm_sequence_length) 的 torch.Tensor)— 在视觉-语言编码器的图像槽位处为 True,每个槽位代表一组 2x2 的潜在 token。
  • encoder_hidden_states_mask (torch.Tensor, 可选) — (batch_size, text_sequence_length) 布尔值,用于标记有效的文本 token。填充位置将被排除在注意力计算之外。
  • kv_cache (QwenImage21KVCache,可选) — 缓存容器。与 kv_cache_mode 一起传入以启用前缀 KV 缓存。
  • kv_cache_mode (str,可选) — "extract" 用于预填充缓存(首次去噪步骤),"cached" 用于从缓存中解码(后续步骤)。 需要 causal_condition=True。
  • attention_kwargs (dict, 可选) — 转发给注意力处理器,并携带 LoRA 层的 scale。
  • return_dict (bool,可选,默认为 True) — 是否返回 Transformer2DModelOutput 而不是普通元组。

返回值

Transformer2DModelOutput 或 tuple

如果 return_dict 为 True,则返回 Transformer2DModelOutput,否则返回一个普通 tuple,其第一个元素为去噪后的潜变量。

Transformer2DModelOutput

类 diffusers.models.modeling_outputs.Transformer2DModelOutput

< >

( sample: torch.Tensor )

参数

  • sample(形状为 (batch_size, num_channels, height, width) 的 torch.Tensor,或者如果 Transformer2DModel 是离散的,则为 (batch size, num_vector_embeds - 1, num_latent_pixels))— 以 encoder_hidden_states 输入为条件的隐藏状态输出。如果是离散的,则返回未加噪潜在像素的概率分布。

Transformer2DModel 的输出。

Update on GitHub