wanvae

Compatibility imports for the family-local Wan VAE.

Keep aliases, including the cache context variables, identical to the canonical objects. New Wan code belongs in fastvideo.models.wan.vae.

fastvideo.models.vaes.wanvae.AutoencoderKLWan

AutoencoderKLWan(config: WanVAEConfig)

A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos. Introduced in [Wan 2.1].

fastvideo.models.vaes.wanvae.AutoencoderKLWan.clear_cache

clear_cache() -> None

fastvideo.models.vaes.wanvae.AutoencoderKLWan.optimize_memory_format

optimize_memory_format() -> None

Optional CUDA-only weight layout optimization used by MG3 LightVAE.

fastvideo.models.vaes.wanvae.AutoencoderKLWan.encode

encode(x: torch.Tensor) -> torch.Tensor

fastvideo.models.vaes.wanvae.AutoencoderKLWan.tiled_encode

tiled_encode(x: torch.Tensor) -> torch.Tensor

fastvideo.models.vaes.wanvae.AutoencoderKLWan.spatial_tiled_encode

spatial_tiled_encode(x: torch.Tensor) -> torch.Tensor

fastvideo.models.vaes.wanvae.AutoencoderKLWan.decode

decode(z: torch.Tensor) -> torch.Tensor

fastvideo.models.vaes.wanvae.AutoencoderKLWan.tiled_decode

tiled_decode(z: torch.Tensor) -> torch.Tensor

fastvideo.models.vaes.wanvae.AutoencoderKLWan.spatial_tiled_decode

spatial_tiled_decode(z: torch.Tensor) -> torch.Tensor

fastvideo.models.vaes.wanvae.AutoencoderKLWan.parallel_tiled_decode

parallel_tiled_decode(z: torch.FloatTensor) -> torch.FloatTensor

fastvideo.models.vaes.wanvae.AutoencoderKLWan.get_streaming_cache

get_streaming_cache() -> list[torch.Tensor | None]

fastvideo.models.vaes.wanvae.AutoencoderKLWan.streaming_decode

streaming_decode(z: torch.Tensor, cache: list[torch.Tensor | None], is_first_chunk: bool = False) -> tuple[torch.Tensor, list[torch.Tensor | None]]

Parameters

z: `torch.Tensor`

Latent tensor of shape [B, C, T, H, W].

cache: `list[torch.Tensor | None]`

The VAE cache.

is_first_chunk: `bool`

Whether this is the first chunk in the sequence.

Returns

tuple[torch.Tensor, list[torch.Tensor | None]]

A tuple of (decoded_frames, updated_cache).

fastvideo.models.vaes.wanvae.AutoencoderKLWan.forward

forward(sample: torch.Tensor, sample_posterior: bool = False, generator: torch.Generator | None = None) -> torch.Tensor

Parameters

sample: `torch.Tensor`

Input sample.

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

Whether or not to return a [DecoderOutput] instead of a plain tuple.

fastvideo.models.vaes.wanvae.AvgDown3D

AvgDown3D(in_channels, out_channels, factor_t, factor_s = 1)

fastvideo.models.vaes.wanvae.AvgDown3D.forward

forward(x: torch.Tensor) -> torch.Tensor

fastvideo.models.vaes.wanvae.DiagonalGaussianDistribution

DiagonalGaussianDistribution(parameters: torch.Tensor, deterministic: bool = False)

fastvideo.models.vaes.wanvae.DiagonalGaussianDistribution.sample

sample(generator: torch.Generator | None = None) -> torch.Tensor

fastvideo.models.vaes.wanvae.DiagonalGaussianDistribution.kl

kl(other: Optional[DiagonalGaussianDistribution] = None) -> torch.Tensor

fastvideo.models.vaes.wanvae.DiagonalGaussianDistribution.nll

nll(sample: torch.Tensor, dims: tuple[int, ...] = (1, 2, 3)) -> torch.Tensor

fastvideo.models.vaes.wanvae.DiagonalGaussianDistribution.mode

mode() -> torch.Tensor

fastvideo.models.vaes.wanvae.DupUp3D

DupUp3D(in_channels: int, out_channels: int, factor_t, factor_s = 1)

fastvideo.models.vaes.wanvae.DupUp3D.forward

forward(x: torch.Tensor) -> torch.Tensor

