Skip to content

vllm_omni.diffusion.models.sana_video

Modules:

Name Description
pipeline_output
pipeline_sana_video
pipeline_sana_video_i2v
transformer_sana_video

SanaImageToVideoPipeline

Bases: SanaVideoPipeline, SupportImageInput

Native vLLM-Omni SANA-Video 2B image-to-video pipeline.

The first latent frame is encoded from the input image and remains fixed; denoising updates only the subsequent video frames.

diffuse 

diffuse(
    latents: Tensor,
    timesteps: Tensor,
    prompt_embeds: Tensor,
    prompt_attention_mask: Tensor,
    negative_prompt_embeds: Tensor | None,
    negative_prompt_attention_mask: Tensor | None,
    guidance_scale: float,
    conditioning_mask: Tensor,
    extra_step_kwargs: dict,
    dtype: dtype,
    output_slice: int | None,
) -> Tensor

forward 

forward(req: DiffusionRequestBatch) -> DiffusionOutput

SanaVideoPipeline

Bases: Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery

vLLM-Omni pipeline for text-to-video generation using Sana.

Parameters:

Name Type Description Default
tokenizer [`GemmaTokenizer`] or [`GemmaTokenizerFast`]

The tokenizer used to tokenize the prompt.

 None
text_encoder [`Gemma2PreTrainedModel`]

Text encoder model to encode the input prompts.

 None
vae [`DistributedAutoencoderKLWan`] or [`DistributedAutoencoderKLLTX2Video`]

Variational Auto-Encoder (VAE) Model to encode and decode videos to and from latent representations.

 None
transformer [`SanaVideoTransformer3DModel`]

Conditional Transformer to denoise the input latents.

 None
scheduler [`DPMSolverMultistepScheduler`]

A scheduler to be used in combination with transformer to denoise the encoded video latents.

 None

default_num_inference_steps class-attribute instance-attribute 

default_num_inference_steps = 50

device instance-attribute 

device = get_local_device()

do_classifier_free_guidance property 

do_classifier_free_guidance

dummy_run_num_frames class-attribute instance-attribute 

dummy_run_num_frames = 25

guidance_scale property 

guidance_scale

interrupt property 

interrupt

num_timesteps property 

num_timesteps

od_config instance-attribute 

od_config = od_config

scheduler instance-attribute 

scheduler = scheduler

supports_step_execution class-attribute instance-attribute 

supports_step_execution = False

text_encoder instance-attribute 

text_encoder = text_encoder

tokenizer instance-attribute 

tokenizer = tokenizer

transformer instance-attribute 

transformer = transformer

vae instance-attribute 

vae = vae

vae_scale_factor instance-attribute 

vae_scale_factor = self.vae_scale_factor_spatial

vae_scale_factor_spatial instance-attribute 

vae_scale_factor_spatial = (
    self.vae.config.spatial_compression_ratio
)

vae_scale_factor_temporal instance-attribute 

vae_scale_factor_temporal = (
    self.vae.config.temporal_compression_ratio
)

video_processor instance-attribute 

video_processor = VideoProcessor(
    vae_scale_factor=self.vae_scale_factor_spatial
)

weights_sources instance-attribute 

weights_sources: list[ComponentSource] = []

check_inputs 

check_inputs(
    prompt,
    height,
    width,
    negative_prompt=None,
    prompt_embeds=None,
    negative_prompt_embeds=None,
    prompt_attention_mask=None,
    negative_prompt_attention_mask=None,
)

combine_cfg_noise 

combine_cfg_noise(
    positive_noise_pred,
    negative_noise_pred,
    true_cfg_scale,
    cfg_normalize=False,
    kwargs=None,
)

diffuse 

diffuse(
    latents: Tensor,
    timesteps: Tensor,
    prompt_embeds: Tensor,
    prompt_attention_mask: Tensor,
    negative_prompt_embeds: Tensor | None,
    negative_prompt_attention_mask: Tensor | None,
    guidance_scale: float,
    extra_step_kwargs: dict,
    dtype: dtype,
    output_slice: int | None,
) -> Tensor

encode_prompt 

encode_prompt(
    prompt: str | list[str],
    do_classifier_free_guidance: bool = True,
    negative_prompt: str = "",
    num_videos_per_prompt: int = 1,
    device: device | None = None,
    prompt_embeds: Tensor | None = None,
    negative_prompt_embeds: Tensor | None = None,
    prompt_attention_mask: Tensor | None = None,
    negative_prompt_attention_mask: Tensor | None = None,
    clean_caption: bool = False,
    max_sequence_length: int = 300,
    complex_human_instruction: list[str] | None = None,
)

Encodes the prompt into text encoder hidden states.

Parameters:

Name Type Description Default
prompt `str` or `list[str]`, *optional*

prompt to be encoded

 required
