diffusers/pr_14539/en/api/models/autoencoder_tiny_video.md

Tiny Video AutoEncoder

Tiny AutoEncoder for Hunyuan Video (TAEHV) was introduced in madebyollin/taehv by Ollin Boer Bohan. It is a family of tiny causal video autoencoders distilled from full video VAEs — taew2_2 decodes the Wan 2.2 latent space of AutoencoderKLWan about 50× faster than the full model — for previews and real-time decoding. Latents are the normalized (roughly unit Gaussian) latents of the full VAE.

Decode a video chunk by chunk with a TinyVideoDecodeCache: each call decodes only the new latent frames, continuing from the previous calls, and the result is identical to a single decode of all frames.

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")

cache = TinyVideoDecodeCache()
for latents in latent_chunks:  # [B, 48, T, h, w], normalized Wan 2.2 latents
    frames = vae.decode(latents, cache=cache).sample  # [B, 3, 4 * T, 16 * h, 16 * w] in [-1, 1]

AutoencoderTinyVideo[[diffusers.AutoencoderTinyVideo]]

diffusers.AutoencoderTinyVideo[[diffusers.AutoencoderTinyVideo]]

diffusers.AutoencoderTinyVideo(in_channels: int = 3, latent_channels: int = 48, patch_size: int = 2, encoder_block_out_channels: tuple = (64, 64, 64), decoder_block_out_channels: tuple = (256, 128, 64, 64), encoder_time_downscale: tuple = (True, True, False), decoder_time_upscale: tuple = (False, True, True))

Source

Parameters:

in_channels (int, defaults to 3) : Number of channels in the input video.

latent_channels (int, defaults to 48) : Number of channels in the latent space.

patch_size (int, defaults to 2) : Pixel-(un)shuffle factor applied to the frames before the encoder and after the decoder.

encoder_block_out_channels (tuple[int, ...], defaults to (64, 64, 64)) : Output channels of the encoder blocks; each block halves the spatial size.

decoder_block_out_channels (tuple[int, ...], defaults to (256, 128, 64, 64)) : Channels of the decoder: the first entry is the width after the input conv, each following block doubles the spatial size and outputs the next entry.

encoder_time_downscale (tuple[bool, ...], defaults to (True, True, False)) : Whether each encoder block halves the number of frames.

decoder_time_upscale (tuple[bool, ...], defaults to (False, True, True)) : Whether each decoder block doubles the number of frames.

A tiny causal video autoencoder (TAEHV, madebyollin/taehv) that encodes to and decodes from the latent space of a full video VAE — e.g. taew2_2 for the Wan 2.2 VAE — orders of magnitude faster than the full model, for previews and real-time decoding. Latents are the normalized (roughly unit Gaussian) latents of the full VAE.

This model inherits from ModelMixin. Check the superclass documentation for its generic methods implemented for all models (such as downloading or saving).

decode[[diffusers.AutoencoderTinyVideo.decode]]

decode(z: Tensor, return_dict: bool = True, cache: diffusers.models.autoencoders.autoencoder_tiny_video.TinyVideoDecodeCache | None = None)

Source

Parameters:

z (torch.Tensor) : Input batch of latent vectors.

return_dict (bool, optional, defaults to True) : Whether to return a ~models.vae.DecoderOutput instead of a plain tuple.

cache (TinyVideoDecodeCache, optional) : Decode a video chunk by chunk: pass the same cache to consecutive calls and each call decodes only the frames in z, continuing from the previous calls.

Decode a batch of latents [B, C, T, h, w] to videos in [-1, 1]. T latent frames decode to T * temporal_upsampling_ratio - (temporal_upsampling_ratio - 1) frames: the first frames produced by the decoder are warm-up frames and are dropped.

encode[[diffusers.AutoencoderTinyVideo.encode]]

encode(x: Tensor, return_dict: bool = True)

Source

Encode a batch of videos [B, C, T, H, W] in [-1, 1]. The frames are padded at the end, by repeating the last one, to a multiple of temporal_compression_ratio.

forward[[diffusers.AutoencoderTinyVideo.forward]]

forward(sample: Tensor, return_dict: bool = True)

Source

Parameters:

sample (torch.Tensor) : Input video [B, C, T, H, W] in [-1, 1].

return_dict (bool, optional, defaults to True) : Whether or not to return a DecoderOutput instead of a plain tuple.

TinyVideoDecodeCache[[diffusers.models.autoencoders.autoencoder_tiny_video.TinyVideoDecodeCache]]

diffusers.models.autoencoders.autoencoder_tiny_video.TinyVideoDecodeCache[[diffusers.models.autoencoders.autoencoder_tiny_video.TinyVideoDecodeCache]]

diffusers.models.autoencoders.autoencoder_tiny_video.TinyVideoDecodeCache()

Source

Per-layer memory for decoding a video chunk by chunk with AutoencoderTinyVideo.decode().

Pass the same cache to consecutive decode(z, cache=cache) calls: each call decodes only the latent frames in z, continuing from the frames decoded by the previous calls, and the concatenated result is identical to decoding all frames in a single call. Create a new cache for every new video.