Diffusers documentation
QwenImage21Transformer2DModel
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
< source >( 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
< source >( image_pad_mask: Tensorimg_shapes: list ) → tuple[torch.Tensor, torch.Tensor]
为联合序列中的每个 token 标注其所属的图像块。
块边界由 img_shapes 中的 token 数量决定,而不是由 image_pad_mask 中 True 的连续段决定:如果两张条件图像恰好相邻且中间没有文本,它们会形成一个连续段,但必须保持为独立的块,否则它们会进行双向注意力计算。
forward
< source >( 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
< source >( 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 的输出。