Hugging Face's logo

Diffusers documentation

Kandinsky 6 Transformers

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.41.0).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Kandinsky 6 Transformers

Kandinsky 6 uses a multimodal diffusion transformer that denoises video and audio latents together for text/image-to-video-and-audio generation, and a text-free diffusion transformer for video super-resolution.

Kandinsky6Transformer3DModel

The multimodal transformer used by Kandinsky6TI2VAPipeline.

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

< >

( 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 )

Parameters

  • in_visual_dim (int, optional, defaults to 16) — Number of input video latent channels.
  • out_visual_dim (int, optional, defaults to 16) — Number of output video latent channels.
  • in_text_dim (int, optional, defaults to 3584) — Text token embedding dimension.
  • in_text_dim2 (int, optional, defaults to 768) — Pooled text embedding dimension.
  • time_dim (int, optional, defaults to 1024) — Time embedding dimension.
  • patch_size (tuple[int, int, int], optional, defaults to (1, 2, 2)) — Video patch size.
  • model_dim (int, optional, defaults to 4096) — Video transformer hidden dimension.
  • ff_dim (int, optional, defaults to 16384) — Video feed-forward hidden dimension.
  • num_text_blocks (int, optional, defaults to 4) — Number of text blocks per modality.
  • num_visual_blocks (int, optional, defaults to 60) — Number of fused video/audio blocks.
  • axes_dims (tuple[int, int, int], optional, defaults to (32, 48, 48)) — RoPE dimensions for video.
  • visual_cond (bool, optional, defaults to True) — Whether video conditioning channels are present.
  • in_audio_dim (int, optional, defaults to 20) — Number of input audio latent channels.
  • out_audio_dim (int, optional, defaults to 20) — Number of output audio latent channels.
  • model_dim_a (int, optional) — Audio transformer hidden dimension. Defaults to model_dim.
  • time_dim_a (int, optional) — Audio time embedding dimension. Defaults to time_dim.
  • ff_dim_a (int, optional) — Audio feed-forward hidden dimension. Defaults to ff_dim.
  • axes_dims_a (tuple[int, int, int], optional) — Audio RoPE dimensions. Defaults to axes_dims.
  • audio_freqs_scaling (float, optional, defaults to 1.0) — Audio RoPE frequency scaling.
  • scale_factor (tuple[float, float, float], optional, defaults to (1.0, 2.0, 2.0)) — Per-axis (t, h, w) RoPE frequency scaling applied to the video positions.
  • text_token_padding (bool, optional, defaults to False) — Checkpoint metadata recording whether the reference model’s text sequences are padded. forward always accepts an optional encoder_attention_mask regardless of this flag; whether one is actually passed is entirely up to the caller.
  • ca_rope (bool, optional, defaults to False) — Whether to use cross-modal audio RoPE.
  • cross_gates (bool, optional, defaults to False) — Whether to use cross-modal residual gates.
  • fix_modulation (bool, optional, defaults to False) — Whether to use the fixed modulation variant.
  • visual_token_type_num_embeddings (int, optional, defaults to 0) — Number of visual token type embeddings.

Kandinsky 6 multimodal transformer for text/image-to-video-and-audio generation.

Video and audio are denoised together through fused, cross-modal transformer blocks, each conditioned on its own text branch (Qwen2.5-VL tokens + CLIP pooled embedding). Passing no audio_hidden_states denoises video alone while still running the fused block’s video self-attention, cross-attention, and feed-forward stages. Rotary embeddings are computed inside forward from the token grid.

Released checkpoints set text_token_padding, ca_rope, cross_gates, and fix_modulation to True (see each checkpoint’s transformer/config.json); the False defaults only describe an architecture variant this repo does not ship weights for.

forward

< >

