为任何（中文）文本生成生动的图像

最新消息！ ImageReward 的论文已被 NeurIPS 2023 接收！

最新消息！ ImageReward 的代码（论文链接）已在 https://github.com/THUDM/ImageReward 上发布！ImageReward 是首个通用的文本到图像人类偏好 RM。

最新消息！ CogView2 的代码（论文链接）已在 https://github.com/THUDM/CogView2 上发布！

最新消息！ 更好、更快的 CogView2（正式版，2022年3月）的演示现已上线！最新模型也支持英文输入，但通常将其翻译成中文会得到更好的效果。

最新消息！ 更好、更快的 CogView2（新版本）的演示现已上线！

最新消息！ CogView 的论文已被 NeurIPS 2021 接收！

CogView 是一个预训练的（4B参数）Transformer 模型，用于通用领域的文本到图像生成。

• 阅读 我们的论文 CogView: 通过 Transformer 掌握文本到图像生成 以获得正式介绍。PB-relax 和 Sandwich-LN 也有助于稳定地训练大型和深层 Transformer（例如消除 NaN 损失）。
• 访问 我们的演示页面：Github 页面 或 Wudao！（目前不进行后选或超分辨率处理，仅支持简体中文输入，但用户可以将其他语言的文本翻译成中文后再输入。注意：Wudao 为中国大陆用户提供更快速的访问。）
• 下载 我们的预训练模型，请访问 清华大学云。
• 引用 我们的论文，如果您觉得我们的工作有所帮助：