negative_prompt `str` or `list[str]`, *optional*

The prompt not to guide the video generation. If not defined, one has to pass negative_prompt_embeds instead. Ignored when not using guidance (i.e., ignored if guidance_scale is less than 1). For PixArt-Alpha, this should be "".

 ''
do_classifier_free_guidance `bool`, *optional*, defaults to `True`

whether to use classifier free guidance or not

 True
num_videos_per_prompt `int`, *optional*, defaults to 1

number of videos that should be generated per prompt

 1
device device | None

(torch.device, optional): torch device to place the resulting embeddings on

 None
prompt_embeds `torch.Tensor`, *optional*

Pre-generated text embeddings. Can be used to easily tweak text inputs, e.g. prompt weighting. If not provided, text embeddings will be generated from prompt input argument.

 None
negative_prompt_embeds `torch.Tensor`, *optional*

Pre-generated negative text embeddings. For Sana, it's should be the embeddings of the "" string.

 None
clean_caption `bool`, defaults to `False`

If True, the function will preprocess and clean the provided caption before encoding.

 False
max_sequence_length `int`, defaults to 300

Maximum sequence length to use for the prompt.

 300
complex_human_instruction `list[str]`, defaults to `complex_human_instruction`

If complex_human_instruction is not empty, the function will use the complex Human instruction for the prompt.

 None

forward 

forward(req: DiffusionRequestBatch) -> DiffusionOutput

load_weights 

load_weights(
    weights: Iterable[tuple[str, Tensor]],
) -> set[str]

predict_noise 

predict_noise(*args, **kwargs) -> Tensor

prepare_extra_step_kwargs 

prepare_extra_step_kwargs(generator, eta)

prepare_latents 

prepare_latents(
    batch_size: int,
    num_channels_latents: int = 16,
    height: int = 480,
    width: int = 832,
    num_frames: int = 81,
    dtype: dtype | None = None,
    device: device | None = None,
    generator: Generator | list[Generator] | None = None,
    latents: Tensor | None = None,
) -> Tensor

SanaVideoPipelineOutput dataclass

Bases: BaseOutput

Output class for Sana-Video pipelines.

Parameters:

