Flash-VAED

Independent reproduction of Flash-VAED: Plug-and-Play VAE Decoders for Efficient Video Generation for the diffusers Wan 2.1 VAE.

• Paper: arXiv:2602.19161
• Official work-in-progress repo: Aoko955/Flash-VAED
• Detailed experiment: OpenVidHD 480p report

This is an independent implementation created by the repository owner with OpenAI Codex. It is not an official release from the paper authors.

Results

Trained on 9,800 OpenVidHD reconstruction pairs at 17×480×832; evaluated on 200 held-out clips. Decoder speed uses one H100, bf16, batch 1, and 81 output frames at 480×832.

Decoder performance

Validation quality

One-quarter retention is the recommended tradeoff: 3.89× faster decoder and 57% lower peak memory with modest fine-detail smoothing. One-eighth retention is intended for previews or throughput-first workloads. Phase 2 is the quality-first option.

Decode examples

Held-out OpenVidHD middle frames. Columns are Source | Wan | Phase 2 | 1/4 retained | 1/8 retained; open an image to inspect fine detail.

These are decoder-only reconstruction results. Actual Wan denoised latents and end-to-end generation still need prompt-diverse validation before production deployment.

Install

pip install -e .

This installs all inference, training, video, evaluation, and development dependencies. Requires Python 3.10+ and a diffusers-format Wan 2.1 VAE.

Reproduce the 10K 480p dataset

The dataset itself is not bundled. It was built reproducibly from the OpenVidHD archives in nkp37/OpenVid-1M:

CUDA_VISIBLE_DEVICES=0 python scripts/build_openvid_pairs.py \
  --out data/openvidhd_480p_10k \
  --archive-dir data/openvidhd_archives \
  --count 10000 --val-count 200 \
  --frames 17 --fps 8 --height 480 --width 832 \
  --dtype bf16

For each unique source video, the builder deterministically selects one 17-frame clip at 8 FPS using seed 20260729, center-crops/resizes it to 480×832 without upscaling, and rejects clips that are too short, too small, low-contrast (std < 0.06), or low-motion (motion < 0.008). The Wan 2.1 VAE encoder then produces a (16, 5, 60, 104) latent. The first 9,800 accepted clips form the training split and the final 200 form a held-out validation split.

Each .pt pair contains the fp16 video and latent plus the source archive, member filename, and clip start time. manifest.jsonl records the same provenance and quality measurements, so construction is resumable and auditable. The completed run accepted 10,000 clips and rejected 1,069 candidates. OpenVid-1M is CC-BY-4.0; users must also follow the licenses of its constituent sources.

Train

Pairs are .pt files containing:

{
    "latent": Tensor[16, T, H // 8, W // 8],
    "video": Tensor[3, 1 + 4 * (T - 1), H, W],  # range [-1, 1]
}

Build reconstruction pairs:

python scripts/prepare_pairs.py \
  --mode reconstruct --videos data/videos --out data/pairs

Three-phase training:

torchrun --standalone --nproc_per_node=4 scripts/train.py \
  --data data/pairs \
  --out runs/wan21 \
  --prune-ratio 0.25 \
  --steps 10000 5000 15000 \
  --latent-frames 5 \
  --dtype bf16 \
  --ckpt-every 500 \
  --preview-pair data/val/example.pt

The phases are:

1. Replace dominant 3D operators and distill deep features.
2. Select and strengthen retained channels.
3. Physically prune up2/up3 and recover output quality.

Training supports multi-GPU gradient averaging, distributed sampling, bf16, streaming activation checkpointing, 500-step checkpoints, and fixed target | Wan | Flash-VAED previews.

Evaluate and benchmark

python scripts/evaluate.py \
  --flash-ckpt runs/wan21/phase3_final.pt \
  --data data/val --latent-frames 5

python scripts/evaluate.py \
  --flash-ckpt runs/wan21/phase3_final.pt \
  --data data/val --latent-frames 5 --reference

python scripts/benchmark.py \
  --flash-ckpt runs/wan21/phase3_final.pt \
  --resolution 480 832 --frames 81 --dtype bfloat16

Evaluation reports PSNR, SSIM, and LPIPS against either paired source videos or the original Wan decoder. Benchmarking reports latency, FPS, peak CUDA memory, parameter count, and speedup.

Diffusers usage

Load a training checkpoint:

import torch
from diffusers import AutoencoderKLWan
from flash_vaed import load_flash_vaed

vae = AutoencoderKLWan.from_pretrained(
    "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
    subfolder="vae",
    torch_dtype=torch.bfloat16,
)
flash = load_flash_vaed(vae, "runs/wan21/phase3_final.pt")
video = flash.decode(latents)

Export and load a standalone diffusers component:

from flash_vaed import FlashAutoencoderKLWan, save_pretrained_flash_vaed

save_pretrained_flash_vaed(flash, "runs/wan21/diffusers_vae")
vae = FlashAutoencoderKLWan.from_pretrained(
    "runs/wan21/diffusers_vae",
    torch_dtype=torch.bfloat16,
)
pipe.vae = vae
video = vae.decode(latents).sample

The export uses standard config.json and diffusion_pytorch_model.safetensors files and does not require the original Wan checkpoint when loaded.

Design

The Wan preset follows the paper:

• mid, up0, up1: depthwise-separable causal 3D convolutions
• up2, up3: spatial 2D convolutions
• configurable retained-channel ratio in up2/up3
• fitted reconstruction shortcuts and feature projections
• fused RMSNorm-SiLU CUDA path with PyTorch fallback
• differentiable streaming decode with correct Wan temporal caches

Main modules:

flash_vaed/
  ops.py           replacement operators
  surgery.py       operator replacement and channel slicing
  calibrate.py     feature collection and channel selection
  prune.py         pruning solution and gradient masks
  distill.py       three-phase trainer
  runtime.py       differentiable streaming decode
  metrics.py       PSNR, SSIM, LPIPS
  bench.py         latency, memory, parameter benchmarks
  diffusers_model.py  standalone diffusers VAE

Tests

python -m pytest -q

The tests cover operator initialization, channel selection, pruning topology, gradient masks, projection initialization, checkpoint round-trips, dtype handling, native diffusers loading, and an end-to-end three-phase run.

License

Code is released under the MIT License. Model weights, datasets, papers, and third-party dependencies remain subject to their own licenses.