Github
Huggingface
论文项目介绍：imagen-pytorch
imagen-pytorch 是一个用 Pytorch 实现的基于文本生成图像的神经网络项目。它实现了 Google 的 Imagen，这是目前在文本到图像合成领域领先的技术，相较于著名的 DALL-E2，它在使用架构上更简单但效果更好。
项目背景
Imagen 使用了一个级联的扩散模型（DDPM），通过从 Google 预训练的 T5 模型中提取的文本嵌入进行条件化生成。这种架构避免了使用像 CLIP 这样的先验网络。此外，该模型还引入了动态裁剪以改善无分类器指导、噪声水平条件化以及高效的内存使用设计。
功能特点
Imagen 主要特点包括：
- 使用大规模预训练 T5 模型进行文本嵌入。
- 采用级联扩散模型，结合文本条件生成图像。
- 引入动态裁剪技术提高模型无分类器指导能力。
- 嵌入内存高效的卷积神经网络 (Unet) 设计。
安装与使用
安装
要使用 imagen-pytorch，只需运行以下命令即可安装：
$ pip install imagen-pytorch
使用示例
以下是一个简单的使用示例，展示如何使用此框架进行训练。
import torch
from imagen_pytorch import Unet, Imagen
unet1 = Unet(dim=32, cond_dim=512, dim_mults=(1, 2, 4, 8), num_resnet_blocks=3, layer_attns=(False, True, True, True))
unet2 = Unet(dim=32, cond_dim=512, dim_mults=(1, 2, 4, 8), num_resnet_blocks=(2, 4, 8, 8), layer_attns=(False, False, False, True))
imagen = Imagen(unets=(unet1, unet2), image_sizes=(64, 256), timesteps=1000, cond_drop_prob=0.1).cuda()
# 模拟图像和文本嵌入
text_embeds = torch.randn(4, 256, 768).cuda()
images = torch.randn(4, 3, 256, 256).cuda()
# 开始训练
for i in (1, 2):
loss = imagen(images, text_embeds=text_embeds, unet_number=i)
loss.backward()
贡献与感谢
此项目在社区的协作下得以不断完善，感谢多位贡献者的努力：
- StabilityAI 和其他赞助者的慷慨支持。
- 🤗 Huggingface 提供的优秀变压器库。
- Jonathan Ho 引领生成式人工智能变革的重要论文。
- Katherine Crowson 的代码，帮助理解连续时间高斯扩散。
使用场景
imagen-pytorch 可以被应用于多种场景，包括：
- 高质量图像生成，例如根据描述生成景观、物品或场景。
- 医学图像的合成研究。
- 视频生成及编辑，通过擅长的空间-时间扩散技术。
多GPU支持
使用 Huggingface 的 Accelerate 库，imagen-pytorch 可以轻松进行多 GPU 训练。只需通过命令行工具调整训练脚本配置，并使用 accelerate launch 命令运行即可。
未来发展
项目计划进一步优化，包括：
- 利用闪存注意力（Flash Attention）提升性能。
- 适配与 Dream Booth 相关的训练和微调。
- 支持更多场景的综合应用，如帧速率控制的视频扩散。
许可
imagen-pytorch 项目采用 MIT 许可证。用户可以自由使用、修改和分发此项目代码。
项目致力于保持透明和开放，希望每位有志于在人工智能领域探索和实践的人都能利用这一工具，创造出色的 AI 模型和应用。
