Open-MAGVIT2 - 自回归视觉生成新突破大幅提升图像分词性能

Open-MAGVIT2: 民主化自回归视觉生成

尽管VQGAN在码本大小和利用率方面存在限制，低估了其潜力，但它在自回归视觉生成中仍然至关重要。MAGVIT2通过无查找技术和大型码本（$2^{18}$）解决了这些问题，在图像和视频生成方面展现了promising的结果，并在VideoPoet中发挥了关键作用。然而，我们目前无法获取这个分词器。:broken_heart:

在我们的代码库中，我们用PyTorch重新实现了MAGVIT2分词器，closely复现了原始结果。我们希望我们的努力能够在自回归视觉生成领域促进创新和创造力。:green_heart:

📰 新闻

[2024.06.17] :fire::fire::fire: 我们发布了图像分词器的训练代码和不同分辨率的检查点，与VQGAN、MaskGIT以及最近的TiTok、LlamaGen和OmniTokenizer相比，达到了最先进的性能（8倍下采样的rFID为0.39）。

🎤 待办事项

通过扩大训练规模改进图像分词器。
完成自回归模型的训练。
视频分词器及相应的自回归模型。

🤗 Open-MAGVIT2仍处于早期阶段，正在积极开发中。敬请期待更新！

📖 实现

图1. Open-MAGVIT2分词器的框架，由编码器、无查找量化器（LFQ）和解码器组成。

🛠️ 安装

环境：我们已在Python 3.8.8和CUDA 11.7上测试过（其他版本也可能适用）。
依赖项：pip install -r requirements
数据集

imagenet
└── train/
    ├── n01440764
        ├── n01440764_10026.JPEG
        ├── n01440764_10027.JPEG
        ├── ...
    ├── n01443537
    ├── ...
└── val/
    ├── ...

第一阶段：视觉分词器的训练

我们遵循MAGVIT-2中生成器的设计，但在GAN训练中使用PatchGAN而非StyleGAN作为判别器。我们结合了MAGVIT-2和VQGAN中使用的损失函数，以获得更好的训练稳定性和重建质量。所有训练细节可在配置文件中找到。请注意，我们使用32个V100进行模型训练。

🍺 定量比较

表1. 不同分词器在256×256 ImageNet 50k验证集上的重建性能。Open-MAGVIT2在不同下采样率上达到了最先进的结果。

方法	标记类型	标记数	训练数据	码本大小	rFID	PSNR	码本利用率	检查点
VQGAN	2D	16 × 16	256 × 256 ImageNet	1024	7.94	19.4	-	-
SD-VQGAN	2D	16 × 16	OpenImages	16384	5.15	-	-	-
MaskGIT	2D	16 × 16	256 × 256 ImageNet	1024	2.28	-	-	-
LlamaGen	2D	16 × 16	256 × 256 ImageNet	16384	2.19	20.79	97%	-
:fire:Open-MAGVIT2	2D	16 × 16	256 × 256 ImageNet	262144	1.53	21.53	100%	IN256_Base
ViT-VQGAN	2D	32 × 32	256 × 256 ImageNet	8192	1.28	-	-	-
VQGAN	2D	32 × 32	OpenImages	16384	1.19	23.38	-	-
SD-VQGAN	2D	32 × 32	OpenImages	16384	1.14	-	-	-
OmniTokenizer-VQ	2D	32 × 32	256 × 256 ImageNet	8192	1.11	-	-	-
LlamaGen	2D	32 × 32	256 × 256 ImageNet	16384	0.59	24.45	-	-
:fire:Open-MAGVIT2*	2D	32 × 32	128 × 128 ImageNet	262144	0.39	25.78	100%	IN128_Base
SD-VQGAN	2D	64 × 64	OpenImages	16384	0.58	-	-	-
TiTok-L	1D	32	256 × 256 ImageNet	4096	2.21	-	-	-
TiTok-B	1D	64	256 × 256 ImageNet	4096	1.70	-	-	-
TiTok-S	1D	128	256 × 256 ImageNet	4096	1.71	-	-	-

(*)表示结果来自直接使用128×128分辨率训练的模型进行推理，未经微调。 表2. 通过在128×128分辨率下训练和测试,与原始MAGVIT2进行比较,如其原始论文中所使用。使用ImageNet 50k验证集进行测试。

方法	令牌类型	令牌数量	数据	LFQ	大型码本	上/下采样器	rFID	URL
MAGVIT2	2D	16×16	128×128 ImageNet	√	√	√	1.21	-
Open-MAGVIT2	2D	16×16	128×128 ImageNet	√	√	√	1.56	IN128_Base

:eyes: 重建可视化

图2. 在256×256分辨率下训练并在256×256分辨率下测试的Open-MAGVIT2分词器的可视化效果(imagenet_256_Base版本)。(a)表示原始图像,而(b)表示重建图像。

图3. 在128×128分辨率下训练并在512×512分辨率下测试的Open-MAGVIT2分词器的可视化效果(imagenet_128_Base版本)。(a)表示原始图像,而(b)表示重建图像。

🚀 训练脚本

128×128分词器训练

bash run_128_B.sh

256×256分词器训练

bash run_256_B.sh

🚀 评估脚本

128×128分词器评估

python evaluation.py --config_file configs/imagenet_lfqgan_128_B.yaml --ckpt_path "Your Path" --image_size 128

256×256分词器评估

python evaluation.py --config_file configs/imagenet_lfqgan_256_B.yaml --ckpt_path "Your Path" --image_size 256

阶段II: 自回归生成训练

MAGVIT2使用非自回归transformer进行图像生成。相反,我们希望利用相对较大的码本来探索自回归视觉生成的潜力。我们目前正在探索阶段II的训练。

❤️ 致谢

我们感谢Lijun Yu的鼓励性讨论。我们参考了很多来自VQGAN和MAGVIT的内容。感谢他们出色的工作。

✏️ 引用

如果您发现这个代码库有帮助,请引用它。

@software{Luo_Open-MAGVIT2_2024,
author = {Luo, Zhuoyan and Shi, Fengyuan and Ge, Yixiao},
month = jun,
title = {{Open-MAGVIT2}},
url = {https://github.com/TencentARC/Open-MAGVIT2},
version = {1.0},
year = {2024}
}

@inproceedings{
yu2024language,
title={Language Model Beats Diffusion - Tokenizer is key to visual generation},
author={Lijun Yu and Jose Lezama and Nitesh Bharadwaj Gundavarapu and Luca Versari and Kihyuk Sohn and David Minnen and Yong Cheng and Agrim Gupta and Xiuye Gu and Alexander G Hauptmann and Boqing Gong and Ming-Hsuan Yang and Irfan Essa and David A Ross and Lu Jiang},
booktitle={The Twelfth International Conference on Learning Representations},
year={2024},
url={https://openreview.net/forum?id=gzqrANCF4g}
}