为任何(中文)文本生成生动的图像
新闻! 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 参数)的变压器,用于通用领域的文本到图像生成。
- 阅读 我们在 ArXiv 上的论文 CogView: Mastering Text-to-Image Generation via Transformers 了解正式介绍。PB-relax 和 Sandwich-LN 也可以帮助您稳定地训练大规模和深层的变压器(例如,消除 NaN 损失)。
- 访问 我们在 Github 页面 或 悟道 的演示!(目前仅支持简体中文输入,无后期选择或超分辨率,但可以将其他语言的文本翻译成中文输入。注意:悟道 为中国大陆用户提供更快的访问。)
- 下载 我们的预训练模型从 清华云盘。
- 引用 如果您发现我们的工作有帮助,请引用我们的论文~
@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}
快速开始
设置
-
硬件:推荐使用带 Nvidia V100 或 A100 的 Linux 服务器,但也可以在较不强大的 GPU 上运行预训练模型,使用较小的
--max-inference-batch-size
或训练较小的模型。 -
环境(选项 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 # 在容器中
下载
- 从 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
-
从 Project Wudao-Wenhui 下载模型。
文件名 描述 cogview-base.tar 预训练文本到图像模型。 cogview-caption.tar 调优后的图像到文本模型,也用于重新排序。 cogview-sr.tar 调优后的超分辨率模型。(警告:运行速度较慢。) 将它们解压到
pretrained/cogview/
文件夹中。以下命令应根据模型名称进行修改。tar -xvf cogview-{base, sr, caption}.tar -C pretrained/cogview/
-
(仅用于训练教程,跳过推理部分。)从我们的链接在清华云盘下载一个小的“鸟类和动物”示例数据集。
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 [path or "interactive"]
。输入文件的路径,也可以是“interactive”,这将启动一个 CLI。--output-path [path]
。结果保存的文件夹。--batch-size [int]
。每个查询生成的样本数。--max-inference-batch-size [int]
。每个前向过程的最大批量大小。如遇 OOM,可减少此值。--debug
。只保存所有生成样本的拼接图像,并以输入文本和日期命名。--with-id
。启用时,您必须在每个输入前指定一个“id”,例如001\t一个漂亮的女孩
,\t 表示 TAB(不是空格)。它将在每个输入的文件夹中生成batch-size
拆分图像。与--debug
冲突。--device [int]
。在哪个 GPU 上运行。
超分辨率
运行以下脚本并输入 text\t{image_path}
,其中 {image_path}
表示先前生成的图像的路径。
./scripts/super_resolution.sh
注意:该功能仅对从我们的图像分词器生成的图像有效(由于 token 分布)。
图像到文本
输入是一行一个图像路径,结果将打印到标准输出。
./scripts/image2text.sh
注意:此任务未优化,因此可能不太具有竞争力(但还可以)。我们将考虑在未来发布一个版本,在此任务上进行更长时间的微调。(待完成)
后选择
此应用程序仅接受文件输入,每行格式为 {text}\t{image_path1}\t{image_path2}\t{image_path3}...
。
输出是 {output_path}/scores.txt
,每一行为一个分数列表,跟随输入行。
./scripts/post_selection.sh
注意:在发布的代码中,为了简化,我们没有公开原始 API,它支持一些高级生成模式,例如文本和部分图像。
训练
在本教程中,我们使用了鸟类和动物数据集的一个子集。二进制数据集由我们的 cogdata toolkit 生成。请等待 cogdata 的正式发布及其教程(尽管现在已经可用)。
单节点
下载数据集后,直接运行
./scripts/pretrain_single_node.sh
多节点
如果您想在通过 infiniband 互连的多个服务器上训练模型,而无需共享文件系统(您可能需要 pdsh
来加速此过程):
- 在每个服务器上,使用
git clone
下载此仓库,并确保数据(LMDB 格式)移动到data
子文件夹中。 - 在每个服务器上,
echo "ip1 ip2 <其他 IP>" > ./docker/ip_list.txt
,然后通过./env/start_docker.sh
启动 docker。 - 通过
docker exec -it bg-cogview bash
进入第一个节点容器中的 docker。 - 进入
/root/cogview
,运行./scripts/pretrain_multiple_nodes.sh
。您可能需要更改 shell 脚本中的配置(尤其是OPTIONS_NCCL
)。
详见 arguments.py
的高级函数训练。
待完成