README.md

metadata

license: mit
pipeline_tag: any-to-any
library_name: transformers

UniTok: A Unified Tokenizer
for Visual Generation and Understanding

Chuofan Ma^1,2 · Yi Jiang^2† · Junfeng Wu^2,3 · Jihan Yang¹
Xin Yu¹ · Zehuan Yuan^2* · Bingyue Peng² · Xiaojuan Qi^1†*

¹HKU   ²ByteDance   ³HUST
†project lead   *corresponding author

This repository implements UniTok, a unified visual tokenizer well-suited for both generation and understanding tasks. It is compatible 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 with the Liquid framework, which sets a new state-of-the-art among unified autoregressive MLLMs.

Abstract

Visual generative and understanding models typically rely on distinct tokenizers to process images, presenting a key challenge for unifying them within a single framework. Recent studies attempt to address this by connecting the training of VQVAE (for autoregressive generation) and CLIP (for understanding) to build a unified tokenizer. However, directly combining these training objectives has been observed to cause severe loss conflicts. In this paper, we show that reconstruction and semantic supervision do not inherently conflict. Instead, the underlying bottleneck stems from limited representational capacity of discrete token space. Building on these insights, we introduce UniTok, a unified tokenizer featuring a novel multi-codebook quantization mechanism that effectively scales up the vocabulary size and bottleneck dimension. In terms of final performance, UniTok sets a new record of 0.38 rFID and 78.6% zero-shot accuracy on ImageNet. Besides, UniTok can be seamlessly integrated into MLLMs to unlock native visual generation capability, without compromising the understanding performance. Additionally, we show that UniTok favors cfg-free generation, reducing gFID from 14.6 to 2.5 on ImageNet 256$\times$256 benchmark. GitHub: this https URL .

News

2025-09-18: UniTok is accepted at NeurIPS 2025 as a spotlight.

2025-05-19: We find UniTok favors generation without classifier-free-guidance -- it reduces gFID (without cfg) from 14.6 to 2.51 on ImageNet 256x256 using LlamaGen-XXL as the generator. Please refer to the updated EVAL.md for more details.

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

Unified MLLM Inference

The simplest code for Lumina-mGPT inference:

from inference_solver import FlexARInferenceSolver
from PIL import Image

# ******************** Image Generation ********************
inference_solver = FlexARInferenceSolver(
    model_path="Alpha-VLLM/Lumina-mGPT-7B-768",
    precision="bf16",
    target_size=768,
)

q1 = f"Generate an image of 768x768 according to the following prompt:
" \
     f"Image of a dog playing water, and a waterfall is in the background."

# generated: tuple of (generated response, list of generated images)
generated = inference_solver.generate(
    images=[],
    qas=[[q1, None]],
    max_gen_len=8192,
    temperature=1.0,
    logits_processor=inference_solver.create_logits_processor(cfg=4.0, image_top_k=2000),
)

a1, new_image = generated[0], generated[1][0]


# ******************* Image Understanding ******************
inference_solver = FlexARInferenceSolver(
    model_path="Alpha-VLLM/Lumina-mGPT-7B-512",
    precision="bf16",
    target_size=512,
)

# "<|image|>" symbol will be replaced with sequence of image tokens before fed to LLM
q1 = "Describe the image in detail. <|image|>"

images = [Image.open("image.png")]
qas = [[q1, None]]

# `len(images)` should be equal to the number of appearance of "<|image|>" in qas
generated = inference_solver.generate(
    images=images,
    qas=qas,
    max_gen_len=8192,
    temperature=1.0,
    logits_processor=inference_solver.create_logits_processor(cfg=4.0, image_top_k=2000),
)

a1 = generated[0]
# generated[1], namely the list of newly generated images, should typically be empty in this case.


# ********************* Omni-Potent *********************
inference_solver = FlexARInferenceSolver(
    model_path="Alpha-VLLM/Lumina-mGPT-7B-768-Omni",
    precision="bf16",
    target_size=768,
)

# Example: Depth Estimation
# For more instructions, see demos/demo_image2image.py
q1 = "Depth estimation. <|image|>"
images = [Image.open("image.png")]
qas = [[q1, None]]

generated = inference_solver.generate(
    images=images,
    qas=qas,
    max_gen_len=8192,
    temperature=1.0,
    logits_processor=inference_solver.create_logits_processor(cfg=1.0, image_top_k=200),
)

a1 = generated[0]
new_image = generated[1][0]

Training

• 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 under ./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.

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