Diffusers documentation
SanaVideoTransformer3DModel
SanaVideoTransformer3DModel
来自 NVIDIA 和 MIT HAN Lab 的 SANA-Video: Efficient Video Generation with Block Linear Diffusion Transformer 中用于 3D 数据（视频）的 Diffusion Transformer 模型，作者为 Junsong Chen, Yuyang Zhao, Jincheng Yu, Ruihang Chu, Junyu Chen, Shuai Yang, Xianbang Wang, Yicheng Pan, Daquan Zhou, Huan Ling, Haozhe Liu, Hongwei Yi, Hao Zhang, Muyang Li, Yukang Chen, Han Cai, Sanja Fidler, Ping Luo, Song Han, Enze Xie。
该论文的摘要如下：
我们推出了 SANA-Video，一个小型扩散模型，能够高效生成分辨率高达 720x1280、时长可达分钟级的视频。SANA-Video 以极快的速度合成高分辨率、高质量且长时长的视频，并具备强大的文本-视频对齐能力，可部署在 RTX 5090 GPU 上。两项核心设计确保了高效、有效且长时长的视频生成：(1) Linear DiT：我们利用线性注意力作为核心操作，鉴于视频生成中处理的 token 数量庞大，这比原生注意力机制更高效。(2) 用于块线性注意力的恒定内存 KV 缓存：我们通过采用源自线性注意力累积特性的恒定内存状态，设计了用于长视频生成的块级自回归方法。该 KV 缓存以固定的内存成本为 Linear DiT 提供全局上下文，消除了对传统 KV 缓存的需求，从而实现了高效的分钟级视频生成。此外，我们探索了有效的数据过滤和模型训练策略，将训练成本降低至 64 块 H100 GPU 上 12 天，仅为 MovieGen 成本的 1%。凭借低成本，SANA-Video 在实测延迟上快 16 倍的同时，相较于现代最先进的中小型扩散模型（例如 Wan 2.1-1.3B 和 SkyReel-V2-1.3B）实现了具有竞争力的性能。此外，SANA-Video 可以以 NVFP4 精度部署在 RTX 5090 GPU 上，将生成 5 秒 720p 视频的推理速度从 71 秒加速至 29 秒（2.4 倍加速）。总之，SANA-Video 实现了低成本、高质量的视频生成。
可以使用以下代码片段加载该模型。
from diffusers import SanaVideoTransformer3DModel
import torch
transformer = SanaVideoTransformer3DModel.from_pretrained("Efficient-Large-Model/SANA-Video_2B_480p_diffusers", subfolder="transformer", dtype=torch.bfloat16)SanaVideoTransformer3DModel
class diffusers.SanaVideoTransformer3DModel
< source >( in_channels: int = 16out_channels: int | None = 16num_attention_heads: int = 20attention_head_dim: int = 112num_layers: int = 20num_cross_attention_heads: int | None = 20cross_attention_head_dim: int | None = 112cross_attention_dim: int | None = 2240caption_channels: int = 2304mlp_ratio: float = 2.5dropout: float = 0.0attention_bias: bool = Falsesample_size: int = 30patch_size: tuple = (1, 2, 2)norm_elementwise_affine: bool = Falsenorm_eps: float = 1e-06interpolation_scale: int | None = Noneguidance_embeds: bool = Falseguidance_embeds_scale: float = 0.1qk_norm: str | None = 'rms_norm_across_heads'rope_max_seq_len: int = 1024 )
参数
- in_channels (
int，默认为16) — 输入数据的通道数。 - out_channels (
int，可选，默认为16) — 输出中的通道数。 - num_attention_heads (
int，默认为20) — 用于多头注意力的注意力头数量。 - attention_head_dim (
int，默认为112) — 每个注意力头中的通道数。 - num_layers（
int，默认为20）— 要使用的 Transformer 块层数。 - num_cross_attention_heads (
int，可选，默认为20) — 用于交叉注意力的注意力头数量。 - cross_attention_head_dim (
int，可选，默认为112) — 交叉注意力中每个头的通道数。 - cross_attention_dim (
int，可选，默认为2240) — 交叉注意力输出的通道数。 - caption_channels (
int，默认为2304) — caption 嵌入的通道数。 - mlp_ratio (
float，默认为2.5) — GLUMBConv 层中使用的扩展比率。 - dropout (
float，默认为0.0) — Dropout 概率。 - attention_bias (
bool，默认为False) — 是否在注意力层中使用偏置。 - sample_size (
int，默认为32) — 输入潜变量的基础尺寸。 - patch_size (
int，默认为1) — 用于 patch 嵌入层的 patch 大小。 - norm_elementwise_affine (
bool，默认为False) — 是否在归一化层中使用逐元素仿射变换。 - norm_eps (
float，默认为1e-6) — 归一化层的 epsilon 值。 - qk_norm (
str, 可选，默认为None) — 用于查询和键的归一化方法。
在 Sana-Video 系列模型中引入的 3D Transformer 模型。
forward
< source >( hidden_states: Tensorencoder_hidden_states: Tensortimestep: Tensorguidance: typing.Optional[torch.Tensor] = Noneencoder_attention_mask: typing.Optional[torch.Tensor] = Noneattention_mask: typing.Optional[torch.Tensor] = Noneattention_kwargs: dict[str, typing.Any] | None = Nonecontrolnet_block_samples: tuple[torch.Tensor] | None = Nonereturn_dict: bool = True )
参数
- hidden_states（形状为
(batch_size, in_channels, num_frames, height, width)的torch.Tensor）— 输入的hidden_states。 - encoder_hidden_states (
torch.Tensor，形状为(batch_size, sequence_len, embed_dims)) — 要使用的条件嵌入（从输入条件（例如提示词）计算得出的嵌入）。 - timestep (
torch.LongTensor) — 用于指示去噪步骤。 - guidance (
torch.Tensor，可选) — 引导尺度嵌入。 - encoder_attention_mask (
torch.Tensor, 可选) — 应用于encoder_hidden_states的交叉注意力掩码。 - attention_mask (
torch.Tensor, 可选) — 应用于hidden_states的自注意力掩码。 - attention_kwargs (
dict, 可选) — 一个 kwargs 字典，如果指定，将传递给 diffusers.models.attention_processor 中定义的AttentionProcessor，该处理器在self.processor下定义。 - controlnet_block_samples（
tupleoftorch.Tensor，可选）— 一个张量列表，如果指定，将被添加到 transformer 块的残差中。 - return_dict (
bool，可选，默认为True) — 是否返回~models.transformer_2d.Transformer2DModelOutput而不是普通的元组。
SanaVideoTransformer3DModel 的前向方法。
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 的输出。