Name Type Description Default
frames `torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]

List of video outputs - It can be a nested list of length batch_size, with each sub-list containing denoised PIL image sequences of length num_frames. It can also be a NumPy array or Torch tensor of shape (batch_size, num_frames, channels, height, width).

 required

frames instance-attribute 

frames: Tensor

SanaVideoTransformer3DModel

Bases: Module

A 3D Transformer model introduced in Sana-Video family of models.

Parameters:

Name Type Description Default
in_channels `int`, defaults to `16`

The number of channels in the input.

 16
out_channels `int`, *optional*, defaults to `16`

The number of channels in the output.

 16
num_attention_heads `int`, defaults to `20`

The number of heads to use for multi-head attention.

 20
attention_head_dim `int`, defaults to `112`

The number of channels in each head.

 112
num_layers `int`, defaults to `20`

The number of layers of Transformer blocks to use.

 20
num_cross_attention_heads `int`, *optional*, defaults to `20`

The number of heads to use for cross-attention.

 20
cross_attention_head_dim `int`, *optional*, defaults to `112`

The number of channels in each head for cross-attention.

 112
cross_attention_dim `int`, *optional*, defaults to `2240`

The number of channels in the cross-attention output.

 2240
caption_channels `int`, defaults to `2304`

The number of channels in the caption embeddings.

 2304
mlp_ratio `float`, defaults to `2.5`

The expansion ratio to use in the GLUMBConv layer.

 2.5
dropout `float`, defaults to `0.0`

The dropout probability.

 0.0
attention_bias `bool`, defaults to `False`

Whether to use bias in the attention layer.

 False
sample_size `int`, defaults to `32`

The base size of the input latent.

 30
patch_size `int`, defaults to `1`

The size of the patches to use in the patch embedding layer.

 (1, 2, 2)
norm_elementwise_affine `bool`, defaults to `False`

Whether to use elementwise affinity in the normalization layer.

 False
norm_eps `float`, defaults to `1e-6`

The epsilon value for the normalization layer.

 1e-06
qk_norm `str`, *optional*, defaults to `None`

The normalization to use for the query and key.

 'rms_norm_across_heads'

caption_norm instance-attribute 

caption_norm = SanaRMSNorm(
    inner_dim, eps=1e-05, elementwise_affine=True
)

caption_projection instance-attribute 

caption_projection = SanaPixArtAlphaTextProjection(
    in_features=caption_channels, hidden_size=inner_dim
)

config instance-attribute 

config = SanaVideoTransformerConfig(
    in_channels=in_channels,
    out_channels=out_channels,
    num_attention_heads=num_attention_heads,
    attention_head_dim=attention_head_dim,
    num_layers=num_layers,
    num_cross_attention_heads=num_cross_attention_heads,
    cross_attention_head_dim=cross_attention_head_dim,
    cross_attention_dim=cross_attention_dim,
    caption_channels=caption_channels,
    mlp_ratio=mlp_ratio,
    dropout=dropout,
    attention_bias=attention_bias,
    sample_size=sample_size,
    patch_size=tuple(patch_size),
    norm_elementwise_affine=norm_elementwise_affine,
    norm_eps=norm_eps,
    interpolation_scale=interpolation_scale,
    guidance_embeds=guidance_embeds,
    guidance_embeds_scale=guidance_embeds_scale,
    qk_norm=qk_norm,
    rope_max_seq_len=rope_max_seq_len,
)

dtype property 

dtype: dtype

norm_out instance-attribute 

norm_out = SanaModulatedNorm(
    inner_dim, elementwise_affine=False, eps=1e-06
)

patch_embedding instance-attribute 

patch_embedding = nn.Conv3d(
    in_channels,
    inner_dim,
    kernel_size=patch_size,
    stride=patch_size,
)

proj_out instance-attribute 

proj_out = nn.Linear(
    inner_dim, math.prod(patch_size) * out_channels
)

rope instance-attribute 

rope = WanRotaryPosEmbed(
    attention_head_dim, patch_size, rope_max_seq_len
)

scale_shift_table instance-attribute 

scale_shift_table = nn.Parameter(
    torch.randn(2, inner_dim) / inner_dim**0.5
)

time_embed instance-attribute 

time_embed = SanaCombinedTimestepGuidanceEmbeddings(
    inner_dim
)

transformer_blocks instance-attribute 

transformer_blocks = nn.ModuleList(
    [
        SanaVideoTransformerBlock(
            inner_dim,
            num_attention_heads,
            attention_head_dim,
            dropout=dropout,
            num_cross_attention_heads=num_cross_attention_heads,
            cross_attention_head_dim=cross_attention_head_dim,
            cross_attention_dim=cross_attention_dim,
            attention_bias=attention_bias,
            norm_elementwise_affine=norm_elementwise_affine,
            norm_eps=norm_eps,
            mlp_ratio=mlp_ratio,
            qk_norm=qk_norm,
            prefix=f"transformer_blocks.{layer_idx}",
        )
        for layer_idx in range(num_layers)
    ]
)

forward 

forward(
    hidden_states: Tensor,
    encoder_hidden_states: Tensor,
    timestep: Tensor,
    guidance: Tensor | None = None,
    encoder_attention_mask: Tensor | None = None,
    attention_mask: Tensor | None = None,
    controlnet_block_samples: tuple[Tensor] | None = None,
    return_dict: bool = True,
) -> tuple[Tensor, ...] | SanaVideoTransformerOutput

The [SanaVideoTransformer3DModel] forward method.

Parameters:

Name Type Description Default
hidden_states `torch.Tensor` of shape `(batch_size, in_channels, num_frames, height, width)`

Input hidden_states.

 required
encoder_hidden_states `torch.Tensor` of shape `(batch_size, sequence_len, embed_dims)`

Conditional embeddings (embeddings computed from the input conditions such as prompts) to use.

 required
timestep `torch.LongTensor`

Used to indicate denoising step.

 required
guidance `torch.Tensor`, *optional*

Guidance scale embedding.

 None
encoder_attention_mask `torch.Tensor`, *optional*

Cross-attention mask applied to encoder_hidden_states.

 None
attention_mask `torch.Tensor`, *optional*

Retained for Diffusers API compatibility. SANA linear self-attention does not consume this mask.

 None
controlnet_block_samples `tuple` of `torch.Tensor`, *optional*

A list of tensors that if specified are added to the residuals of transformer blocks.

 None
return_dict `bool`, *optional*, defaults to `True`

Whether to return SanaVideoTransformerOutput instead of a plain tuple.

 True

Returns:

Type Description
tuple[Tensor, ...] | SanaVideoTransformerOutput

If return_dict is True, a [SanaVideoTransformerOutput] is returned, otherwise a
tuple[Tensor, ...] | SanaVideoTransformerOutput

tuple where the first element is the sample tensor.

from_config classmethod 

from_config(
    config: dict | SanaVideoTransformerConfig,
) -> SanaVideoTransformer3DModel

load_weights 

load_weights(
    weights: Iterable[tuple[str, Tensor]],
) -> set[str]

get_sana_video_i2v_post_process_func 

get_sana_video_i2v_post_process_func(
    od_config: OmniDiffusionConfig,
)

get_sana_video_i2v_pre_process_func 

get_sana_video_i2v_pre_process_func(
    od_config: OmniDiffusionConfig,
)

get_sana_video_post_process_func 

get_sana_video_post_process_func(
    od_config: OmniDiffusionConfig,
)