Hugging Face's logo

Diffusers documentation

Kandinsky 6 Transformers

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→