自回归模型超越扩散模型：🦙 Llama 用于可扩展图像生成

此仓库包含预训练模型权重以及用于以下研究的 PyTorch（torch>=2.1.0）训练和采样代码：

自回归模型超越扩散模型：Llama 用于可扩展图像生成
孙培泽 (Peize Sun)、姜毅 (Yi Jiang)、陈守法 (Shoufa Chen)、张士龙 (Shilong Zhang)、彭冰悦 ()、罗平 (Ping Luo)、袁泽寰 (Zehuan Yuan) 香港大学、字节跳动

更多可视化效果请访问

🔥 更新

• [2024.06.28] 文本条件图像生成的图像编码器和自回归模型已发布！快来试试吧！
• [2024.06.15] vLLM 已支持从 1 亿到 30 亿参数的所有模型！
• [2024.06.11] 类别条件图像生成的图像编码器和自回归模型已发布！
• [2024.06.11] 代码和演示已发布！

🌿 简介

我们推出了 LlamaGen，这是一个全新的图像生成模型家族，它将大型语言模型中原始的“下一个词预测”范式应用于视觉生成领域。这有力地回答了这样一个问题：如果经过适当的规模扩展，像 Llama 这样的纯自回归模型——即对视觉信号“没有归纳偏置”的模型——是否能够达到最先进的图像生成性能？我们重新审视了图像编码器的设计空间、图像生成模型的可扩展性特性及其训练数据的质量。

在本仓库中，我们发布了：

• 下采样比例为 16 和 8 的两种图像编码器。
• 七种从 1 亿到 30 亿参数不等的类别条件生成模型。
• 两种 7 亿参数的文本条件生成模型。
• 在 上提供的在线演示，用于运行预训练模型。
• 支持 vLLM 服务框架，可实现 300%–400% 的加速。

🦄 ImageNet 上的类别条件图像生成

VQ-VAE 模型

方法	参数	编码	rFID (256x256)	权重
vq_ds16_c2i	72M	16x16	2.19	vq_ds16_c2i.pt
vq_ds16_c2i	72M	24x24	0.94	上述
vq_ds16_c2i	72M	32x32	0.70	上述
vq_ds8_c2i	70M	32x32	0.59	vq_ds8_c2i.pt

AR 模型

方法	参数	训练	编码	FID (256x256)	权重
LlamaGen-B	111M	DDP	16x16	5.46	c2i_B_256.pt
LlamaGen-B	111M	DDP	24x24	6.09	c2i_B_384.pt
LlamaGen-L	343M	DDP	16x16	3.80	c2i_L_256.pt
LlamaGen-L	343M	DDP	24x24	3.07	c2i_L_384.pt
LlamaGen-XL	775M	DDP	24x24	2.62	c2i_X_384L.pt
LlamaGen-XXL	1.4B	FSDP	24x24	2.34	c2i_XXL_384.pt
LlamaGen-3B	3.1B	FSDP	24x24	2.18	c2i_3B_384.pt

演示

请下载模型，将其放入 ./pretrained_models 文件夹中，然后运行：

python3 autoregressive/sample/sample_c2i.py --vq-ckpt ./pretrained_models/vq_ds16_c2i.pt --gpt-ckpt ./pretrained_models/c2i_L_384.pt --gpt-model GPT-L --image-size 384
# 或
python3 autoregressive/sample/sample_c2i.py --vq-ckpt ./pretrained_models/vq_ds16_c2i.pt --gpt-ckpt ./pretrained_models/c2i_XXL_384.pt --gpt-model GPT-XXL --from-fsdp --image-size 384

生成的图像将保存为 sample_c2i.png。

Gradio 演示

您也可以使用我们的在线 Gradio 演示 或在本地运行：

python app.py

🚀 文本条件图像生成

VQ-VAE 模型

方法	参数	编码	数据	权重
vq_ds16_t2i	72M	16x16	LAION COCO (50M) + 内部数据 (10M)	vq_ds16_t2i.pt

AR 模型

方法	参数	编码	数据	权重
LlamaGen-XL	775M	16x16	LAION COCO (50M)	t2i_XL_stage1_256.pt
LlamaGen-XL	775M	32x32	内部数据 (10M)	t2i_XL_stage2_512.pt

演示

在运行演示之前，请参阅 语言说明文档 以安装所需的软件包和语言模型。

请下载模型，将其放入 ./pretrained_models 文件夹中，然后运行：

python3 autoregressive/sample/sample_t2i.py --vq-ckpt ./pretrained_models/vq_ds16_t2i.pt --gpt-ckpt ./pretrained_models/t2i_XL_stage1_256.pt --gpt-model GPT-XL --image-size 256
# 或
python3 autoregressive/sample/sample_t2i.py --vq-ckpt ./pretrained_models/vq_ds16_t2i.pt --gpt-ckpt ./pretrained_models/t2i_XL_stage2_512.pt --gpt-model GPT-XL --image-size 512

生成的图像将保存为 sample_t2i.png。

本地 Gradio 演示

