We present an open-source pytorch implementation of Google's MAGVIT-v2 visual tokenizer named O2-MAGVIT2, which stands for handling dual modality (Image and Video) tokenization with a single Tokenizer. O2-MAGVIT2 is aligned with MAGVIT-v2 to a large extent. It uses Lookup-free quantizer(LFQ) with codebook size of $2^{18}$ and the exact same architecture of the encoder, decoder and discriminator described in the original paper. To facilitate training, we use huggingface's accelerate to wrap the trainer. We also release a preview version of the video tokenizer trained on a Panda-70M subset to validate its performance.

We re-implemented the MAGVIT-v2's architecture exactly. Below is from the magvit-v2's attachments.

The whole training includes two stages. In stage I, we train an image tokenizer with OpenImage dataset (which contains 8M training samples) for 10 epochs with batch size 256. For stage II, we random sampled 9.3M samples from panda-70M and train the video tokenizer for 1 epoch with batch size 128.

We adopt almost the same hyper-parameter setting as MAGVIT-v2 with minimal change. See configs/magvit2_3d_model_config.yaml for model setup details and configs/magvit2_3d_train_config.yaml for training setup.

We release a pretrained checkpoint of the video tokenizer on huggingface as a preview. Note that due to much fewer training steps, the model is certainly under-trained and thus may not provide a good enough performance if you use it directly. We recommend considering it a step stone to continue training to get better results.

The checkpoint of O2-MAGVIT2-preview can be found here.

We refer some ideas and implementations from MAGVIT, vector-quantize-pytorch, praxis, LlamaGen, pytorch-image-models, and VQGAN. Thanks a lot for their excellent work.

If you found our work interesting, please cite the following references and give us a star.

@misc{Fang_O2-MAGVIT2,
author = {Fang, Xuezhi and Yao, Yiqun and Jiang, Xin and Li, Xiang and Yu, Naitong and Wang, Yequan},
license = {Apache-2.0},
title = {O2-MAGVIT2},
year = {2024},
url = {https://github.com/cofe-ai/O2-MAGVIT2}
}