Diffusers 中 Flux2Transformer2DModel 全面解析架构、配置与源码级原理【免费下载链接】diffusers Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers本文以 Hugging Face Diffusers 仓库中 Flux2Transformer2DModel 官方 API 文档 为骨架结合 Flux 2 Transformer 源码实现 与 Flux2 模块化管线 深入展开系统讲解 Flux 2 图像 Transformer 的模型结构、全部配置参数、输入输出约定、参考图 KV 缓存机制以及并行化设计。读完本文你将理解Flux2Transformer2DModel在 Flux2 文生图/图生图流程中的位置并具备在自定义代码中正确实例化、调用与调试该模型的能力。Flux2Transformer2DModel 是什么Flux2Transformer2DModel是 Diffusers 对 Flux2 系列模型中图像类数据image-like dataTransformer 主干网络的官方实现位于 src/diffusers/models/transformers/transformer_flux2.py。它属于扩散模型DiT 风格去噪网络负责根据带噪潜变量latents、文本条件与时间步预测噪声 / 速度场是 Flux2 管线中计算量最大的核心模块。从类的继承关系源码第 1059–1067 行可以看到它整合了 Diffusers 的多套基础设施class Flux2Transformer2DModel( ModelMixin, # 模型保存 / 加载from_pretrained / save_pretrained ConfigMixin, # 配置序列化register_to_config PeftAdapterMixin, # PEFT LoRA 适配器挂载 FromOriginalModelMixin,# 从原始 checkpoint 加载 FluxTransformer2DLoadersMixin, # Flux Transformer 专用加载器 CacheMixin, # 缓存工具如 mag cache AttentionMixin, # attention processor 机制 ):该模型同时服务于两条 Flux2 管线家族标准 Flux2pipeline_flux2.py与 Klein 变体pipeline_flux2_klein.py、pipeline_flux2_klein_kv.py差异主要在文本编码器Mistral3 vs Qwen3与是否使用参考图 KV 缓存Transformer 本体是同一套实现。模型配置参数全解Flux2Transformer2DModel.__init__源码第 1141–1158 行通过register_to_config注册全部超参数意味着这些参数会写入model_index.json/ 配置文件可通过from_pretrained自动恢复。默认参数直接对应 Flux2 官方权重规模参数默认值说明patch_size1将输入切成 patch 的大小Flux2 为 1 即不平铺 patchin_channels128输入潜变量的通道数out_channelsNone输出通道数为None时默认等于in_channelsnum_layers8双流double-streamDiT 块数量负责文本与图像流的联合注意力num_single_layers48单流single-streamDiT 块数量两流拼接后统一处理attention_head_dim128每个注意力头的维度num_attention_heads48注意力头数量inner_dim heads × head_dim 6144joint_attention_dim15360联合注意力维度即encoder_hidden_states文本嵌入的特征维度timestep_guidance_channels256时间步 / guidance 嵌入的正弦编码通道数mlp_ratio3.0FFN 隐藏层相对维度倍数axes_dims_rope(32, 32, 32, 32)RoPE 旋转位置编码各轴的维度分配rope_theta2000RoPE 的 theta 基频eps1e-6LayerNorm / RMSNorm 的 epsilonguidance_embedsTrue是否使用 guidance 嵌入guidance-distilled 变体需要其中inner_dim6144由num_attention_heads * attention_head_dim推导而来是后续所有模块输入投影、调制、注意力的统一通道数。构造参数对模型结构的影响num_layers与num_single_layers分别控制Flux2TransformerBlock双流与Flux2SingleTransformerBlock单流两个nn.ModuleList的长度源码第 1186–1213 行。Flux2 默认 8 个双流块 48 个单流块这也对应了源码中_repeated_blocks与_no_split_modules的声明用于梯度检查点gradient checkpointing与 offload 时按块切分。guidance_embedsFalse时Flux2TimestepGuidanceEmbeddings内部的guidance_embedder被置为None源码第 1016–1021 行此时即便传入guidance也只会输出纯时间步嵌入。Klein 变体即属此类其调用处传guidanceNone见 denoise.py 第 214 行。前向传播输入输出约定forward签名源码第 1225–1240 行如下def forward( self, hidden_states: torch.Tensor, # (B, img_seq_len, in_channels) 图像潜 token encoder_hidden_states: torch.Tensor None, # (B, txt_seq_len, joint_attention_dim) 文本嵌入 timestep: torch.LongTensor None, # 去噪时间步内部 ×1000 后编码 img_ids: torch.Tensor None, # 图像 token 的 RoPE 位置 id txt_ids: torch.Tensor None, # 文本 token 的 RoPE 位置 id guidance: torch.Tensor None, # guidance 缩放嵌入guidance-distilled 变体 joint_attention_kwargs: dict[str, Any] | None None, # 透传给 attention processor 的额外参数 return_dict: bool True, kv_cache: Flux2KVCache | None None, # 参考图 KV 缓存 kv_cache_mode: str | None None, # extract / cached / None num_ref_tokens: int 0, # 参考图 token 数 ref_fixed_timestep: float 0.0, # 参考 token 调制用固定时间步 ) - torch.Tensor | Flux2Transformer2DModelOutput输出由 Flux2Transformer2DModelOutput 承载字段类型说明sampletorch.Tensor形状(batch_size, num_channels, height, width)条件于encoder_hidden_states的隐藏状态输出即预测的噪声kv_cacheFlux2KVCache | None参考图 token 的 KV 缓存仅在kv_cache_modeextract时返回调用示例从 Flux2LoopDenoiser 的去噪调用 可以看到真实调用方式noise_pred components.transformer( hidden_stateslatent_model_input, # (B, img_seq_len, 128) timesteptimestep / 1000, # 时间步除以 1000 传入forward 内部再 ×1000 guidanceblock_state.guidance, encoder_hidden_statesblock_state.prompt_embeds, # Mistral3 / Qwen3 文本嵌入 txt_idsblock_state.txt_ids, # 文本 token 的 4 维位置 id (T, H, W, L) img_idsimg_ids, # 图像 token 的 4 维位置 id joint_attention_kwargsblock_state.joint_attention_kwargs, return_dictFalse, )[0]注意两个容易踩坑的约定时间步缩放管线传入timestep / 1000而forward内部第一件事就是timestep * 1000源码第 1284 行即模型内部以「千分制」时间步做正弦编码guidance同样被×1000第 1287 行。位置 id 的维度img_ids/txt_ids支持 2D 与 3D 输入3D 时取[0]第 1322–1325 行其最后一维长度必须等于len(axes_dims_rope)默认为 4Flux2PosEmbed会按每个轴分别计算一维 RoPE 再拼接源码第 978–998 行。内部架构双流块与单流块Flux2 Transformer 采用与 Flux.1 类似的「双流 → 单流」分层设计但内部实现有显著区别。双流块 Flux2TransformerBlockFlux2TransformerBlock源码第 876–968 行同时维护图像流与文本流两套隐藏状态每个块包含联合注意力Flux2Attention图像与文本各自经过独立的 RMSNormQK-NormQ 与 K 在投影后均做归一化norm_q/norm_k文本侧通过add_q/k/v_proj生成额外的 KV 并拼接到图像注意力中Flux2AttnProcessor第 364–366 行两个独立的 FFNff图像流与ff_context文本流都使用Flux2FeedForwardAdaLN 调制图像与文本分别用独立的调制参数temb_mod_img/temb_mod_txt每个调制包含 shift / scale / gate 三组且 attention 与 MLP 各一组Flux2Modulation.split(temb_mod_img, 2)第 923–926 行即一个双流块需要 2 组调制参数集。单流块 Flux2SingleTransformerBlockFlux2SingleTransformerBlock源码第 807–873 行先把文本与图像流cat成单一序列然后送入parallel attentionFlux2ParallelSelfAttention源码第 723–804 行。其核心特点是借鉴 ViT-22B 的并行 Transformer 块设计QKV 投影与 MLP 输入投影融合为单个线性层to_qkv_mlp_proj输出维度为3 * inner_dim mlp_hidden_dim * mlp_mult_factor第 768–770 行注意力输出投影与 MLP 输出投影融合为to_out第 782 行只有一组调制参数mod_param_sets1因为 attention 与 FF 并行执行源码第 1178–1179 行的注释说明了这一点。激活函数与 FFNFlux2SwiGLU源码第 285–298 行是一个无训练参数的模块Flux2 把 SwiGLU 的 gate 线性层融合进了前一个线性层因此 FFN 实际为linear_in(2×inner) → SwiGLU(砍半) → linear_out两段式结构Flux2FeedForward第 316–318 行。这也直接反映在张量并行分片计划中ff.linear_in使用PackedColwiseParallel([1, 1])按 gate / linear 两半等分源码第 1132 行。参考图 KV 缓存机制Klein 变体核心这是 Flux2 Transformer 在 Diffusers 中新增的关键能力服务于pipeline_flux2_klein_kv.py的参考图reference image加速场景。缓存数据结构Flux2KVLayerCache源码第 61–85 行单层缓存保存参考 token 经 RoPE 之后的 K / V张量形状为(batch_size, num_ref_tokens, num_heads, head_dim)提供store/get/clear三个方法Flux2KVCache源码第 88–110 行全局容器按双流块与单流块分别维护层缓存列表并记录num_ref_tokens。extract / cached 两种模式在Flux2KVAttnProcessor双流源码第 398–494 行与Flux2KVParallelSelfAttnProcessor单流源码第 633–720 行中extract模式第一次去噪步序列布局为[txt, ref, img]。参考 token 只做自注意力ref self-attend而 txt 与 img token 关注全部 token——由_flux2_kv_causal_attention源码第 113–170 行实现同时把参考 token 的 K / Vclone()存入缓存第 458 行、第 689 行。参考 token 的调制参数使用固定时间步ref_fixed_timestep默认 0.0通过_blend_double_block_mods/_blend_single_block_mods与图像调制参数按位置拼接源码第 1296–1315 行、第 1381–1385 行。cached模式后续去噪步序列布局退化为[txt, img]注意力时把缓存的参考 K / V注入到 txt 与 img 之间第 136–139 行从而省去参考 token 的重复计算。kv_cache_modeNone完全退化为标准前向行为与Flux2AttnProcessor一致。forward 层的装配逻辑见源码第 1335–1350 行extract 模式创建并返回新的Flux2KVCachecached 模式读取外部传入的缓存。该机制与仓库中WanAnimate2的参考帧 KV 缓存设计见 transformer_wan_animate_2.py思路一致是 Diffusers 处理「参考图像条件」的通用加速范式。并行化支持上下文并行与张量并行从源码类属性可直接看到 Flux2 Transformer 对大规模并行的内置支持源码第 1103–1139 行上下文并行Context Parallel_cp_plan将hidden_states、encoder_hidden_states、img_ids、txt_ids沿序列维度split_dim1切分到不同设备proj_out再聚合输出ContextParallelOutput(gather_dim1)。张量并行Tensor Parallel_tp_plan精确声明了每个投影的切分方式——双流块的to_q/k/v与add_q/k/v_proj按列切分colwise、to_out按行切分rowwise、SwiGLU 的linear_in用PackedColwiseParallel([1, 1])等分、单流块的融合投影用PackedColwiseParallel()/PackedRowwiseParallel()。而 AdaLN 调制层与 QK-Norm刻意保持复制注释指出调制需要访问完整隐藏维度QK-Norm 在头已切分后仍作用于完整head_dim。注意力处理器使用unflatten(-1, (-1, attn.head_dim))让-1吸收头数从而在张量并行下每个 rank 只需处理自己的头切片处理器无需感知 TP 度数源码第 347–351 行注释。此外_supports_gradient_checkpointing True与_no_split_modules [Flux2TransformerBlock, Flux2SingleTransformerBlock]表明该模型支持梯度检查点与分层 offload适合大模型训练场景。在 Flux2 管线中的完整调用链Flux2Transformer2DModel在 Diffusers 中是作为管线组件ComponentSpec(transformer, Flux2Transformer2DModel)被模块化管线装配的见 denoise.py。典型流程为文本编码器Mistral3 或 Qwen3生成prompt_embeds与txt_idsVAE 编码器将输入图像压缩为潜变量经 pack 后得到latents与img_ids4 维位置 id调度器FlowMatchEulerDiscreteScheduler逐步去噪每一步调用transformer(hidden_states, timestep, guidance, encoder_hidden_states, txt_ids, img_ids, ...)若为 Klein KV 管线第一步以kv_cache_modeextract提取参考图 KV后续步以kv_cache_modecached复用去噪完成后经 VAE 解码为最终图像。对应测试覆盖见 tests/modular_pipelines/flux2/ 下的test_modular_pipeline_flux2.py与test_modular_pipeline_flux2_klein.py可作行为参考。实战要点小结实例化直接使用Flux2Transformer2DModel.from_pretrained(black-forest-labs/FLUX.2-dev, subfoldertransformer)即可加载官方权重无需手动指定参数——register_to_config保证配置随权重自动恢复。LoRA 适配继承自PeftAdapterMixin且 forward 上标注了apply_lora_scale(joint_attention_kwargs)因此可在joint_attention_kwargs中传入scale控制 LoRA 强度。dtype 注意FP16 推理时单流块输出会执行clip(-65504, 65504)防止溢出源码第 866–867 行这是模型自带的数值稳定性保护。KV 缓存边界kv_cache_modeextract时输出中会额外携带kv_cache且_skip_keys [kv_cache]源码第 1223 行保证该动态对象不会进入状态字典save_pretrained不会序列化它。Flux2Transformer2DModel 完整覆盖了 Flux2 从文本条件到噪声预测的全部主干逻辑理解其双流 / 单流分层、AdaLN 调制、QK-Norm 注意力与 KV 缓存机制是深入使用与二次开发 Flux2 系列管线的基础。【免费下载链接】diffusers Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 Diffusers 中 Flux2Transformer2DModel 全面解析：架构、配置与源码级原理
Diffusers 中 Flux2Transformer2DModel 全面解析架构、配置与源码级原理【免费下载链接】diffusers Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers本文以 Hugging Face Diffusers 仓库中 Flux2Transformer2DModel 官方 API 文档 为骨架结合 Flux 2 Transformer 源码实现 与 Flux2 模块化管线 深入展开系统讲解 Flux 2 图像 Transformer 的模型结构、全部配置参数、输入输出约定、参考图 KV 缓存机制以及并行化设计。读完本文你将理解Flux2Transformer2DModel在 Flux2 文生图/图生图流程中的位置并具备在自定义代码中正确实例化、调用与调试该模型的能力。Flux2Transformer2DModel 是什么Flux2Transformer2DModel是 Diffusers 对 Flux2 系列模型中图像类数据image-like dataTransformer 主干网络的官方实现位于 src/diffusers/models/transformers/transformer_flux2.py。它属于扩散模型DiT 风格去噪网络负责根据带噪潜变量latents、文本条件与时间步预测噪声 / 速度场是 Flux2 管线中计算量最大的核心模块。从类的继承关系源码第 1059–1067 行可以看到它整合了 Diffusers 的多套基础设施class Flux2Transformer2DModel( ModelMixin, # 模型保存 / 加载from_pretrained / save_pretrained ConfigMixin, # 配置序列化register_to_config PeftAdapterMixin, # PEFT LoRA 适配器挂载 FromOriginalModelMixin,# 从原始 checkpoint 加载 FluxTransformer2DLoadersMixin, # Flux Transformer 专用加载器 CacheMixin, # 缓存工具如 mag cache AttentionMixin, # attention processor 机制 ):该模型同时服务于两条 Flux2 管线家族标准 Flux2pipeline_flux2.py与 Klein 变体pipeline_flux2_klein.py、pipeline_flux2_klein_kv.py差异主要在文本编码器Mistral3 vs Qwen3与是否使用参考图 KV 缓存Transformer 本体是同一套实现。模型配置参数全解Flux2Transformer2DModel.__init__源码第 1141–1158 行通过register_to_config注册全部超参数意味着这些参数会写入model_index.json/ 配置文件可通过from_pretrained自动恢复。默认参数直接对应 Flux2 官方权重规模参数默认值说明patch_size1将输入切成 patch 的大小Flux2 为 1 即不平铺 patchin_channels128输入潜变量的通道数out_channelsNone输出通道数为None时默认等于in_channelsnum_layers8双流double-streamDiT 块数量负责文本与图像流的联合注意力num_single_layers48单流single-streamDiT 块数量两流拼接后统一处理attention_head_dim128每个注意力头的维度num_attention_heads48注意力头数量inner_dim heads × head_dim 6144joint_attention_dim15360联合注意力维度即encoder_hidden_states文本嵌入的特征维度timestep_guidance_channels256时间步 / guidance 嵌入的正弦编码通道数mlp_ratio3.0FFN 隐藏层相对维度倍数axes_dims_rope(32, 32, 32, 32)RoPE 旋转位置编码各轴的维度分配rope_theta2000RoPE 的 theta 基频eps1e-6LayerNorm / RMSNorm 的 epsilonguidance_embedsTrue是否使用 guidance 嵌入guidance-distilled 变体需要其中inner_dim6144由num_attention_heads * attention_head_dim推导而来是后续所有模块输入投影、调制、注意力的统一通道数。构造参数对模型结构的影响num_layers与num_single_layers分别控制Flux2TransformerBlock双流与Flux2SingleTransformerBlock单流两个nn.ModuleList的长度源码第 1186–1213 行。Flux2 默认 8 个双流块 48 个单流块这也对应了源码中_repeated_blocks与_no_split_modules的声明用于梯度检查点gradient checkpointing与 offload 时按块切分。guidance_embedsFalse时Flux2TimestepGuidanceEmbeddings内部的guidance_embedder被置为None源码第 1016–1021 行此时即便传入guidance也只会输出纯时间步嵌入。Klein 变体即属此类其调用处传guidanceNone见 denoise.py 第 214 行。前向传播输入输出约定forward签名源码第 1225–1240 行如下def forward( self, hidden_states: torch.Tensor, # (B, img_seq_len, in_channels) 图像潜 token encoder_hidden_states: torch.Tensor None, # (B, txt_seq_len, joint_attention_dim) 文本嵌入 timestep: torch.LongTensor None, # 去噪时间步内部 ×1000 后编码 img_ids: torch.Tensor None, # 图像 token 的 RoPE 位置 id txt_ids: torch.Tensor None, # 文本 token 的 RoPE 位置 id guidance: torch.Tensor None, # guidance 缩放嵌入guidance-distilled 变体 joint_attention_kwargs: dict[str, Any] | None None, # 透传给 attention processor 的额外参数 return_dict: bool True, kv_cache: Flux2KVCache | None None, # 参考图 KV 缓存 kv_cache_mode: str | None None, # extract / cached / None num_ref_tokens: int 0, # 参考图 token 数 ref_fixed_timestep: float 0.0, # 参考 token 调制用固定时间步 ) - torch.Tensor | Flux2Transformer2DModelOutput输出由 Flux2Transformer2DModelOutput 承载字段类型说明sampletorch.Tensor形状(batch_size, num_channels, height, width)条件于encoder_hidden_states的隐藏状态输出即预测的噪声kv_cacheFlux2KVCache | None参考图 token 的 KV 缓存仅在kv_cache_modeextract时返回调用示例从 Flux2LoopDenoiser 的去噪调用 可以看到真实调用方式noise_pred components.transformer( hidden_stateslatent_model_input, # (B, img_seq_len, 128) timesteptimestep / 1000, # 时间步除以 1000 传入forward 内部再 ×1000 guidanceblock_state.guidance, encoder_hidden_statesblock_state.prompt_embeds, # Mistral3 / Qwen3 文本嵌入 txt_idsblock_state.txt_ids, # 文本 token 的 4 维位置 id (T, H, W, L) img_idsimg_ids, # 图像 token 的 4 维位置 id joint_attention_kwargsblock_state.joint_attention_kwargs, return_dictFalse, )[0]注意两个容易踩坑的约定时间步缩放管线传入timestep / 1000而forward内部第一件事就是timestep * 1000源码第 1284 行即模型内部以「千分制」时间步做正弦编码guidance同样被×1000第 1287 行。位置 id 的维度img_ids/txt_ids支持 2D 与 3D 输入3D 时取[0]第 1322–1325 行其最后一维长度必须等于len(axes_dims_rope)默认为 4Flux2PosEmbed会按每个轴分别计算一维 RoPE 再拼接源码第 978–998 行。内部架构双流块与单流块Flux2 Transformer 采用与 Flux.1 类似的「双流 → 单流」分层设计但内部实现有显著区别。双流块 Flux2TransformerBlockFlux2TransformerBlock源码第 876–968 行同时维护图像流与文本流两套隐藏状态每个块包含联合注意力Flux2Attention图像与文本各自经过独立的 RMSNormQK-NormQ 与 K 在投影后均做归一化norm_q/norm_k文本侧通过add_q/k/v_proj生成额外的 KV 并拼接到图像注意力中Flux2AttnProcessor第 364–366 行两个独立的 FFNff图像流与ff_context文本流都使用Flux2FeedForwardAdaLN 调制图像与文本分别用独立的调制参数temb_mod_img/temb_mod_txt每个调制包含 shift / scale / gate 三组且 attention 与 MLP 各一组Flux2Modulation.split(temb_mod_img, 2)第 923–926 行即一个双流块需要 2 组调制参数集。单流块 Flux2SingleTransformerBlockFlux2SingleTransformerBlock源码第 807–873 行先把文本与图像流cat成单一序列然后送入parallel attentionFlux2ParallelSelfAttention源码第 723–804 行。其核心特点是借鉴 ViT-22B 的并行 Transformer 块设计QKV 投影与 MLP 输入投影融合为单个线性层to_qkv_mlp_proj输出维度为3 * inner_dim mlp_hidden_dim * mlp_mult_factor第 768–770 行注意力输出投影与 MLP 输出投影融合为to_out第 782 行只有一组调制参数mod_param_sets1因为 attention 与 FF 并行执行源码第 1178–1179 行的注释说明了这一点。激活函数与 FFNFlux2SwiGLU源码第 285–298 行是一个无训练参数的模块Flux2 把 SwiGLU 的 gate 线性层融合进了前一个线性层因此 FFN 实际为linear_in(2×inner) → SwiGLU(砍半) → linear_out两段式结构Flux2FeedForward第 316–318 行。这也直接反映在张量并行分片计划中ff.linear_in使用PackedColwiseParallel([1, 1])按 gate / linear 两半等分源码第 1132 行。参考图 KV 缓存机制Klein 变体核心这是 Flux2 Transformer 在 Diffusers 中新增的关键能力服务于pipeline_flux2_klein_kv.py的参考图reference image加速场景。缓存数据结构Flux2KVLayerCache源码第 61–85 行单层缓存保存参考 token 经 RoPE 之后的 K / V张量形状为(batch_size, num_ref_tokens, num_heads, head_dim)提供store/get/clear三个方法Flux2KVCache源码第 88–110 行全局容器按双流块与单流块分别维护层缓存列表并记录num_ref_tokens。extract / cached 两种模式在Flux2KVAttnProcessor双流源码第 398–494 行与Flux2KVParallelSelfAttnProcessor单流源码第 633–720 行中extract模式第一次去噪步序列布局为[txt, ref, img]。参考 token 只做自注意力ref self-attend而 txt 与 img token 关注全部 token——由_flux2_kv_causal_attention源码第 113–170 行实现同时把参考 token 的 K / Vclone()存入缓存第 458 行、第 689 行。参考 token 的调制参数使用固定时间步ref_fixed_timestep默认 0.0通过_blend_double_block_mods/_blend_single_block_mods与图像调制参数按位置拼接源码第 1296–1315 行、第 1381–1385 行。cached模式后续去噪步序列布局退化为[txt, img]注意力时把缓存的参考 K / V注入到 txt 与 img 之间第 136–139 行从而省去参考 token 的重复计算。kv_cache_modeNone完全退化为标准前向行为与Flux2AttnProcessor一致。forward 层的装配逻辑见源码第 1335–1350 行extract 模式创建并返回新的Flux2KVCachecached 模式读取外部传入的缓存。该机制与仓库中WanAnimate2的参考帧 KV 缓存设计见 transformer_wan_animate_2.py思路一致是 Diffusers 处理「参考图像条件」的通用加速范式。并行化支持上下文并行与张量并行从源码类属性可直接看到 Flux2 Transformer 对大规模并行的内置支持源码第 1103–1139 行上下文并行Context Parallel_cp_plan将hidden_states、encoder_hidden_states、img_ids、txt_ids沿序列维度split_dim1切分到不同设备proj_out再聚合输出ContextParallelOutput(gather_dim1)。张量并行Tensor Parallel_tp_plan精确声明了每个投影的切分方式——双流块的to_q/k/v与add_q/k/v_proj按列切分colwise、to_out按行切分rowwise、SwiGLU 的linear_in用PackedColwiseParallel([1, 1])等分、单流块的融合投影用PackedColwiseParallel()/PackedRowwiseParallel()。而 AdaLN 调制层与 QK-Norm刻意保持复制注释指出调制需要访问完整隐藏维度QK-Norm 在头已切分后仍作用于完整head_dim。注意力处理器使用unflatten(-1, (-1, attn.head_dim))让-1吸收头数从而在张量并行下每个 rank 只需处理自己的头切片处理器无需感知 TP 度数源码第 347–351 行注释。此外_supports_gradient_checkpointing True与_no_split_modules [Flux2TransformerBlock, Flux2SingleTransformerBlock]表明该模型支持梯度检查点与分层 offload适合大模型训练场景。在 Flux2 管线中的完整调用链Flux2Transformer2DModel在 Diffusers 中是作为管线组件ComponentSpec(transformer, Flux2Transformer2DModel)被模块化管线装配的见 denoise.py。典型流程为文本编码器Mistral3 或 Qwen3生成prompt_embeds与txt_idsVAE 编码器将输入图像压缩为潜变量经 pack 后得到latents与img_ids4 维位置 id调度器FlowMatchEulerDiscreteScheduler逐步去噪每一步调用transformer(hidden_states, timestep, guidance, encoder_hidden_states, txt_ids, img_ids, ...)若为 Klein KV 管线第一步以kv_cache_modeextract提取参考图 KV后续步以kv_cache_modecached复用去噪完成后经 VAE 解码为最终图像。对应测试覆盖见 tests/modular_pipelines/flux2/ 下的test_modular_pipeline_flux2.py与test_modular_pipeline_flux2_klein.py可作行为参考。实战要点小结实例化直接使用Flux2Transformer2DModel.from_pretrained(black-forest-labs/FLUX.2-dev, subfoldertransformer)即可加载官方权重无需手动指定参数——register_to_config保证配置随权重自动恢复。LoRA 适配继承自PeftAdapterMixin且 forward 上标注了apply_lora_scale(joint_attention_kwargs)因此可在joint_attention_kwargs中传入scale控制 LoRA 强度。dtype 注意FP16 推理时单流块输出会执行clip(-65504, 65504)防止溢出源码第 866–867 行这是模型自带的数值稳定性保护。KV 缓存边界kv_cache_modeextract时输出中会额外携带kv_cache且_skip_keys [kv_cache]源码第 1223 行保证该动态对象不会进入状态字典save_pretrained不会序列化它。Flux2Transformer2DModel 完整覆盖了 Flux2 从文本条件到噪声预测的全部主干逻辑理解其双流 / 单流分层、AdaLN 调制、QK-Norm 注意力与 KV 缓存机制是深入使用与二次开发 Flux2 系列管线的基础。【免费下载链接】diffusers Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考