@article{ding2021cogview,
  title={CogView: Mastering Text-to-Image Generation via Transformers},
  author={Ding, Ming and Yang, Zhuoyi and Hong, Wenyi and Zheng, Wendi and Zhou, Chang and Yin, Da and Lin, Junyang and Zou, Xu and Shao, Zhou and Yang, Hongxia and Tang, Jie},
  journal={arXiv preprint arXiv:2105.13290},
  year={2021}

• Google Colab 两位贡献者已成功在 Colab 上搭建了 CogView ！

入门指南

环境搭建

• 硬件：建议使用配备 Nvidia V100 或 A100 显卡的 Linux 服务器，但也可以使用较小的 --max-inference-batch-size 来运行预训练模型，或者在性能较弱的 GPU 上训练小型模型。

• 环境（选项 1）：请先安装 PyTorch（>=1.7.0）和 apex，然后通过 pip install -r requirements.txt 安装其他依赖项。

• 环境（选项 2）：我们准备了一个 Docker 镜像，以防您无法正确配置环境。拉取镜像并创建后台容器后进入：

  docker pull cogview/cuda111_torch181_deepspeed040
  ./env/start_docker.sh && docker exec -it bg-cogview bash
  
  cd /root/cogview # 在容器内
  

下载

0. 从 BAAI 网站 或清华大学云下载图像分词器 vqvae_hard_biggerset_011.pt。将文件放置在 pretrained/vqvae 目录下。

wget 'https://cloud.tsinghua.edu.cn/f/71607a5dca69417baa8c/?dl=1' -O pretrained/vqvae/vqvae_hard_biggerset_011.pt

1. 从 Project Wudao-Wenhui 下载模型。

   文件名	描述
   cogview-base.tar	预训练的文本到图像模型。
   cogview-caption.tar	微调后的图像到文本模型，也用于重新排序。
   cogview-sr.tar	微调后的超分辨率模型。（警告：运行速度较慢。）

   将它们解压到 pretrained/cogview/ 目录中。以下命令需根据模型名称进行调整。

   tar -xvf cogview-{base, sr, caption}.tar -C pretrained/cogview/
   

2. （仅适用于训练教程，推理时可跳过。）从我们在清华大学云上的链接下载一个小的“鸟类与动物”示例数据集。

wget https://cloud.tsinghua.edu.cn/f/1e4963ec8ac84941ba68/?dl=1 -O data/bird_animal.bin

运行 CogView！（模型推理）

我们将生成函数封装成了脚本。详细信息请参阅 generate_samples.py 和 arguments.py。

文本到图像生成

将文本查询（每行一条）写入 input.txt，然后运行：

./scripts/text2image.sh --debug

结果将保存在新的文件夹 samples_text2image/ 中。

推理时常用的参数主要有：

• --input-source [路径或 "interactive"]。输入文件的路径，也可以设置为“interactive”，此时会启动一个命令行界面。
• --output-path [路径]。保存结果的文件夹。
• --batch-size [整数]。每次查询生成的样本数量。
• --max-inference-batch-size [整数]。每次前向传播的最大批量大小。如果出现内存不足问题，请适当降低该值。
• --debug。仅保存所有生成样本的拼接图像，并按输入文本和日期命名。
• --with-id。启用此选项时，必须在每个输入前指定一个“id”，例如 001\t一个漂亮的女孩，其中 \t 表示制表符（不是空格）。它会为每个输入生成一个名为“id”的文件夹，里面包含 batch-size 张分割图像。此选项与 --debug 冲突。
• --device [整数]。指定使用的 GPU 编号。

超分辨率

运行以下脚本，并输入 text\t{image_path}，其中 {image_path} 是之前生成的图像路径。

./scripts/super_resolution.sh

注意：此功能仅对由我们的图像分词器生成的图像有效（由于标记分布的原因）。

图像到文本

输入格式为“每行一个图像路径”，结果将打印到标准输出。

./scripts/image2text.sh

注意：该模型并未针对此任务进行优化，因此可能竞争力不足（但尚可接受）。未来我们将考虑发布一个在此任务上长期微调的版本。（待办事项）

后选

此应用仅接受文件输入，每行格式为 {text}\t{image_path1}\t{image_path2}\t{image_path3}...。 输出为 {output_path}/scores.txt，每一行对应输入的一行，并列出相应的分数。

./scripts/post_selection.sh

注意：在发布的代码中，为了简化起见，我们并未公开原始 API，而该 API 支持一些高级生成模式，例如文本和部分图像的结合。

训练

在这里，我们使用来自鸟兽数据集的一个子集作为教程材料。二进制数据集是由我们的 cogdata 工具包 生成的。请耐心等待 cogdata 的正式发布及其相关教程（尽管目前已有可用版本）。

单节点

下载数据集后，直接运行：

./scripts/pretrain_single_node.sh

多节点

如果你想在通过 InfiniBand 互连的多台服务器上训练模型，且没有共享文件系统（你可能需要 pdsh 来加速这一过程）：

1. 在每台服务器上，使用 git clone 克隆本仓库，并确保将数据（LMDB 格式）移至 data 子文件夹中。
2. 在每台服务器上，执行 echo "ip1 ip2 <其他IP>" > ./docker/ip_list.txt，然后通过 ./env/start_docker.sh 启动 Docker 容器。
3. 使用 docker exec -it bg-cogview bash 进入第一节点的 Docker 容器。
4. 进入 /root/cogview 目录，运行 ./scripts/pretrain_multiple_nodes.sh。你可能需要修改该 Shell 脚本中的配置（尤其是 OPTIONS_NCCL）。

有关训练的高级功能，请参阅 arguments.py。 待办

图库

CogView 快速上手指南

CogView 是一个基于 Transformer 的预训练文本生成图像模型（40 亿参数），支持通过中文描述生成生动图像。本指南将帮助您快速在本地部署并运行该模型。

环境准备

系统要求

• 操作系统：Linux
• GPU：推荐 Nvidia V100 或 A100。显存较小的显卡可通过减小 --max-inference-batch-size 参数运行。
• 软件依赖：
  • PyTorch >= 1.7.0
  • NVIDIA Apex

前置依赖安装

请确保已安装 CUDA 驱动。随后安装 PyTorch 和 Apex，最后安装其他 Python 依赖：

# 安装 PyTorch (请参考 pytorch.org 选择对应 CUDA 版本的命令)
# 安装 Apex
git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --no-cache-dir ./

# 返回项目根目录后安装其他依赖
pip install -r requirements.txt

提示：如果环境配置困难，推荐使用官方提供的 Docker 镜像（见下文安装步骤选项二）。

安装步骤

1. 获取代码

git clone https://github.com/THUDM/CogView.git
cd CogView

2. 环境配置（二选一）

选项一：本地环境 完成上述“环境准备”中的依赖安装即可。

选项二：Docker 环境（推荐） 拉取预构建镜像并启动容器：

docker pull cogview/cuda111_torch181_deepspeed040
./env/start_docker.sh && docker exec -it bg-cogview bash
cd /root/cogview

3. 下载模型权重

国内用户推荐使用清华云盘或悟道平台加速下载。

步骤 3.1：下载图像分词器 (VQVAE)

mkdir -p pretrained/vqvae
wget 'https://cloud.tsinghua.edu.cn/f/71607a5dca69417baa8c/?dl=1' -O pretrained/vqvae/vqvae_hard_biggerset_011.pt

步骤 3.2：下载主模型 从 清华云盘 或 悟道平台 下载以下文件：

• cogview-base.tar (文生图主模型)
• cogview-caption.tar (图生文/重排序模型，可选)
• cogview-sr.tar (超分辨率模型，可选，推理较慢)

将下载的 .tar 文件解压到 pretrained/cogview/ 目录：

mkdir -p pretrained/cogview
tar -xvf cogview-base.tar -C pretrained/cogview/
# 如需超分功能，继续解压
# tar -xvf cogview-sr.tar -C pretrained/cogview/

基本使用

文生图 (Text-to-Image)

这是最常用的功能。您需要准备一个输入文件，每行包含一段中文描述。

1. 创建输入文件 input.txt：

   一只在草地上奔跑的金色 retriever 犬
   赛博朋克风格的城市夜景，霓虹灯闪烁
   

2. 运行生成脚本：

   ./scripts/text2image.sh --debug
   

   • --debug：将生成的所有图片拼接保存，方便预览。
   • 结果将保存在 samples_text2image/ 文件夹中。

3. 常用参数说明：

   • --input-source "interactive"：跳过文件，直接进入命令行交互模式输入提示词。
   • --batch-size [int]：每个提示词生成的图片数量。
   • --max-inference-batch-size [int]：单次前向传播的最大批次大小。若遇到显存不足 (OOM)，请减小此数值。
   • --device [int]：指定使用的 GPU ID。

超分辨率 (Super-Resolution)

对 CogView 生成的低分辨率图像进行高清化处理。

1. 准备输入文件，格式为：文本描述<TAB>图片路径。
2. 运行脚本：

   ./scripts/super_resolution.sh
   

   注意：该功能仅对由 CogView 自身生成的图像效果最佳。

交互式命令行示例

如果您想直接测试而不创建文件，可以使用交互模式：

./scripts/text2image.sh --input-source interactive --debug

启动后，直接在终端输入中文提示词并回车即可看到生成结果。