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)