AsymmetricMagVitV2

Lightweight open-source reproduction of MagVitV2, fully aligned with the paper’s functionality. Supports image and video joint encoding and decoding, as well as videos of arbitrary length and resolution.

All spatio-temporal operators are implemented using causal 3D to avoid video instability caused by 2D+1D, ensures that the FVD does not sudden increases.
The Encoder and Decoder support arbitrary resolutions, support auto-regressive inference for arbitrary durations.
Training employs multi-resolution and dynamic-duration mixed training, allowing decoding of videos with arbitrary odd frames as long as GPU memory permits, demonstrating temporal extrapolation capability.
The model is closely aligned with MagVitV2 but with reduced parameter, particularly in the lightweight Encoder, reducing the burden of caching VAE features.

Demo

16 channel VAE image reconstruction

Converting MP4 to GIF may result in detail loss, pixelation, and incomplete duration. It is recommended to watch the original video for the best experience.

60s 3840x2160

bilibili_Black Myth：Wu KongULR 16zVAE

60s 1920x1080

bilibili_tokyo_walk ULR 16zVAE

image reconstruction

Installation
Model Weights
Metric
Inference
TODO List
Contact Us
Reference

Installation

1. Clone the repo

git clonehttps://github.com/bornfly-detachment/AsymmetricMagVitV2.git
cd AsymmetricMagVitV2

2. Setting up the virtualenv

This is assuming you have navigated to the AsymmetricMagVitV2 root after cloning it.

# install required packages from pypi
python3 -m venv .pt2
source .pt2/bin/activate
pip3 install -r requirements/pt2.txt

Model Weights

model	downsample (THW)	Encoder Size	Decoder Size
svd 2Dvae	1x8x8	34M	64M
AsymmetricMagVitV2	4x8x8	100M	159M

model	Data	#iterations	URL
AsymmetricMagVitV2_4z	20M Intervid	2node 1200k	AsymmetricMagVitV2_4z
AsymmetricMagVitV2_16z	20M Intervid	4node 860k	AsymmetricMagVitV2_16z

Metric

model	temporal-frame	fvd	fid	psnr	ssim
SVD VAE	1	190.6	1.8	28.2	1.0
openSoraPlan	1	249.8	1.04	29.6	0.99
openSoraPlan	17	725.4	3.17	23.4	0.89
openSoraPlan	33	756.8	3.5	23	0.89
AsymmetricMagVitV2_16z	1	106.7	0.2	36.3	1.0
AsymmetricMagVitV2_16z	17	131.4	0.8	30.7	1.0
AsymmetricMagVitV2_16z	33	208.2	1.4	28.9	1.0

Note:

The test video is the original scale of data/videos/tokyo_walk.mp4. Previously, preprocessing with resize+CenterCrop256 resolution was also tested on a larger test set, and the results showed consistent trends. Now, it has been found that high-resolution and original-sized videos pose the most challenging task for 3DVAE. Therefore, only this one video was tested, configured at 8fps, and evaluated for the first 10 seconds.
The evaluation code can be referenced in models/evaluation.py. However, it has been a while since I last ran it, and there have been modifications to the inference code. Calculating FID and FVD scores depends on the model, original image preprocessing, inference hyperparameters, and the randomness introduced by sampling encoder KL. As a result, scores cannot be accurately reproduced. Nonetheless, this can serve as a reference for designing one’s own benchmark.

Inference

Use AsymmetricMagVitV2 in your own code


from models.vae import AsymmetricMagVitV2Pipline
import torch
from models.utils.image_op import imdenormalize, imnormalize, read_video, read_image
import torchvision.transforms as transforms


device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.bfloat16
encoder_init_window = 17
input_path = "data/videos/tokyo_walk.mp4"
img_transform = transforms.Compose([transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
input, last_frame_id = read_video(input_path, encoder_init_window, sample_fps=8, img_transform, start=0)

model = AsymmetricMagVitV2Pipline.from_pretrained("BornFly/AsymmetricMagVitV2_16z").to(device, dtype).eval()
init_z, reg_log = model.encode(input, encoder_init_window, is_init_image=True, return_reg_log=True, unregularized=False)
init_samples = model.decode(init_z.to(device, dtype), decode_batch_size=1, is_init_image=True)

High-resolution video encoding and decoding, greater than 720p（spatial-temporal slice）

About Encoder hyperparameter configuration

slice frame spatial using： --max_siz --min_size
slice video temporal using： --encoder_init_window --encoder_window

If the GPU VRAM is not sufficient, metrics for evaluation can be adjusted to be between 256 and 512 at maximum.

About Decoder hyperparameter configuration

slice latent spatial using： --min_latent_size --max_latent_size

(default GPU VRAM needs to exceed 28GB. If the GPU VRAM is not sufficient, metrics for evaluation can be adjusted to be between 32=256p/8 and 64=512p/8 at maximum.)

slice latent temporal using： --decoder_init_window,

5 frames of latent space corresponds to 17 frames of the original video. The calculation formula is as follows: latent_T_dim = (frame_T_dim - 1) / temporal_downsample_num; in this model, temporal_downsample_num=4

1. encode & decode video

python infer_vae.py --input_path data/videos/tokyo-walk.mp4 --model_path vae_16z_bf16_hf  --output_folder vae_eval_out/vae_4z_bf16_hf_videos > infer_vae_video.log 2>&1

2. encode & decode image

python infer_vae.py --input_path data/images --model_path vae_16z_bf16_hf  --output_folder vae_eval_out/vae_4z_bf16_hf_images > infer_vae_image.log 2>&1

TODO List

Reproducing Sora, a 16-channel VAE integrated with SD3. Due to limited computational resources, the focus is on generating 1K high-definition dynamic wallpapers.
Reproducing VideoPoet, supporting multimodal joint representation. Due to limited computational resources, the focus is on generating music videos.

Contact Us

If there are any code-related questions, feel free to contact me via email——bornflyborntochange@outlook.com.
You need to scan the image to join the WeChat group or if it is expired, add this student as a friend first to invite you.

Reference

Open-Sora-Plan: https://github.com/PKU-YuanGroup/Open-Sora-Plan
Open-Sora: https://github.com/hpcaitech/Open-Sora
SVD: https://github.com/Stability-AI/generative-models

BornFly
/

AsymmetricMagVitV2_16z