Skip to content

FoundationVision/UniTok

Repository files navigation

UniTok: A Unified Tokenizer
for Visual Generation and Understanding

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

Paper PDF Project Page

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

teaser

Built upon UniTok, we construct an MLLM capable of both multimodal generation and understanding with the Liquid framework, which sets a new state-of-the-art among unified autoregressive MLLMs.

teaser

News

2025-04-15: The gradio demo of UniTok MLLM is available on Huggingface now!

2025-04-02: A new checkpoint of UniTok is released, which has better downstream task performance by replacing the causal attention projection layer with full attention. The model weights of our unified MLLM are also available on Huggingface now!

2025-02-28: Paper, code, model, and project page for UniTok are all released.

Performance

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.41 70.8
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 Weights

Model Res. #Token Code Shape rFID Checkpoint
UniTok-Large 256 256 16 $\times$ 16 $\times$ 8 0.41 Download

Usage

Requirements

  • Python ≥ 3.10
  • PyTorch ≥ 2.3.1

Installation

git clone https://github.com/FoundationVision/UniTok.git
cd UniTok
pip install -r requirements.txt

Inference

Please download the checkpoint and fill in the ckpt_path.

python inference.py \
    --ckpt_path /path/to/unitok_tokenizer.pth \
    --src_img /path/to/test_img --rec_img /path/to/rec_img

Training

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.

Unified MLLM

We show that UniTok significantly boosts the performance of unified MLLMs.

Visual Understanding Performance on VQA Benchmarks.

Method LLM Res. VQAv2 GQA TextVQA POPE MME MM-Vet
Show-o Phi-1.5-1.3B 256 59.3 48.7 - 73.8 948 -
Liquid Gemma-7B 512 71.3 58.4 42.4 81.1 1119 -
VILA-U Llama-2-7B 256 75.3 58.3 48.3 83.9 1336 27.7
UniTok Llama-2-7B 256 76.8 61.1 51.6 83.2 1448 33.9

Visual Generation Performance on GenAI-Bench.

Method Type Count Differ Compare Logical Overall
Negate Universal
Show-o Discrete Diff. 0.70 0.62 0.71 0.51 0.65 0.60
VILA-U Autoregressive 0.70 0.71 0.74 0.53 0.66 0.64
Liquid Autoregressive 0.76 0.73 0.74 0.46 0.74 0.65
UniTok Autoregressive 0.76 0.79 0.74 0.46 0.73 0.67

Please refer to EVAL.md for more details.

Evaluation

We also 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.

Acknowledgement

UniTok is built upon the awesome works VAR, DataComp, Liquid, LLaVA, LlamaGen, and ViTamin.

LICENSE

This project is licensed under the MIT License - see the LICENSE file for details.

Citation

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