Chuofan Ma1,2 · Yi Jiang2† · Junfeng Wu2,3 · Jihan Yang1
Xin Yu1 · Zehuan Yuan2* · Bingyue Peng2 · Xiaojuan Qi1†*
1HKU 2ByteDance 3HUST
†project lead *corresponding author
This repo implements UniTok, a unified visual tokenizer well-suited for both generation and understanding tasks. It is compatiable with autoregressive generative models (e.g. LlamaGen), multimodal understanding models (e.g. LLaVA), and unified MLLMs (e.g. Chameleon and Liquid).
Built upon UniTok, we construct an MLLM capable of both multimodal generation and understanding, which sets a new state-of-the-art among unified autoregressive MLLMs. The weights of our MLLM will be released soon.
2025-02-14: Paper, code, and model weights for UniTok are all released.
Method | #Tokens | rFID ↓ | Accuracy |
---|---|---|---|
VQVAE Model | |||
VQ-GAN | 256 | 4.98 | -- |
RQ-VAE | 256 | 1.30 | -- |
VAR | 680 | 0.90 | -- |
CLIP Model | |||
CLIP | 256 | -- | 76.2 |
SigLIP | 256 | -- | 80.5 |
ViTamin | 256 | -- | 81.2 |
Unified Model | |||
TokenFlow † | 680 | 1.37 | -- |
VILA-U † | 256 | 1.80 | 73.3 |
UniTok | 256 | 0.39 | 70.5 |
UniTok † | 256 | 0.38 | 78.6 |
† indicates the model uses pretrained CLIP weights for initialization. Although CLIP weight initialization boosts ImageNet zero-shot accuracy, we notice that random initialization leads to better downstream understanding performance. We thus release the model checkpoint of UniTok that is trained from scratch.
Model | Res. | #Token | Code Shape | rFID | Checkpoint |
---|---|---|---|---|---|
UniTok-Large | 256 | 256 | 16 |
0.39 | Download |
- Python ≥ 3.10
- PyTorch ≥ 2.3.1
git clone https://github.com/FoundationVision/UniTok.git
cd UniTok
pip install -r requirements.txt
Please download the checkpoint and fill in the ckpt_path
.
python inference.py \
--ckpt_path /path/to/unitok/checkpoint \
--src_img /path/to/test_img --rec_img /path/to/rec_img
-
We train UniTok on DataComp-1B. Please follow the instructions to download and prepare the data.
-
Download the models used for loss calculation and put them in
./external
. -
Download the ImageNet validation set for zero-shot accuracy evaluation.
-
Download the ImageNet 256$\times$256 reference batch for FID evaluation.
Configure nnodes, nproc_per_node, node_rank, master_addr, master_port
in launch.sh
and run:
bash launch.sh \
--output_dir '/path/to/save/checkpoints/' \
--train_data '/path/to/datacomp/shards/{00000000..00140146}.tar' \
--imagenet_val '/path/to/imagenet_val/' \
--fid_eval_src '/path/to/imagenet_reference_batch' \
--fid_eval_dst '/path/to/save/imagenet_reconstructed_batch'
Note: For more hyper-parameter configurations, please check utils/config.py
.
We benchmark UniTok in terms of both understanding performance using the LLaVA framework and generation performance using the LLamaGen framework. Please refer to EVAL.md for more details.
UniTok is built upon the awesome works VAR, DataComp, LLaVA, LlamaGen, and ViTamin.
This project is licensed under the MIT License - see the LICENSE file for details.
If you find this project useful, please consider citing:
@article{unitok,
title={UniTok: A Unified Tokenizer for Visual Generation and Understanding},
author={Ma, Chuofan and Jiang, Yi and Wu, Junfeng and Yang, Jihan and Yu, Xin and Yuan, Zehuan and Peng, Bingyue and Qi, Xiaojuan},
journal={arXiv preprint arXiv:2502.20321},
year={2025}
}