Diffusers documentation
Kandinsky 6 Transformers
Kandinsky 6 Transformers
Kandinsky 6 は、テキスト/画像から動画と音声の生成を行う多モーダル拡散トランスフォーマー(動画と音声の潜在変数を同時にノイズ除去する)と、動画の超解像を行うテキストフリーの拡散トランスフォーマーを使用しています。
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、optional、デフォルトは 16) — 出力される動画潜在変数のチャネル数です。
- in_text_dim (int、任意、デフォルトは 3584) — テキストトークンの埋め込み次元数。
- in_text_dim2 (int、任意、デフォルトは 768) — プールされたテキスト埋め込みの次元数。
- time_dim (int、optional、デフォルトは 1024) — 時間埋め込みの次元数。
- patch_size (tuple[int, int, int]、optional、デフォルトは
(1, 2, 2)) — ビデオのパッチサイズ。 - model_dim (int、optional、デフォルトは 4096) — Video transformer の隠れ次元。
- ff_dim (int、optional、デフォルト値は 16384) — Video feed-forward の隠れ層の次元数。
- num_text_blocks (int、optional、デフォルト値は 4)— モダリティごとのテキストブロック数。
- num_visual_blocks (int、任意、デフォルトは 60) — 融合された動画/音声ブロックの数。
- axes_dims (tuple[int, int, int]、optional、デフォルトは
(32, 48, 48)) — 動画用の RoPE 次元です。 - visual_cond (bool、optional、デフォルトは True)— 動画の条件付けチャネルが存在するかどうかを示します。
- in_audio_dim (int、任意、デフォルト値は 20) — 入力オーディオ潜在チャネルの数。
- out_audio_dim (int、任意、デフォルトは 20) — 出力オーディオ潜在チャネルの数。
- model_dim_a (int、任意) — Audio transformer の隠れ次元。デフォルト値は model_dim です。
- time_dim_a (int、任意) — 音声の時間埋め込み次元。デフォルトは time_dim です。
- ff_dim_a (int、任意) — 音声フィードフォワードの隠れ層の次元数。デフォルト値は ff_dim です。
- axes_dims_a (tuple[int, int, int]、任意) — Audio RoPE の次元。デフォルト値は axes_dims です。
- audio_freqs_scaling (float、optional、デフォルト値は 1.0) — Audio RoPE の周波数スケーリング。
- scale_factor (tuple[float, float, float]、optional、デフォルトは (1.0, 2.0, 2.0)) — 動画の位置に適用される、軸ごとの (t, h, w) RoPE 周波数スケーリング。
- text_token_padding (bool、optional、デフォルトは False) — 参照モデルのテキストシーケンスがパディングされているかどうかを記録するチェックポイントのメタデータです。forward はこのフラグに関係なく、常に任意の encoder_attention_mask を受け取ります。実際に渡すかどうかは呼び出し側が完全に判断します。
- ca_rope (bool、任意、デフォルトは False) — クロスモーダルオーディオ RoPE を使用するかどうかを指定します。
- cross_gates (bool、任意、デフォルトは False) — クロスモーダル残差ゲートを使用するかどうか。
- fix_modulation (bool、任意、デフォルトは False) — 固定 modulation バリアントを使用するかどうか。
- visual_token_type_num_embeddings (int、optional、デフォルト値は 0) — 視覚トークンタイプ埋め込みの数。
Kandinsky 6 は、テキスト/画像から動画とオーディオを生成するためのマルチモーダル・トランスフォーマーです。
動画と音声は、融合されたクロスモーダル 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 (
torch.Tensor、shape(batch_size, num_frames, height, width, in_channels)) —(B, T, H, W, C)レイアウトの動画潜在変数。visual_cond=Trueの場合、in_channelsは2 * in_visual_dim + 1となり、ノイズ付き潜在変数、条件付け潜在変数、および条件付けマスクを表します。 - encoder_hidden_states (
torch.Tensor、形状(batch_size, sequence_length, in_text_dim)) — テキストトークンの埋め込みです。動画と音声のテキストブランチで共有されます。 - pooled_projections(
torch.Tensor、形状(batch_size, in_text_dim2))— 両方のブランチのタイムステップ埋め込みに加算されるプーリング済みテキスト埋め込み。 - timestep (
torch.Tensor、shape(batch_size,)) —[0, num_train_timesteps]のスケールにおける Diffusion の timestep で、両方のモダリティで共有されます。 - audio_hidden_states (
torch.Tensor、形状(batch_size, audio_length, in_audio_dim)、省略可) — オーディオの潜在変数です。省略した場合、融合ブロックは動画パスのみを実行します。 - visual_rope_pos (
tuple[torch.Tensor, torch.Tensor, torch.Tensor], 省略可) — パッチ化された動画トークンの軸別(t, h, w)回転位置インデックスです。デフォルトでは、各軸に対してarangeが使用されます。パイプラインは、位置0を再利用する参照フレームを追加する際に、明示的な時間インデックスを渡します。 - encoder_attention_mask (
torch.Tensor、形状(batch_size, sequence_length)、任意) —encoder_hidden_states用のブール値パディングマスクです。プロンプトにパディングがない場合はNoneを渡してください。 - visual_token_type_ids (
torch.Tensor、shape(batch_size, num_frames)、任意) — フレームごとのトークンタイプ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 が、アップスケールされた動画のタイルを1つずつ精緻化するために使用するテキストフリーの 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)形式のパッチサイズ。 - 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)) — トークン位置に適用される、軸ごとの RoPE 周波数スケーリングです。 - nabla_threshold (
float、デフォルトは0.8) — NABLAブロック選択における累積アテンションの閾値。 - nabla_window (
tuple[int, int, int]、デフォルトは(11, 7, 7)) — 各 8x8 トークンブロックが常にアテンションを行う、スライディングタイル事前分布の奇数(t, h, w)範囲です。 - tile_sizes (
tuple[tuple[int, int], ...]、デフォルトは((512, 512), (512, 768), (768, 512))) — モデルが学習された動画タイルのピクセル(height, width)サイズです。Kandinsky6SRPipeline は、入力動画のアスペクト比に最も近いサイズのタイルで、すべてのタイルを精細化します。
Kandinsky 6のビデオ超解像度向けのテキストフリー拡散トランスフォーマー。
モデルは、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、shape(batch_size, num_frames, height, width, 2 * in_visual_dim + 1)) —(B, T, H, W, C)レイアウトの潜在タイルです。ノイズ付き潜在変数、アンカー潜在変数、アンカーマスクをチャネル軸方向に連結したものです。 - timestep (
torch.Tensor、shape(batch_size,)) —[0, num_train_timesteps]スケール上の Diffusion timestep。 - return_dict (
bool, defaults toTrue) — Whether to return a Transformer2DModelOutput instead of a plain tuple.
戻り値
Transformer2DModelOutput または tuple
形状(batch_size, num_frames, height, width, out_visual_dim)の予測。