O2-MAGVIT2

Introduction

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.

Architecture

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

Quick start

Inference: edit the arguments in scripts/run_inference.sh and run the following command to see the reconstruction result:
```
bash scripts/run_inference.sh
```
run python inference.py -h for more details.

Training: edit the config under the configs/ then run the following command to train the model:

NODE_RANK=0
MASTER_ADDR=localhost:25001
NUM_NODES=1
NUM_GPUS=8

bash scripts/run_train_3d.sh $NODE_RANK $MASTER_ADDR $NUM_NODES $NUM_GPUS

Training Procedure

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.

Hyper parameters

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.

Pretrained Models

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.

Acknowledgement

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.

Citation

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}
}