Diffusers documentation
Kandinsky 6 Transformers
Kandinsky 6 Transformers
Kandinsky 6 使用多模态扩散 Transformer,对视频和音频潜变量进行联合去噪,以实现文本/图像到视频和音频的生成,并使用无文本扩散 Transformer 进行视频超分辨率。
Kandinsky6Transformer3DModel
Kandinsky6TI2VAPipeline 所使用的多模态 transformer。
import torch
from diffusers import Kandinsky6Transformer3DModel
transformer = Kandinsky6Transformer3DModel.from_pretrained(
"kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers", subfolder="transformer", torch_dtype=torch.bfloat16
)class diffusers.Kandinsky6Transformer3DModel
< source >( in_visual_dim: int = 16out_visual_dim: int = 16in_text_dim: int = 3584in_text_dim2: int = 768time_dim: int = 1024patch_size: tuple = (1, 2, 2)model_dim: int = 4096ff_dim: int = 16384num_text_blocks: int = 4num_visual_blocks: int = 60axes_dims: tuple = (32, 48, 48)visual_cond: bool = Truein_audio_dim: int = 20out_audio_dim: int = 20model_dim_a: int | None = Nonetime_dim_a: int | None = Noneff_dim_a: int | None = Noneaxes_dims_a: tuple | None = Noneaudio_freqs_scaling: float = 1.0scale_factor: tuple | list[float] = (1.0, 2.0, 2.0)text_token_padding: bool = Falseca_rope: bool = Falsecross_gates: bool = Falsefix_modulation: bool = Falsevisual_token_type_num_embeddings: int = 0 )
参数
- in_visual_dim (int,可选,默认为 16) — 输入视频潜在通道的数量。
- out_visual_dim (int,可选,默认为 16) — 输出视频潜在通道的数量。
- in_text_dim (int,可选,默认为 3584) — 文本 token 嵌入维度。
- in_text_dim2 (int,可选,默认为 768) — 池化文本嵌入的维度。
- time_dim (int,可选,默认为 1024) — 时间嵌入维度。
- patch_size (tuple[int, int, int],可选,默认为
(1, 2, 2)) — 视频 patch 大小。 - model_dim (int,可选,默认为 4096) — 视频 Transformer 的隐藏维度。
- ff_dim (int,可选,默认为 16384) — 视频前馈隐藏层维度。
- num_text_blocks (int,可选,默认为 4) — 每个模态的文本块数量。
- num_visual_blocks (int,可选,默认为 60) — 融合的视频/音频块数量。
- axes_dims (tuple[int, int, int],可选,默认为
(32, 48, 48)) — 视频的 RoPE 维度。 - visual_cond (bool,可选,默认为 True) — 指示是否存在视频条件通道。
- in_audio_dim (int,可选,默认为 20) — 输入音频潜在通道的数量。
- out_audio_dim (int,可选,默认为 20) — 输出音频潜在通道的数量。
- model_dim_a (int,可选) — 音频 Transformer 的隐藏维度。默认值为 model_dim。
- time_dim_a (int,可选) — 音频时间嵌入维度。默认为 time_dim。
- ff_dim_a (int,可选) — 音频前馈隐藏维度。默认值为 ff_dim。
- axes_dims_a (tuple[int, int, int],可选) — 音频 RoPE 维度。默认值为 axes_dims。
- audio_freqs_scaling (float,可选,默认为 1.0) — 音频 RoPE 频率缩放。
- scale_factor (tuple[float, float, float],可选,默认为 (1.0, 2.0, 2.0)) — 应用于视频位置的按轴 (t, h, w) RoPE 频率缩放。
- text_token_padding(bool,可选,默认为 False)—— 记录参考模型的文本序列是否经过填充的 Checkpoint 元数据。无论此标志如何,forward 始终接受一个可选的 encoder_attention_mask;是否实际传入该参数完全由调用方决定。
- ca_rope (bool,可选,默认为 False) — 是否使用跨模态音频 RoPE。
- cross_gates(bool,可选,默认为 False)——是否使用跨模态残差门控。
- fix_modulation (bool,可选,默认为 False) — 是否使用固定调制变体。
- visual_token_type_num_embeddings (int,可选,默认为 0) — 视觉 token 类型嵌入的数量。
Kandinsky 6 多模态 Transformer,用于文本/图像到视频和音频的生成。
视频和音频通过融合的跨模态 Transformer 块共同进行去噪,每个块都以其自身的文本分支(Qwen2.5-VL 标记 + CLIP 池化嵌入)作为条件。如果不传入 audio_hidden_states,则仅对视频进行去噪,但仍会执行融合块中的视频自注意力、交叉注意力和前馈阶段。旋转位置嵌入在 forward 内部根据标记网格计算得出。
发布的检查点将 text_token_padding、ca_rope、cross_gates 和 fix_modulation 设置为 True(参见各检查点的 transformer/config.json);False 的默认值仅描述了一种本仓库未提供权重的架构变体。
forward
< source >( hidden_states: Tensorencoder_hidden_states: Tensorpooled_projections: Tensortimestep: Tensoraudio_hidden_states: Tensor | None = Nonevisual_rope_pos: tuple[Tensor, Tensor, Tensor] | None = Noneencoder_attention_mask: Tensor | None = Nonevisual_token_type_ids: Tensor | None = Nonereturn_dict: bool = True ) → Kandinsky6TransformerOutput or tuple
参数
- hidden_states(形状为
(batch_size, num_frames, height, width, in_channels)的torch.Tensor)— 采用(B, T, H, W, C)布局的视频潜在表示。当visual_cond=True时,in_channels为2 * in_visual_dim + 1:包含带噪潜在表示、条件潜在表示以及条件掩码。 - encoder_hidden_states(形状为
(batch_size, sequence_length, in_text_dim)的torch.Tensor)— 文本 token 嵌入,由视频和音频的文本分支共享。 - pooled_projections(形状为
(batch_size, in_text_dim2)的torch.Tensor)— 添加到两个分支的时间步嵌入中的池化文本嵌入。 - timestep(形状为
(batch_size,)的torch.Tensor)— 在[0, num_train_timesteps]尺度上的扩散时间步,两种模态共用。 - audio_hidden_states (
torch.Tensor,形状为(batch_size, audio_length, in_audio_dim),可选) — 音频潜在表示。如果省略,融合块将仅执行视频路径。 - visual_rope_pos (
tuple[torch.Tensor, torch.Tensor, torch.Tensor], 可选) — 分块视频 token 沿各轴(t, h, w)的旋转位置索引。默认情况下,对每个轴使用arange。当流水线添加一个复用位置0的参考帧时,会传入显式的时间索引。 - encoder_attention_mask(形状为
(batch_size, sequence_length)的torch.Tensor,可选) —encoder_hidden_states的布尔填充掩码。当提示词未进行填充时,请传入None。 - visual_token_type_ids(形状为
(batch_size, num_frames)的torch.Tensor,可选) — 逐帧的 token 类型 ID,通过visual_token_type_embeddings进行嵌入。要求visual_token_type_num_embeddings > 0。 - return_dict (
bool,默认为True) — 是否返回Kandinsky6TransformerOutput而不是普通元组。
返回值
Kandinsky6TransformerOutput 或 tuple
在 (B, T, H, W, out_visual_dim) 布局下的预测视频速度,以及当提供 audio_hidden_states 时,形状为 (B, audio_length, out_audio_dim) 的预测音频速度。
Kandinsky6SRTransformer3DModel
Kandinsky6SRPipeline 用于逐块精修放大后视频的无文本 Transformer。
import torch
from diffusers import Kandinsky6SRTransformer3DModel
transformer = Kandinsky6SRTransformer3DModel.from_pretrained(
"kandinskylab/Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers", subfolder="transformer", torch_dtype=torch.bfloat16
)
# The transformer always runs NABLA sparse attention (`nabla_threshold`, 0.8 by default) on the `flex` backend.
# Compile it, otherwise flex falls back to an eager implementation that needs far more memory at video resolutions.
transformer.compile_repeated_blocks(fullgraph=True)class diffusers.Kandinsky6SRTransformer3DModel
< source >( in_visual_dim: int = 64out_visual_dim: int = 640time_dim: int = 512patch_size: tuple[int, int, int] = (1, 1, 1)model_dim: int = 1792ff_dim: int = 7168num_visual_blocks: int = 32axes_dims: tuple[int, int, int] = (16, 24, 24)scale_factor: tuple[float, float, float] = (1.0, 2.0, 2.0)nabla_threshold: float = 0.8nabla_window: tuple[int, int, int] = (11, 7, 7)tile_sizes: tuple[tuple[int, int], ...] = ((512, 512), (512, 768), (768, 512)) )
参数
- in_visual_dim (
int,默认为64) — K-VAE 的潜在通道数。 - out_visual_dim (
int,默认为640) — 输出通道数。 - time_dim (
int,默认为512) — 时间步嵌入的维度。 - patch_size (
tuple[int, int, int],默认为(1, 1, 1)) — 以(temporal, height, width)形式表示的 patch 大小。 - model_dim (
int,默认为1792) — Transformer 的隐藏维度。 - ff_dim (
int,默认为7168) — 前馈网络的内部维度。 - num_visual_blocks (
int,默认为32) — Transformer 块的数量。 - axes_dims (
tuple[int, int, int],默认为(16, 24, 24)) —(t, h, w)各轴上的 RoPE 维度;它们的总和即为注意力头维度。 - scale_factor (
tuple[float, float, float],默认为(1.0, 2.0, 2.0)) — 应用于 token 位置的逐轴 RoPE 频率缩放。 - nabla_threshold (
float,默认为0.8) — NABLA 块选择的累积注意力阈值。 - nabla_window (
tuple[int, int, int],默认为(11, 7, 7)) — 每个 8x8 token 块始终关注的滑动 tile 先验的奇数(t, h, w)范围。 - tile_sizes (
tuple[tuple[int, int], ...],默认为((512, 512), (512, 768), (768, 512))) — 模型训练时使用的视频瓦片的像素(height, width)尺寸。Kandinsky6SRPipeline 会以宽高比最接近输入视频的瓦片尺寸来细化每个瓦片。
用于 Kandinsky 6 视频超分辨率的无文本扩散 Transformer。
该模型对 K-VAE 潜在视频的图块进行去噪。其输入将含噪潜在表示与用于条件化超分辨率的锚点潜在表示及锚点掩码进行拼接(2 * in_visual_dim + 1 个通道),其输出包含 out_visual_dim 个通道:对于普通 velocity 检查点为 in_visual_dim,对于 PiflowScheduler 蒸馏检查点则为扩展后的 n_grid * in_visual_dim。视频自注意力通过 flex 注意力后端上的 NABLA 稀疏块模式运行,这要求令牌网格(patch_size 之后的 height 和 width)能被 8 整除。
forward
< source >( hidden_states: Tensortimestep: Tensorreturn_dict: bool = True ) → Transformer2DModelOutput or tuple
参数
- hidden_states (
torch.Tensor,形状为(batch_size, num_frames, height, width, 2 * in_visual_dim + 1)) — 采用(B, T, H, W, C)布局的潜在瓦片:沿通道轴拼接的含噪潜在值、锚点潜在值和锚点掩码。 - timestep(形状为
(batch_size,)的torch.Tensor)— 在[0, num_train_timesteps]尺度上的扩散时间步。 - return_dict (
bool,默认为True) — 是否返回 Transformer2DModelOutput 而不是普通元组。
返回值
Transformer2DModelOutput 或 tuple
形状为 (batch_size, num_frames, height, width, out_visual_dim) 的预测结果。