⚡ 服务

我们使用服务框架 vLLM 来提高吞吐量。请参阅 服务说明文档 以安装所需的软件包。

python3 autoregressive/serve/sample_c2i.py --vq-ckpt ./pretrained_models/vq_ds16_c2i.pt --gpt-ckpt ./pretrained_models/c2i_XXL_384.pt --gpt-model GPT-XXL --from-fsdp --image-size 384

生成的图像将保存为 sample_c2i_vllm.png。

入门指南

有关安装、训练和评估，请参阅 入门指南。

许可证

本项目的大部分内容采用 MIT 许可证授权。项目中的部分内容则依据所引用项目的单独许可证提供，具体详情请参阅相应文件。

BibTeX

@article{sun2024autoregressive,
  title={自回归模型胜过扩散模型：Llama 用于可扩展的图像生成},
  author={孙培泽、蒋毅、陈守法、张世龙、彭冰悦、罗平、袁泽寰},
  journal={arXiv 预印本 arXiv:2406.06525},
  year={2024}
}

LlamaGen 快速上手指南

LlamaGen 是一个将大语言模型（LLM）的“下一词预测”范式应用于图像生成的自回归模型系列。它证明了无需视觉归纳偏置的纯自回归模型，在适当扩展后也能达到最先进的图像生成效果。本指南帮助你快速在本地运行预训练模型。

1. 环境准备

在开始之前，请确保你的开发环境满足以下要求：

• 操作系统: Linux (推荐) 或 macOS
• Python: >= 3.8
• PyTorch: >= 2.1.0 (必须支持 CUDA)
• 硬件: 推荐配备 NVIDIA GPU (显存建议 16GB 以上以运行较大模型，如 3B 版本；小模型可在较小显存上运行)

前置依赖安装：

建议使用 pip 安装核心依赖。国内用户推荐使用清华源或阿里源加速下载。

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
pip install -r requirements.txt

注意：如果项目中未提供 requirements.txt，请至少安装以下关键库：

pip install transformers accelerate gradio einops timm

若需使用 vLLM 进行加速推理（速度提升 300%-400%），请额外安装：

pip install vllm

2. 安装步骤与模型下载

步骤一：克隆代码库

git clone https://github.com/FoundationVision/LlamaGen.git
cd LlamaGen

步骤二：下载预训练模型

你需要下载 图像分词器 (VQ-VAE) 和 自回归生成模型 (AR Model) 的权重文件。

1. 在项目根目录下创建文件夹：

   mkdir -p pretrained_models
   

2. 从 Hugging Face 下载模型并放入 pretrained_models 目录。

   • 类别条件生成 (Class-conditional) 示例 (以 LlamaGen-L 384x384 为例):

     • VQ 模型: vq_ds16_c2i.pt
     • GPT 模型: c2i_L_384.pt

   • 文本条件生成 (Text-conditional) 示例:

     • 需参考 language/README.md 安装额外的语言模型依赖。
     • VQ 模型: vq_ds16_t2i.pt
     • GPT 模型 (Stage 1): t2i_XL_stage1_256.pt

   (国内网络访问 Hugging Face 较慢时，可使用镜像站如 hf-mirror.com 或手动下载后上传至服务器)

3. 基本使用

以下展示最基础的类别条件图像生成流程。生成的图片将保存为 sample_c2i.png。

运行示例 (类别条件生成)

确保模型文件已位于 ./pretrained_models 目录下，然后运行以下命令：

python3 autoregressive/sample/sample_c2i.py --vq-ckpt ./pretrained_models/vq_ds16_c2i.pt --gpt-ckpt ./pretrained_models/c2i_L_384.pt --gpt-model GPT-L --image-size 384

如果你使用的是通过 FSDP 训练的大模型（如 XXL 或 3B 版本），需添加 --from-fsdp 参数：

python3 autoregressive/sample/sample_c2i.py --vq-ckpt ./pretrained_models/vq_ds16_c2i.pt --gpt-ckpt ./pretrained_models/c2i_XXL_384.pt --gpt-model GPT-XXL --from-fsdp --image-size 384

运行示例 (文本条件生成)

前提：已完成 language/README.md 中的额外依赖安装。

python3 autoregressive/sample/sample_t2i.py --vq-ckpt ./pretrained_models/vq_ds16_t2i.pt --gpt-ckpt ./pretrained_models/t2i_XL_stage1_256.pt --gpt-model GPT-XL --image-size 256

启动本地 Web Demo (Gradio)

如果你想通过网页界面交互式地生成图像，可以运行：

python app.py

启动后，终端会显示本地访问地址（通常为 http://127.0.0.1:7860），在浏览器打开即可使用。

高性能推理 (可选)

若已安装 vLLM，可使用以下命令获得更快的生成速度：

python3 autoregressive/serve/sample_c2i.py --vq-ckpt ./pretrained_models/vq_ds16_c2i.pt --gpt-ckpt ./pretrained_models/c2i_XXL_384.pt --gpt-model GPT-XXL --from-fsdp --image-size 384