( 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

Parameters

  • hidden_states (torch.Tensor of shape (batch_size, num_frames, height, width, in_channels)) — Video latents in the (B, T, H, W, C) layout. With visual_cond=True, in_channels is 2 * in_visual_dim + 1: the noisy latent, the conditioning latent and a conditioning mask.
  • encoder_hidden_states (torch.Tensor of shape (batch_size, sequence_length, in_text_dim)) — Text token embeddings, shared by the video and audio text branches.
  • pooled_projections (torch.Tensor of shape (batch_size, in_text_dim2)) — Pooled text embedding added to the timestep embedding of both branches.
  • timestep (torch.Tensor of shape (batch_size,)) — Diffusion timestep on the [0, num_train_timesteps] scale, shared by both modalities.
  • audio_hidden_states (torch.Tensor of shape (batch_size, audio_length, in_audio_dim), optional) — Audio latents. When omitted the fused blocks run their video path only.
  • visual_rope_pos (tuple[torch.Tensor, torch.Tensor, torch.Tensor], optional) — Per-axis (t, h, w) rotary position indices of the patchified video tokens. Defaults to arange over each axis. The pipeline passes an explicit temporal index when it appends a reference frame that reuses position 0.
  • encoder_attention_mask (torch.Tensor of shape (batch_size, sequence_length), optional) — Boolean padding mask for encoder_hidden_states. Pass None when no prompt is padded.
  • visual_token_type_ids (torch.Tensor of shape (batch_size, num_frames), optional) — Per-frame token type ids, embedded through visual_token_type_embeddings. Requires visual_token_type_num_embeddings > 0.
  • return_dict (bool, defaults to True) — Whether to return a Kandinsky6TransformerOutput instead of a plain tuple.

Returns

Kandinsky6TransformerOutput or tuple

The predicted video velocity in the (B, T, H, W, out_visual_dim) layout and, when audio_hidden_states was given, the predicted audio velocity of shape (B, audio_length, out_audio_dim).

Kandinsky6SRTransformer3DModel

The text-free transformer used by Kandinsky6SRPipeline to refine one tile of the upscaled video at a time.

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

< >

( 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)) )

Parameters

  • in_visual_dim (int, defaults to 64) — Number of latent channels of the K-VAE.
  • out_visual_dim (int, defaults to 640) — Number of output channels.
  • time_dim (int, defaults to 512) — Dimension of the timestep embedding.
  • patch_size (tuple[int, int, int], defaults to (1, 1, 1)) — Patch size as (temporal, height, width).
  • model_dim (int, defaults to 1792) — Hidden dimension of the transformer.
  • ff_dim (int, defaults to 7168) — Inner dimension of the feed-forward networks.
  • num_visual_blocks (int, defaults to 32) — Number of transformer blocks.
  • axes_dims (tuple[int, int, int], defaults to (16, 24, 24)) — RoPE dimensions per (t, h, w) axis; their sum is the attention head dimension.
  • scale_factor (tuple[float, float, float], defaults to (1.0, 2.0, 2.0)) — Per-axis RoPE frequency scaling applied to the token positions.
  • nabla_threshold (float, defaults to 0.8) — Cumulative-attention threshold of the NABLA block selection.
  • nabla_window (tuple[int, int, int], defaults to (11, 7, 7)) — Odd (t, h, w) extents of the sliding-tile prior that every 8x8 token block always attends to.
  • tile_sizes (tuple[tuple[int, int], ...], defaults to ((512, 512), (512, 768), (768, 512))) — Pixel (height, width) sizes of the video tiles the model was trained on. Kandinsky6SRPipeline refines every tile at the size whose aspect ratio is closest to the input video’s.

Text-free diffusion transformer for Kandinsky 6 video super-resolution.

The model denoises tiles of the K-VAE latent video. Its input concatenates the noisy latent with the anchor latent and the anchor mask that condition the super-resolution (2 * in_visual_dim + 1 channels), and its output holds out_visual_dim channels: in_visual_dim for a plain velocity checkpoint, or a widened n_grid * in_visual_dim for the PiflowScheduler distilled checkpoints. Video self-attention runs through the NABLA sparse block pattern on the flex attention backend, which needs the token grid (height and width after patch_size) to be divisible by 8.

forward

< >

( hidden_states: Tensortimestep: Tensorreturn_dict: bool = True ) → Transformer2DModelOutput or tuple

Parameters

  • hidden_states (torch.Tensor of shape (batch_size, num_frames, height, width, 2 * in_visual_dim + 1)) — Latent tiles in the (B, T, H, W, C) layout: the noisy latent, the anchor latent and the anchor mask concatenated along the channel axis.
  • timestep (torch.Tensor of shape (batch_size,)) — Diffusion timestep on the [0, num_train_timesteps] scale.
  • return_dict (bool, defaults to True) — Whether to return a Transformer2DModelOutput instead of a plain tuple.

Returns

Transformer2DModelOutput or tuple

The prediction of shape (batch_size, num_frames, height, width, out_visual_dim).

Update on GitHub

←JoyImageEditTransformer3DModel Krea2Transformer2DModel→