William Lin
86ef6bfadf
- fastvideo-kernel/wan22_vae_decode: the Wan2.2 (TI2V-5B) VAE decoder on sm_100a / sm_103a at 704x1248, built on first use against fastvideo-kernel's CUTLASS submodule. It frees intermediates early and carries 4 latent frames per pass: a 5 s video peaks at 20.3 GiB, down from 94.6 GiB in the kernel's first version, with the same output. - WanVAEConfig.decode_kernel routes AutoencoderKLWan.decode through it where it applies (bf16 with fp32 accumulation, about 7x faster than the fp32 decode on a B300, PSNR 61-63 dB, SSIM 0.998 against it); decode_kernel_parallel splits the decode over the sequence-parallel group (2, 4 or 8 GPUs) and broadcasts the video, with the ranks agreeing on whether every one of them could build the kernel. - PAN2's DiT runs fused Triton kernels for Q/K RMSNorm + RoPE and LayerNorm + modulation (pan2_fusions, set_fused_kernels). - PAN2PipelineConfig turns all three on (vae_decode_kernel, vae_decode_kernel_parallel, dit_fused_kernels). - VideoGenerator copies a decoded RGB video's frames to the host once instead of once per frame. - The PAN2 presets sample at guidance 2.5, the reference launcher's default since 2026-10-23.
92 lines
2.8 KiB
Python
92 lines
2.8 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Architecture and runtime configuration for the Wan VAE."""
|
|
|
|
from dataclasses import dataclass, field
|
|
|
|
import torch
|
|
|
|
from fastvideo.configs.models.vaes.base import VAEArchConfig, VAEConfig
|
|
|
|
|
|
@dataclass
|
|
class WanVAEArchConfig(VAEArchConfig):
|
|
base_dim: int = 96
|
|
decoder_base_dim: int | None = None
|
|
z_dim: int = 16
|
|
dim_mult: tuple[int, ...] = (1, 2, 4, 4)
|
|
num_res_blocks: int = 2
|
|
attn_scales: tuple[float, ...] = ()
|
|
temperal_downsample: tuple[bool, ...] = (False, True, True)
|
|
dropout: float = 0.0
|
|
latents_mean: tuple[float, ...] = (
|
|
-0.7571,
|
|
-0.7089,
|
|
-0.9113,
|
|
0.1075,
|
|
-0.1745,
|
|
0.9653,
|
|
-0.1517,
|
|
1.5508,
|
|
0.4134,
|
|
-0.0715,
|
|
0.5517,
|
|
-0.3632,
|
|
-0.1922,
|
|
-0.9497,
|
|
0.2503,
|
|
-0.2921,
|
|
)
|
|
latents_std: tuple[float, ...] = (
|
|
2.8184,
|
|
1.4541,
|
|
2.3275,
|
|
2.6558,
|
|
1.2196,
|
|
1.7708,
|
|
2.6052,
|
|
2.0743,
|
|
3.2687,
|
|
2.1526,
|
|
2.8652,
|
|
1.5579,
|
|
1.6382,
|
|
1.1253,
|
|
2.8251,
|
|
1.9160,
|
|
)
|
|
is_residual: bool = False
|
|
in_channels: int = 3
|
|
out_channels: int = 3
|
|
patch_size: int | None = None
|
|
scale_factor_temporal: int = 4
|
|
scale_factor_spatial: int = 8
|
|
clip_output: bool = True
|
|
|
|
def __post_init__(self):
|
|
self.scaling_factor: torch.tensor = 1.0 / torch.tensor(self.latents_std).view(1, self.z_dim, 1, 1, 1)
|
|
self.shift_factor: torch.tensor = torch.tensor(self.latents_mean).view(1, self.z_dim, 1, 1, 1)
|
|
self.temporal_compression_ratio = self.scale_factor_temporal
|
|
self.spatial_compression_ratio = self.scale_factor_spatial
|
|
|
|
|
|
@dataclass
|
|
class WanVAEConfig(VAEConfig):
|
|
arch_config: WanVAEArchConfig = field(default_factory=WanVAEArchConfig)
|
|
use_feature_cache: bool = True
|
|
|
|
use_tiling: bool = False
|
|
use_temporal_tiling: bool = False
|
|
use_parallel_tiling: bool = False
|
|
|
|
use_light_vae: bool = False
|
|
# Decode with fastvideo-kernel's Wan2.2 decoder (wan22_vae_decode: B200 / GB200 / B300, the Wan2.2-TI2V-5B VAE,
|
|
# 704 x 1248 video, batch 1) in bf16 where it applies; other inputs and devices take the decode here.
|
|
decode_kernel: bool = False
|
|
# With decode_kernel, split each decode over the sequence-parallel group (2, 4 or 8 GPUs) instead of decoding on
|
|
# every GPU alone. Every rank of the group must decode the same latents at the same time, as a pipeline's decoding
|
|
# stage does; the group's first rank assembles the video and broadcasts it, so every rank returns it as before.
|
|
decode_kernel_parallel: bool = False
|
|
|
|
def __post_init__(self):
|
|
self.blend_num_frames = (self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames) * 2
|