Spaces:
Running on L40S
Running on L40S
File size: 4,803 Bytes
9f818c5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1
"""FLOPs estimation for the Wan 2.2 VAE encoder (Encoder3d)."""
from decimal import Decimal
def compute_wan_vae_encoder_flops(
B: int | Decimal,
T: int,
H: int,
W: int,
*,
dim: int = 160,
z_dim: int = 48,
dim_mult: list[int] | None = None,
num_res_blocks: int = 2,
temperal_downsample: list[bool] | None = None,
) -> Decimal:
"""Compute forward-pass FLOPs for the Wan 2.2 VAE encoder (Encoder3d).
The encoder converts a pixel-space video [B, 3, T, H, W] into a latent
[B, z_dim, T//4, H//16, W//16]. It is frozen during training so only
forward-pass FLOPs are counted (no backward).
The architecture: patchify(2) -> conv1 -> 4 downsample stages (each with
``num_res_blocks`` residual blocks + optional spatial/temporal downsample)
-> middle block (ResBlock + single-head spatial attention + ResBlock)
-> head (RMSNorm + SiLU + conv) -> pointwise 1x1 conv.
Args:
B: Batch size.
T: Number of pixel-space temporal frames.
H: Pixel-space height (must be divisible by 16).
W: Pixel-space width (must be divisible by 16).
dim: Base channel dimension of the encoder (default 160).
z_dim: Latent channel dimension (default 48, encoder outputs 2*z_dim).
dim_mult: Channel multiplier per stage (default [1, 2, 4, 4]).
num_res_blocks: Residual blocks per downsample stage (default 2).
temperal_downsample: Per-stage temporal downsampling flags (default
[False, True, True]).
Returns:
Total forward-pass FLOPs as a Decimal.
"""
if dim_mult is None:
dim_mult = [1, 2, 4, 4]
if temperal_downsample is None:
temperal_downsample = [False, True, True]
B = int(B)
flops = Decimal(0)
def _causalconv3d_flops(c_in: int, c_out: int, kt: int, kh: int, kw: int, bt: int, bh: int, bw: int) -> int:
return 2 * c_out * c_in * kt * kh * kw * B * bt * bh * bw
def _resblock_flops(in_dim: int, out_dim: int, bt: int, bh: int, bw: int) -> int:
vol = B * bt * bh * bw
f = 0
f += 5 * in_dim * vol # RMS_norm(in_dim)
f += 2 * out_dim * in_dim * 27 * vol # CausalConv3d(in_dim, out_dim, 3)
f += 5 * out_dim * vol # RMS_norm(out_dim)
f += 2 * out_dim * out_dim * 27 * vol # CausalConv3d(out_dim, out_dim, 3)
if in_dim != out_dim:
f += 2 * out_dim * in_dim * vol # shortcut CausalConv3d(in_dim, out_dim, 1)
return f
def _attnblock_flops(d: int, bt: int, bh: int, bw: int) -> int:
vol = B * bt * bh * bw
seq = bh * bw
f = 0
f += 5 * d * vol # RMS_norm
f += 2 * (d * 3) * d * vol # to_qkv Conv2d(d, 3d, 1)
f += 4 * B * bt * seq * seq * d # QK^T + Attn*V
f += 2 * d * d * vol # proj Conv2d(d, d, 1)
return f
# After patchify(patch_size=2): [B, 12, T, H/2, W/2]
t, h, w = T, H // 2, W // 2
# conv1: CausalConv3d(12, dims[0], 3)
dims = [dim * u for u in [1] + dim_mult] # [160, 160, 320, 640, 640]
flops += _causalconv3d_flops(12, dims[0], 3, 3, 3, t, h, w)
# Downsample stages
for i, (in_d, out_d) in enumerate(zip(dims[:-1], dims[1:])):
t_down = temperal_downsample[i] if i < len(temperal_downsample) else False
down_flag = i != len(dim_mult) - 1
cur_in = in_d
for _ in range(num_res_blocks):
flops += _resblock_flops(cur_in, out_d, t, h, w)
cur_in = out_d
if down_flag:
if t_down:
h_new, w_new = h // 2, w // 2
flops += 2 * out_d * out_d * 9 * B * t * h_new * w_new # spatial conv2d
t_new = t // 2
flops += 2 * out_d * out_d * 3 * B * t_new * h_new * w_new # temporal conv3d(3,1,1)
t, h, w = t_new, h_new, w_new
else:
h_new, w_new = h // 2, w // 2
flops += 2 * out_d * out_d * 9 * B * t * h_new * w_new
h, w = h_new, w_new
# Middle block: ResBlock + AttentionBlock + ResBlock
mid_dim = dims[-1]
flops += _resblock_flops(mid_dim, mid_dim, t, h, w)
flops += _attnblock_flops(mid_dim, t, h, w)
flops += _resblock_flops(mid_dim, mid_dim, t, h, w)
# Head: RMS_norm + SiLU + CausalConv3d(mid_dim, z_dim*2, 3)
enc_out_dim = z_dim * 2
flops += 5 * mid_dim * B * t * h * w # RMS_norm
flops += _causalconv3d_flops(mid_dim, enc_out_dim, 3, 3, 3, t, h, w)
# WanVAE_.conv1: CausalConv3d(z_dim*2, z_dim*2, 1) — pointwise 1x1
flops += _causalconv3d_flops(enc_out_dim, enc_out_dim, 1, 1, 1, t, h, w)
return Decimal(flops)
|