PixArt-α: 快速训练的图像生成新范式

PixArt-α:开启图像生成新时代

在人工智能生成内容(AIGC)领域,文本到图像(Text-to-Image, T2I)生成技术一直是研究热点。然而,目前最先进的T2I模型通常需要庞大的训练成本,动辄上百万GPU小时,不仅严重阻碍了AIGC社区的基础创新,还增加了大量的碳排放。为解决这一问题,来自华为诺亚方舟实验室、大连理工大学等机构的研究人员提出了PixArt-α,这是一种基于Transformer的T2I扩散模型,不仅图像生成质量可与当前最先进的图像生成器(如Imagen、SDXL甚至Midjourney)相媲美,而且训练速度显著超过现有的大规模T2I模型。

PixArt-α的核心设计

PixArt-α的成功主要得益于三个核心设计:

训练策略分解: 研究人员将训练过程分解为三个不同的步骤,分别优化像素依赖性、文本-图像对齐和图像美学质量。这种分解策略使得每个阶段都能专注于特定的目标,从而提高整体训练效率。
高效的T2I Transformer: 通过在Diffusion Transformer (DiT)中引入交叉注意力模块来注入文本条件,同时简化了计算密集型的类条件分支。这种设计既保证了模型的生成能力,又显著降低了计算复杂度。
高信息量数据: 研究team强调了文本-图像对中概念密度的重要性。他们利用大型视觉-语言模型自动标注密集的伪标题,以辅助文本-图像对齐学习。这种方法极大地提高了训练数据的质量和信息量。

PixArt-α架构图

惊人的训练效率

得益于这些创新设计,PixArt-α的训练速度远超现有的大规模T2I模型。具体来说:

PixArt-α仅需675个A100 GPU天,相比Stable Diffusion v1.5的6,250个GPU天,节省了89.2%的训练时间。
在经济成本上,PixArt-α仅需约$26,000,而Stable Diffusion v1.5则需要$320,000,节省了近$300,000。
在环境影响方面,PixArt-α减少了90%的碳排放。
与更大规模的SOTA模型RAPHAEL相比,PixArt-α的训练成本仅为其1%。

这些数据充分展示了PixArt-α在训练效率方面的巨大优势,为AIGC社区和初创企业从头构建高质量、低成本的生成模型提供了新的可能性。

高质量的图像生成

尽管大幅降低了训练成本,PixArt-α在图像生成质量上并未妥协。事实上,它生成的图像在质量、艺术性和语义控制方面都表现出色。以下是一些PixArt-α生成的图像示例:

PixArt-α生成的图像示例

这些图像展示了PixArt-α在各种场景下的强大生成能力,从逼真的人物肖像到奇幻的抽象艺术作品,都能呈现出高度的细节和创意。

PixArt-α的应用场景

除了基础的文本到图像生成,PixArt-α还可以与其他技术结合,拓展更多应用场景:

PixArt + Dreambooth: 通过Dreambooth技术,PixArt-α可以学习特定的对象或风格,实现更个性化的图像生成。
PixArt + ControlNet: 结合ControlNet技术,PixArt-α可以根据边缘图、草图等控制信息生成图像,提供更精确的生成控制。
PixArt + LCM (Latent Consistency Model): 通过与LCM结合,PixArt-α可以实现更快速的推理,在不到1秒的时间内生成高质量图像。

这些组合大大扩展了PixArt-α的应用范围,使其能够适应各种不同的图像生成需求。

如何使用PixArt-α

环境配置

要开始使用PixArt-α,首先需要配置适当的环境:

Python >= 3.9 (推荐使用Anaconda或Miniconda)
PyTorch >= 1.13.0+cu11.7

可以通过以下命令快速设置环境:

conda create -n pixart python=3.9
conda activate pixart
pip install torch==2.1.1 torchvision==0.16.1 torchaudio==2.1.1 --index-url https://download.pytorch.org/whl/cu118

git clone https://github.com/PixArt-alpha/PixArt-alpha.git
cd PixArt-alpha
pip install -r requirements.txt

模型下载

PixArt-α提供了多个预训练模型,可以根据需求选择下载:

T5: 用于文本编码
VAE: 用于图像编码和解码
PixArt-α: 不同分辨率的图像生成模型(256x256, 512x512, 1024x1024)
PixArt-δ: 结合LCM的快速推理模型
ControlNet: 用于控制图像生成的辅助模型

这些模型可以从Hugging Face或OpenXLab下载。

训练PixArt-α

对于想要自己训练PixArt-α的研究者,可以参考以下步骤:

准备数据集,可以使用SAM数据集或自定义数据集。
修改配置文件和数据加载器。
运行训练脚本,例如:

python -m torch.distributed.launch --nproc_per_node=2 --master_port=12345 train_scripts/train.py configs/pixart_config/PixArt_xl2_img256_SAM.py --work-dir output/train_SAM_256

此外,PixArt-α还支持Dreambooth、LCM、ControlNet等扩展训练,可以根据需求选择相应的训练方法。

推理和应用

PixArt-α提供了多种使用方式:

使用Gradio进行快速体验:

DEMO_PORT=12345 python app/app.py

在🧨 diffusers中集成:

from diffusers import PixArtAlphaPipeline
import torch

pipe = PixArtAlphaPipeline.from_pretrained("PixArt-alpha/PixArt-XL-2-1024-MS", torch_dtype=torch.float16, use_safetensors=True)
pipe.to("cuda")

prompt = "A small cactus with a happy face in the Sahara desert."
image = pipe(prompt).images[0]
image.save("./cactus.png")

在线演示:

PixArt-α在Hugging Face上提供了在线演示,用户可以直接体验模型的图像生成能力。

结语

PixArt-α作为一种快速、高效的文本到图像生成模型,不仅在生成质量上达到了业界领先水平,更重要的是大幅降低了训练成本和时间。这一突破性进展为AIGC社区和初创企业提供了新的可能性,使得从头构建高质量生成模型变得更加可行。

随着PixArt-α的开源和不断改进,我们可以期待看到更多基于此技术的创新应用和进一步的研究成果。无论是在艺术创作、设计辅助,还是在教育、娱乐等领域,PixArt-α都有望带来深远的影响。

研究团队希望PixArt-α能为AIGC社区提供新的洞见,加速构建高质量、低成本的生成模型,推动整个领域的快速发展。随着技术的不断进步,我们可以期待在不久的将来,AI生成的图像将在质量和多样性上达到新的高度,为人类的创造力提供更强大的工具和支持。