Diffusers 中的 HeliosTransformer3DModel:14B 实时自回归视频扩散 Transformer 全解析
【免费下载链接】diffusers🤗 Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers
HeliosTransformer3DModel 是 🤗 Diffusers 为 Helios 视频生成模型提供的 3D 视频类数据 Transformer 主干实现,支持文生视频(T2V)、图生视频(I2V)与视频生视频(V2V)三种任务,是 HeliosPipeline 与HeliosPyramidPipeline的核心去噪网络。阅读本文后,你将掌握该模型的三种官方权重(Base / Mid / Distilled)的加载方式、完整配置参数含义、内部架构(3D Patch Embedding、多尺度历史记忆分支、3D RoPE、AdaLN 自注意力块)以及它与 Helios 调度器、金字塔采样流水线的协同方式。
模型背景:为什么需要 HeliosTransformer3DModel
Helios 是北京大学与字节跳动等机构提出的 14B 规模实时长视频生成模型(论文Helios: Real Real-Time Long Video Generation Model)。HeliosTransformer3DModel 是它在 Diffusers 中的 Transformer 主干实现,其设计围绕三个核心目标:
- 长视频抗漂移:不依赖 self-forcing、误差库、关键帧采样等常见抗漂移启发式,而是通过显式模拟漂移的训练策略与多尺度历史上下文压缩来保持分钟级视频的连贯性;
- 实时生成:不使用 KV-cache、因果掩码、稀疏注意力等标准加速手段,而是通过大幅压缩历史与噪声上下文、减少采样步数来实现低计算成本;
- 统一输入表示:以自回归扩散方式统一支持 T2V、I2V、V2V。
从源码结构看,HeliosTransformer3DModel 的输入是五维张量(batch, channels, frames, height, width),这与 Diffusers 中大多数面向单帧图像的 2D Transformer(如 Transformer2DModel)有本质区别,是其命名为 "3D" 的原因。
快速上手:加载 HeliosTransformer3DModel
官方文档给出了三种官方权重的加载方式,分别对应"质量优先""折中""效率优先"三个档位。完整的加载代码位于 helios_transformer3d.md:
from diffusers import HeliosTransformer3DModel # Best Quality(最优质量) transformer = HeliosTransformer3DModel.from_pretrained("BestWishYsh/Helios-Base", subfolder="transformer", dtype=torch.bfloat16) # Intermediate Weight(折中权重) transformer = HeliosTransformer3DModel.from_pretrained("BestWishYsh/Helios-Mid", subfolder="transformer", dtype=torch.bfloat16) # Best Efficiency(最优效率) transformer = HeliosTransformer3DModel.from_pretrained("BestWishYsh/Helios-Distilled", subfolder="transformer", dtype=torch.bfloat16)三个变体的差异(依据 pipeline 文档 与中文使用指南 docs/source/zh/using-diffusers/helios.md):
| 权重 | 定位 | 预测目标 | 配合调度器 | 生成方式 |
|---|---|---|---|---|
| Helios-Base | 质量优先 | v-prediction | HeliosScheduler + 标准 CFG | 50 步 |
| Helios-Mid | 折中 | v-prediction | HeliosScheduler + CFG-Zero* | 金字塔 3×20 步 |
| Helios-Distilled | 效率优先 | x0-prediction | HeliosDMDScheduler | 金字塔 3×2 步 |
加载时使用subfolder="transformer"是因为 Helios 仓库把 VAE、文本编码器与 Transformer 分别存放在不同的子目录中。配合完整推理时,通常用HeliosPipeline.from_pretrained(...)一次性加载全部组件,再通过pipeline.transformer访问该模型(下文"与 Pipeline 集成"一节给出示例)。
架构深度解析:从源码看 HeliosTransformer3DModel
核心实现位于 src/diffusers/models/transformers/transformer_helios.py。整体前向流程可归纳为 8 个阶段:输入补丁化 → 历史潜变量分支 → 条件嵌入 → Transformer 块 → 输出归一化 → 反补丁化(unpatchify)。
1. 3D Patch Embedding 与多尺度历史记忆分支
模型首先用nn.Conv3d将噪声潜变量切成 3D patch(默认patch_size=(1, 2, 2),即时间维不切分、空间维 2×2),把形状从(B, C, T, H, W)转为 token 序列:
self.patch_embedding = nn.Conv3d(in_channels, inner_dim, kernel_size=patch_size, stride=patch_size)更具特色的是Multi-Term Memory Patch(多尺度历史记忆分支,has_multi_term_memory_patch=True时启用)。它用三组不同 stride 的 3D 卷积patch_short/patch_mid/patch_long分别压缩短、中、长历史上下文,并在前向时按时间顺序拼接到当前 token 序列之前(对应forward参数latents_history_short/mid/long与indices_latents_history_*)。这样模型可以"看到"不同时间尺度上的历史帧,从根本上抑制长视频漂移与重复运动——这正是 Helios 宣称无需自强制等启发式的架构基础:
self.patch_short = nn.Conv3d(in_channels, inner_dim, kernel_size=patch_size, stride=patch_size) self.patch_mid = nn.Conv3d(in_channels, inner_dim, kernel_size=tuple(2 * p for p in patch_size), stride=tuple(2 * p for p in patch_size)) self.patch_long = nn.Conv3d(in_channels, inner_dim, kernel_size=tuple(4 * p for p in patch_size), stride=tuple(4 * p for p in patch_size))中间与长分支在进入卷积前会先通过pad_for_3d_conv做 replicate 填充,再对 rotary 位置编码做center_down_sample_3d(3D 平均池化)以对齐空间尺度。
2. 3D 旋转位置编码(HeliosRotaryPosEmbed)
HeliosRotaryPosEmbed实现了时间-空间三维旋转位置编码,三个维度(T/Y/X)分别使用独立的频率基底freqs_base_t / freqs_base_y / freqs_base_x,频率公式为1 / theta^(arange(0, dim, 2)/dim)。源码中特别注明:位置网格的 einsum 计算被强制在 float32 下进行,因为 bfloat16 无法表示超过 256 的连续整数,否则长视频中较远帧的位置会坍缩到同一频率,导致位置编码失效——这是长视频生成场景下的一个关键数值精度细节。
前向时,frame_indices与空间 meshgrid 扩展成(B, T, H, W)的位置网格,分别计算 T/Y/X 的 cos/sin 频率并拼接。最终在注意力处理器中通过apply_rotary_emb_transposed作用到 query 与 key 上。
3. HeliosTransformerBlock:AdaLN 与双层注意力
每个HeliosTransformerBlock由三个子模块组成(对应HeliosAttention定义于 transformer_helios.py 的HeliosAttention类):
- 自注意力(
attn1):Q/K 做 RMSNorm(qk_norm="rms_norm_across_heads"),采用 RoPE 位置编码;可选is_amplify_history模式,对历史 token 的 key 施加可学习的缩放(history_key_scale,支持scalar或per_head两种粒度,最大缩放 10.0),用于蒸馏模型放大"首块"历史信息; - 交叉注意力(
attn2):以文本嵌入为 KV(added_kv_proj_dim提供额外的 KV 投影),支持guidance_cross_attn——此时仅对当前块做交叉注意力,历史块保持原样,避免 CFG 对历史上下文的重复计算; - 前馈网络(
ffn):FeedForward搭配 GELU-approx 激活。
时间条件通过 AdaLN 注入:scale_shift_table(形状(1, 6, dim))与时间投影按通道切分成 6 组 shift/scale/gate,分别调制自注意力前的 LayerNorm、自注意力残差门控、FFN 前的 LayerNorm 与 FFN 残差门控。所有 LayerNorm 均采用 FP32 计算(FP32LayerNorm),并在_keep_in_fp32_modules中列出,以保证 14B 规模下 bf16 训练的数值稳定性。
注意力计算统一走HeliosAttnProcessor,基于scaled_dot_product_attention(要求 PyTorch 2.0+),并通过dispatch_attention_fn支持切换 FlashAttention、SageAttention 等后端(见下文速度优化)。
4. 输出层与反补丁化
HeliosOutputNorm同样使用 AdaLN(shift/scale 由时间嵌入按original_context_length截取),之后proj_out线性层把每个 token 映射回out_channels * prod(patch_size)通道,最后 reshape + permute 完成 unpatchify,恢复为(B, C, T, H, W)的视频潜变量,输出封装为Transformer2DModelOutput(sample=...)(该类定义于 modeling_outputs.py)。
5. 历史时间步归零与梯度检查点
zero_history_timestep=True时,历史上下文 token 的时间步被置为 0(timestep_t0),即历史帧按"无噪声"条件参与计算,这与自回归去噪的语义一致。模型还内置了梯度检查点支持(_supports_gradient_checkpointing=True)与 Layerwise Casting 模式(_skip_layerwise_casting_patterns排除 patch 与条件嵌入层),方便在训练/微调场景下节省显存。
配置参数完整说明
以下参数表依据 transformer_helios.py 中HeliosTransformer3DModel.__init__的完整签名整理(括号内为默认值):
| 参数 | 默认值 | 含义 |
|---|---|---|
patch_size | (1, 2, 2) | 3D patch 尺寸(t, h, w),决定 token 数与显存占用 |
num_attention_heads | 40 | 注意力头数 |
attention_head_dim | 128 | 每个注意力头的通道数 |
in_channels | 16 | 输入(潜变量)通道数 |
out_channels | 16 | 输出通道数,为None时回退为in_channels |
text_dim | 4096 | 文本嵌入输入维度 |
freq_dim | 256 | 正弦时间嵌入频率维度 |
ffn_dim | 13824 | 前馈网络中间维度 |
num_layers | 40 | Transformer 块数量 |
cross_attn_norm | True | 交叉注意力前是否归一化 |
qk_norm | "rms_norm_across_heads" | Q/K 归一化方式 |
eps | 1e-6 | 归一化层 epsilon |
added_kv_proj_dim | None | 附加 KV 投影通道数,用于交叉注意力 |
rope_dim | (44, 42, 42) | RoPE 在 T/Y/X 三个维度上的频率维度 |
rope_theta | 10000.0 | RoPE 频率基数 theta |
guidance_cross_attn | True | 是否使用引导式交叉注意力 |
zero_history_timestep | True | 历史 token 时间步是否置零 |
has_multi_term_memory_patch | True | 是否启用短/中/长多尺度历史记忆分支 |
is_amplify_history | False | 是否对历史 key 施加可学习放大(蒸馏模型使用) |
history_scale_mode | "per_head" | 历史放大粒度:scalar或per_head |
测试用例 tests/models/transformers/test_models_transformer_helios.py 中给出了一个小型化配置参考(2 层、2 头、head_dim=12、ffn_dim=32),可用于验证模型结构与前向形状。
forward 输入输出详解
forward的核心签名(带@apply_lora_scale("attention_kwargs")装饰器,支持 LoRA 缩放):
hidden_states:形状(batch_size, num_channels, num_frames, height, width)的加噪视频潜变量;timestep:去噪步数;encoder_hidden_states:形状(batch_size, sequence_len, embed_dims)的文本条件嵌入;indices_hidden_states:当前帧的帧索引,用于计算 RoPE;indices_latents_history_short/mid/long与latents_history_short/mid/long:三个尺度历史潜变量及其帧索引(Stage-1 自回归上下文);return_dict:为True时返回Transformer2DModelOutput,否则返回元组。
多尺度历史分支的处理顺序为 short → mid → long,依次拼接在序列头部;original_context_length记录拼接前的当前块长度,供输出归一化与引导交叉注意力截取。若indices_hidden_states为None,模型会自动生成torch.arange作为默认帧索引,因此做单步前向调试时这些历史参数都可以省略。
与 Pipeline 集成:T2V / I2V / V2V 与内存优化
HeliosTransformer3DModel 本身只是去噪主干,实际使用需通过 HeliosPipeline(Base)或HeliosPyramidPipeline(Mid/Distilled,基于金字塔流匹配采样,对应pyramid_num_inference_steps_list=[20, 20, 20]或[2, 2, 2])。Pipeline 源码位于 src/diffusers/pipelines/helios/,其中 pipeline_helios_pyramid.py 展示了完整的调用范式:文本由UMT5EncoderModel编码、图像/视频经AutoencoderKLWan编码、HeliosScheduler/HeliosDMDScheduler负责步进采样。
内存优化示例(约 6GB 显存,来自 pipeline 文档):
import torch from diffusers import AutoModel, HeliosPipeline from diffusers.hooks.group_offloading import apply_group_offloading from diffusers.utils import export_to_video vae = AutoModel.from_pretrained("BestWishYsh/Helios-Base", subfolder="vae", dtype=torch.float32) pipeline = HeliosPipeline.from_pretrained("BestWishYsh/Helios-Base", vae=vae, dtype=torch.bfloat16) pipeline.enable_group_offload( onload_device=torch.device("cuda"), offload_device=torch.device("cpu"), offload_type="leaf_level", use_stream=True, record_stream=True, ) output = pipeline( prompt=prompt, negative_prompt=negative_prompt, num_frames=99, num_inference_steps=50, guidance_scale=5.0, generator=torch.Generator("cuda").manual_seed(42), ).frames[0] export_to_video(output, "helios_base_t2v_output.mp4", fps=24)推理速度优化:文档给出的组合是set_attention_backend切换注意力后端(如"_flash_3_hub"适配 Hopper 架构 GPU)+ 对 text_encoder / vae / transformer 分别执行torch.compile(mode="max-autotune-no-cudagraphs", dynamic=False)。从源码看,HeliosAttnProcessor通过dispatch_attention_fn分发到不同后端,这解释了为什么注意力后端可以无侵入切换。此外,上下文并行(Context Parallelism)方案在该模型中有内置支持:_cp_plan在 40 个 block 的attn1/attn2/ffn层定义了输入切分(ContextParallelInput,按序列维 split_dim=1)与输出聚合(ContextParallelOutput),用于长视频跨多卡并行处理。
测试与验证
仓库为该模型提供了完备的测试矩阵(tests/models/transformers/test_models_transformer_helios.py):
TestHeliosTransformer3D:核心前向/形状测试,基于hf-internal-testing/tiny-helios-base-transformer微型权重;TestHeliosTransformer3DMemory:显存优化测试;TestHeliosTransformer3DTraining:训练测试,并显式验证梯度检查点被应用;TestHeliosTransformer3DAttention:注意力处理器测试;TestHeliosTransformer3DCompile:torch.compile 测试(注:在确定性算法下因 PyTorch issue #170079 存在已知编译失败场景);TestHeliosTransformer3DLoRA:LoRA 适配测试。
测试的 dummy 输入完整覆盖了hidden_states、timestep、encoder_hidden_states以及 short/mid/long 三档历史潜变量与索引,可直接作为自定义调用该模型的形状参考。此外 tests/pipelines/helios/test_helios.py 与 tests/modular_pipelines/helios/test_modular_pipeline_helios.py 覆盖了端到端 pipeline 与模块化 pipeline 的集成行为。
相关资源
- 模型文档:docs/source/en/api/models/helios_transformer3d.md
- Pipeline 文档:docs/source/en/api/pipelines/helios.md
- 调度器文档:docs/source/en/api/schedulers/helios.md(HeliosScheduler)与 docs/source/en/api/schedulers/helios_dmd.md(HeliosDMDScheduler)
- 使用指南(中文):docs/source/zh/using-diffusers/helios.md
- 核心实现:src/diffusers/models/transformers/transformer_helios.py
- 模块化 Pipeline 组件:src/diffusers/modular_pipelines/helios/
如需深入研究 Helios 的金字塔采样策略,可继续阅读HeliosPyramidPipeline的pyramid_num_inference_steps_list、use_zero_init、zero_steps等参数(详见 pipeline 文档),其"CFG-Zero"与"放大首块"机制分别与模型配置中的zero_history_timestep和is_amplify_history直接对应。
【免费下载链接】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),仅供参考