Diffusers documentation
Flux3Transformer2DModel
Flux3Transformer2DModel
用于 FLUX.3 类图像数据的 Transformer 模型，支持文本生成图像和图像生成图像。三种 token 模式共享一个联合序列：带噪图像 token、经 VAE 编码的参考图像 token 以及文本上下文。每种模式首先通过其专属的按模式划分的 Transformer 块，随后各模式进行拼接以进入联合单块，最后图像 token 通过输出层。时间步是按 token 划分的：图像 token 携带流时间步，而参考 token 和文本上下文均位于 0。
import torch
from diffusers import Flux3Transformer2DModel
transformer = Flux3Transformer2DModel.from_pretrained(
"black-forest-labs/flux-3-image-klein-base", subfolder="transformer", dtype=torch.bfloat16
)Flux3Transformer2DModel
类 diffusers.Flux3Transformer2DModel
< source >( in_channels: int = 128cond_in_channels: int = 128num_layers: int = 5num_single_layers: int = 28attention_head_dim: int = 128num_attention_heads: int = 24joint_attention_dim: int = 20480timestep_channels: int = 256mlp_ratio: float = 3.0axes_dims_rope: tuple[int, ...] = (32, 32, 32, 32)rope_theta: int = 10000eps: float = 1e-06flow_map: bool = False )
参数
- in_channels (
int，默认为128) — 图像 token 的通道宽度（即打包后的 VAE 潜变量，latent_channels * 4）。 - cond_in_channels (
int，默认为128) — 参考图像 token 的通道宽度。 - num_layers (
int，默认为5) — 每种模式的 Transformer 块数量。 - num_single_layers (
int，默认为28) — 联合单 Transformer 块的数量。 - attention_head_dim (
int，默认为128) — 每个注意力头使用的维度数。 - num_attention_heads (
int，默认为24) — 要使用的注意力头数量。 - joint_attention_dim (
int，默认为20480) —encoder_hidden_states（堆叠的 Qwen3-VL 隐藏状态）的嵌入/通道维度。 - timestep_channels (
int，默认为256) — 正弦时间步嵌入的通道数。 - mlp_ratio (
float，默认为3.0) — 前馈子块的扩展比率。 - axes_dims_rope (
tuple[int]，默认为(32, 32, 32, 32)) — 用于旋转位置嵌入的维度，对应(t, h, w, l)位置轴。 - rope_theta (
int，默认为10000) — 旋转位置编码的基础频率。 - eps (
float，默认为1e-6) — 层归一化和 QK 归一化的 epsilon 值。 - flow_map (
bool，默认为False) — 模型是否为流映射（步数蒸馏）：会添加target_timestep_embedder，并要求在 forward 过程中传入target_timestep。
FLUX.3 中用于文本到图像和图像到图像生成的 Transformer 模型。
三种 token 模式共享一个联合序列：带噪图像 token（hidden_states）、可选的 VAE 编码参考图像 token（image_cond_hidden_states）以及文本上下文（encoder_hidden_states）。每种模式首先通过其专属的按模式划分的 transformer 块（图像模式使用 mode_blocks，文本上下文使用 txt_mode_blocks），随后将文本上下文与当前激活的图像模式拼接，用于联合单块处理，且仅图像 token 通过输出层。调制是按模式进行的：图像 token 由 flow 时间步调制，参考图像 token 和文本上下文由时间步 0 调制。Flow-map 检查点（flow_map=True，即步数蒸馏模型）还会嵌入每个去噪步骤所跳转到的时间步，并将其添加到每种模式的调制中。
forward
< source >( hidden_states: Tensorencoder_hidden_states: Tensortimestep: Tensorimg_ids: Tensortxt_ids: Tensorimage_cond_hidden_states: typing.Optional[torch.Tensor] = Noneimage_cond_ids: typing.Optional[torch.Tensor] = Nonetarget_timestep: typing.Optional[torch.Tensor] = Noneattention_kwargs: dict[str, typing.Any] | None = Nonereturn_dict: bool = True )
参数
- hidden_states（形状为
(batch_size, image_seq_len, in_channels)的torch.Tensor）— 带噪声的图像潜在 token。 - encoder_hidden_states（形状为
(batch_size, text_seq_len, joint_attention_dim)的torch.Tensor）— 文本上下文嵌入（堆叠的 Qwen3-VL 隐藏状态）。 - timestep（形状为
(batch_size,)的torch.Tensor）— 图像 token 的 Flow 时间步，取值范围为[0, 1]。参考图像 token 和文本上下文位于时间步0。 - img_ids（形状为
(image_seq_len, 4)的torch.Tensor）— 图像 token 的位置 ID(t, h, w, l)，用于旋转位置嵌入，并在批次中的所有样本间共享。 - txt_ids（形状为
(text_seq_len, 4)的torch.Tensor）— 文本上下文 token 的位置 ID。 - image_cond_hidden_states (
torch.Tensor，形状为(batch_size, cond_seq_len, cond_in_channels)，可选) — 经 VAE 编码的参考图像 token（干净的 conditioning 内容，用于图像到图像生成）。 - image_cond_ids（形状为
(cond_seq_len, 4)的torch.Tensor，可选）— 参考图像 token 的位置 ID。使用image_cond_hidden_states时必填。 - target_timestep（形状为
(batch_size,)的torch.Tensor，可选） — 当前去噪步骤所跳转到的[0, 1]区间内的 Flow 时间步。仅 flow-map 模型（flow_map=True）需要此参数。 - attention_kwargs (
dict, 可选) — 一个 kwargs 字典，如果指定了其scale项，则会在本次前向传播中对 LoRA 层进行缩放。 - return_dict (
bool，可选，默认为True) — 是否返回~models.transformer_2d.Transformer2DModelOutput而不是普通的元组。
Flux3Transformer2DModel 的前向方法。
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 的输出。