fastvideo.models.vaes.wanvae.ParallelTiledVAE

ParallelTiledVAE(config: VAEConfig, **kwargs)

fastvideo.models.vaes.wanvae.ParallelTiledVAE.to

to(device) -> ParallelTiledVAE

fastvideo.models.vaes.wanvae.ParallelTiledVAE.encode

encode(x: torch.Tensor) -> torch.Tensor

fastvideo.models.vaes.wanvae.ParallelTiledVAE.decode

decode(z: torch.Tensor) -> torch.Tensor

fastvideo.models.vaes.wanvae.ParallelTiledVAE.blend_v

blend_v(a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor

fastvideo.models.vaes.wanvae.ParallelTiledVAE.blend_h

blend_h(a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor

fastvideo.models.vaes.wanvae.ParallelTiledVAE.blend_t

blend_t(a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor

fastvideo.models.vaes.wanvae.ParallelTiledVAE.spatial_tiled_encode

spatial_tiled_encode(x: torch.Tensor) -> torch.Tensor

Encode a batch of images using a tiled encoder.

Parameters

x: `torch.Tensor`

Input batch of videos.

Returns

torch.Tensor

torch.Tensor: The latent representation of the encoded videos.

fastvideo.models.vaes.wanvae.ParallelTiledVAE.parallel_tiled_decode

parallel_tiled_decode(z: torch.FloatTensor) -> torch.FloatTensor

Parallel version of tiled_decode that distributes both temporal and spatial computation across GPUs

fastvideo.models.vaes.wanvae.ParallelTiledVAE.spatial_tiled_decode

spatial_tiled_decode(z: torch.Tensor) -> torch.Tensor

Decode a batch of images using a tiled decoder.

Parameters

z: `torch.Tensor`

Input batch of latent vectors.

Returns

torch.Tensor

torch.Tensor: The decoded images.

fastvideo.models.vaes.wanvae.ParallelTiledVAE.tiled_encode

tiled_encode(x: torch.Tensor) -> torch.Tensor

fastvideo.models.vaes.wanvae.ParallelTiledVAE.tiled_decode

tiled_decode(z: torch.Tensor) -> torch.Tensor

fastvideo.models.vaes.wanvae.ParallelTiledVAE.enable_tiling

enable_tiling(tile_sample_min_height: int | None = None, tile_sample_min_width: int | None = None, tile_sample_min_num_frames: int | None = None, tile_sample_stride_height: int | None = None, tile_sample_stride_width: int | None = None, tile_sample_stride_num_frames: int | None = None, blend_num_frames: int | None = None, use_tiling: bool | None = None, use_temporal_tiling: bool | None = None, use_parallel_tiling: bool | None = None) -> None

Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow processing larger images.

Parameters

tile_sample_min_height: `int`, *optional*

The minimum height required for a sample to be separated into tiles across the height dimension.

tile_sample_min_width: `int`, *optional*

The minimum width required for a sample to be separated into tiles across the width dimension.

tile_sample_min_num_frames: `int`, *optional*

The minimum number of frames required for a sample to be separated into tiles across the frame dimension.

tile_sample_stride_height: `int`, *optional*

The minimum amount of overlap between two consecutive vertical tiles. This is to ensure that there are no tiling artifacts produced across the height dimension.

tile_sample_stride_width: `int`, *optional*

The stride between two consecutive horizontal tiles. This is to ensure that there are no tiling artifacts produced across the width dimension.

tile_sample_stride_num_frames: `int`, *optional*

The stride between two consecutive frame tiles. This is to ensure that there are no tiling artifacts produced across the frame dimension.

fastvideo.models.vaes.wanvae.ParallelTiledVAE.disable_tiling

disable_tiling() -> None

Disable tiled VAE decoding. If enable_tiling was previously enabled, this method will go back to computing decoding in one step.

fastvideo.models.vaes.wanvae.WanAttentionBlock

WanAttentionBlock(dim)

Causal self-attention with a single head.

Parameters

dim: int

The number of channels in the input tensor.

fastvideo.models.vaes.wanvae.WanAttentionBlock.forward

forward(x)

fastvideo.models.vaes.wanvae.WanCausalConv3d

WanCausalConv3d(in_channels: int, out_channels: int, kernel_size: int | tuple[int, int, int], stride: int | tuple[int, int, int] = 1, padding: int | tuple[int, int, int] = 0)

A custom 3D causal convolution layer with feature caching support.

This layer extends the standard Conv3D layer by ensuring causality in the time dimension and handling feature caching for efficient inference.

Parameters

in_channels: int

Number of channels in the input image

out_channels: int

Number of channels produced by the convolution

kernel_size: int or tuple

Size of the convolving kernel

stride: int or tuple

Stride of the convolution. Default: 1

padding: int or tuple

Zero-padding added to all three sides of the input. Default: 0

fastvideo.models.vaes.wanvae.WanCausalConv3d.forward

forward(x, cache_x = None)

fastvideo.models.vaes.wanvae.WanDecoder3d

WanDecoder3d(dim = 128, z_dim = 4, dim_mult = (1, 2, 4, 4), num_res_blocks = 2, attn_scales = (), temperal_upsample = (False, True, True), dropout = 0.0, non_linearity: str = 'silu', out_channels: int = 3, is_residual: bool = False)

A 3D decoder module.

Parameters

dim: int

The base number of channels in the first layer.

z_dim: int

The dimensionality of the latent space.

dim_mult: list of int

Multipliers for the number of channels in each block.

num_res_blocks: int

Number of residual blocks in each block.

attn_scales: list of float

Scales at which to apply attention mechanisms.

temperal_upsample: list of bool

Whether to upsample temporally in each block.

dropout: float

Dropout rate for the dropout layers.

non_linearity: str

Type of non-linearity to use.

fastvideo.models.vaes.wanvae.WanDecoder3d.forward

forward(x)

fastvideo.models.vaes.wanvae.WanEncoder3d

WanEncoder3d(in_channels: int = 3, dim = 128, z_dim = 4, dim_mult = (1, 2, 4, 4), num_res_blocks = 2, attn_scales = (), temperal_downsample = (True, True, False), dropout = 0.0, non_linearity: str = 'silu', is_residual: bool = False)

A 3D encoder module.

Parameters

dim: int

The base number of channels in the first layer.

z_dim: int

The dimensionality of the latent space.

dim_mult: list of int

Multipliers for the number of channels in each block.

num_res_blocks: int

Number of residual blocks in each block.

attn_scales: list of float

Scales at which to apply attention mechanisms.

temperal_downsample: list of bool

Whether to downsample temporally in each block.

dropout: float

Dropout rate for the dropout layers.

non_linearity: str

Type of non-linearity to use.

fastvideo.models.vaes.wanvae.WanEncoder3d.forward

forward(x)

fastvideo.models.vaes.wanvae.WanMidBlock

WanMidBlock(dim: int, dropout: float = 0.0, non_linearity: str = 'silu', num_layers: int = 1)

Middle block for WanVAE encoder and decoder.

Parameters

dim: int

Number of input/output channels.

dropout: float

Dropout rate.

non_linearity: str

Type of non-linearity to use.

fastvideo.models.vaes.wanvae.WanMidBlock.forward

forward(x)

fastvideo.models.vaes.wanvae.WanRMS_norm

WanRMS_norm(dim: int, channel_first: bool = True, images: bool = True, bias: bool = False)

A custom RMS normalization layer.

Parameters

dim: int

The number of dimensions to normalize over.

channel_first: bool

Whether the input tensor has channels as the first dimension. Default is True.

images: bool

Whether the input represents image data. Default is True.

bias: bool

Whether to include a learnable bias term. Default is False.

fastvideo.models.vaes.wanvae.WanRMS_norm.forward

forward(x)

fastvideo.models.vaes.wanvae.WanResample

WanResample(dim: int, mode: str, upsample_out_dim: int = None)

A custom resampling module for 2D and 3D data.

Parameters

dim: int

The number of input/output channels.

mode: str

The resampling mode. Must be one of: - 'none': No resampling (identity operation). - 'upsample2d': 2D upsampling with nearest-exact interpolation and convolution. - 'upsample3d': 3D upsampling with nearest-exact interpolation, convolution, and causal 3D convolution. - 'downsample2d': 2D downsampling with zero-padding and convolution. - 'downsample3d': 3D downsampling with zero-padding, convolution, and causal 3D convolution.

fastvideo.models.vaes.wanvae.WanResample.forward

forward(x)

fastvideo.models.vaes.wanvae.WanResidualBlock

WanResidualBlock(in_dim: int, out_dim: int, dropout: float = 0.0, non_linearity: str = 'silu')

A custom residual block module.

Parameters

in_dim: int

Number of input channels.

out_dim: int

Number of output channels.

dropout: float

Dropout rate for the dropout layer. Default is 0.0.

non_linearity: str

Type of non-linearity to use. Default is "silu".

fastvideo.models.vaes.wanvae.WanResidualBlock.forward

forward(x)

fastvideo.models.vaes.wanvae.WanResidualDownBlock

WanResidualDownBlock(in_dim, out_dim, dropout, num_res_blocks, temperal_downsample = False, down_flag = False)

fastvideo.models.vaes.wanvae.WanResidualDownBlock.forward

forward(x)

fastvideo.models.vaes.wanvae.WanResidualUpBlock

WanResidualUpBlock(in_dim: int, out_dim: int, num_res_blocks: int, dropout: float = 0.0, temperal_upsample: bool = False, up_flag: bool = False, non_linearity: str = 'silu')

A block that handles upsampling for the WanVAE decoder. Args: in_dim (int): Input dimension out_dim (int): Output dimension num_res_blocks (int): Number of residual blocks dropout (float): Dropout rate temperal_upsample (bool): Whether to upsample on temporal dimension up_flag (bool): Whether to upsample or not non_linearity (str): Type of non-linearity to use

fastvideo.models.vaes.wanvae.WanResidualUpBlock.forward

forward(x)

Forward pass through the upsampling block. Args: x (torch.Tensor): Input tensor feat_cache (list, optional): Feature cache for causal convolutions feat_idx (list, optional): Feature index for cache management Returns: torch.Tensor: Output tensor

fastvideo.models.vaes.wanvae.WanUpBlock

WanUpBlock(in_dim: int, out_dim: int, num_res_blocks: int, dropout: float = 0.0, upsample_mode: str | None = None, non_linearity: str = 'silu')

A block that handles upsampling for the WanVAE decoder.

Parameters

in_dim: int

Input dimension

out_dim: int

Output dimension

num_res_blocks: int

Number of residual blocks

dropout: float

Dropout rate

upsample_mode: str

Mode for upsampling ('upsample2d' or 'upsample3d')

non_linearity: str

Type of non-linearity to use

fastvideo.models.vaes.wanvae.WanUpBlock.forward

forward(x)

Forward pass through the upsampling block.

Parameters

x: torch.Tensor

Input tensor

feat_cache: list

Feature cache for causal convolutions

feat_idx: list

Feature index for cache management

Returns

torch.Tensor: Output tensor

fastvideo.models.vaes.wanvae.WanUpsample

WanUpsample()

Perform upsampling while ensuring the output tensor has the same data type as the input.

Parameters

x: torch.Tensor

Input tensor to be upsampled.

Returns

torch.Tensor: Upsampled tensor with the same data type as the input.

fastvideo.models.vaes.wanvae.WanUpsample.forward

forward(x)

fastvideo.models.vaes.wanvae.WanVAEConfig

WanVAEConfig(arch_config: WanVAEArchConfig = WanVAEArchConfig(), load_encoder: bool = True, load_decoder: bool = True, tile_sample_min_height: int = 256, tile_sample_min_width: int = 256, tile_sample_min_num_frames: int = 16, tile_sample_stride_height: int = 192, tile_sample_stride_width: int = 192, tile_sample_stride_num_frames: int = 12, blend_num_frames: int = 0, use_tiling: bool = False, use_temporal_tiling: bool = False, use_parallel_tiling: bool = False, use_temporal_scaling_frames: bool = True, use_feature_cache: bool = True, use_light_vae: bool = False, *, _resolved_attention_backend: AttentionBackendEnum | None = None)

fastvideo.models.vaes.wanvae.forward_context

forward_context(first_frame_arg = False, feat_cache_arg = None, feat_idx_arg = None, first_chunk_arg = None, use_light_vae_arg = False)

fastvideo.models.vaes.wanvae.patchify

patchify(x, patch_size)

fastvideo.models.vaes.wanvae.unpatchify

unpatchify(x, patch_size)