taew2_2 — Tiny Video AutoEncoder for Wan 2.2 (Diffusers)

taew2_2 by Ollin Boer Bohan, converted to AutoencoderTinyVideo (scripts/convert_taehv_to_diffusers.py). A tiny causal video autoencoder (11.4M parameters) for the latent space of the Wan 2.2 VAE (AutoencoderKLWan, 48 channels, 16× spatial / 4× temporal): previews and real-time decoding, roughly 50× faster than the full VAE. It works on the normalized (roughly unit Gaussian) Wan 2.2 latents.

import torch
from diffusers import AutoencoderTinyVideo
from diffusers.models.autoencoders.autoencoder_tiny_video import TinyVideoDecodeCache

vae = AutoencoderTinyVideo.from_pretrained("YiYiXu/taew2_2-diffusers", dtype=torch.bfloat16).to("cuda")

# one shot: [B, 48, T, h, w] normalized latents -> [B, 3, 4 * T - 3, 16 * h, 16 * w] frames in [-1, 1]
frames = vae.decode(latents).sample

# chunk by chunk (streaming / world models): identical result, only the new latents are decoded per call
cache = TinyVideoDecodeCache()
for chunk in latent_chunks:
    frames = vae.decode(chunk, cache=cache).sample

Weights: the taew2_2.pth release of madebyollin/taehv (MIT).