PixArt-Σ: 突破性的4K文本到图像生成模型

PixArt-Σ简介

PixArt-Σ是由华为诺亚方舟实验室、大连理工大学、香港大学和香港科技大学的研究人员共同开发的一款先进的文本到图像生成模型。作为PixArt-α的升级版,PixArt-Σ采用了创新的弱到强训练方法,成功突破了之前模型的限制,能够直接生成高达4K分辨率的精美图像。

PixArt-Σ logo

核心特点与优势

PixArt-Σ相比其前身PixArt-α以及其他现有模型有以下几个显著的优势:

超高分辨率输出: PixArt-Σ能够直接生成4K(3840x2160)分辨率的图像,这在文本到图像生成领域是一个重大突破。
更长的文本输入: 模型支持处理最多300个T5 token的输入文本,相比PixArt-α的120个token有了很大提升,能够理解和呈现更加复杂和细致的描述。
升级的VAE: 采用了SDXL的VAE,进一步提升了图像质量和细节表现。
更高的图像保真度: 生成的图像在细节、质量和与文本提示的一致性方面都有显著提升。
灵活的分辨率支持: 除了4K之外,还提供了256px、512px、1024px和2K等多种分辨率的预训练模型。

模型架构与训练

PixArt-Σ基于Diffusion Transformer (DiT)架构,采用了创新的弱到强训练策略。这种方法允许模型从低分辨率逐步学习到高分辨率图像生成,有效克服了直接训练高分辨率模型的困难。

训练过程中使用了大规模的图文对数据集,并结合了多种先进的优化技术,如多尺度训练、渐进式学习等。这些策略共同确保了模型能够生成高质量、高分辨率的图像。

实际效果展示

为了直观地展示PixArt-Σ的强大能力,我们来看几组对比样例:

模型	样例1	样例2	样例3
PixArt-Σ
PixArt-α

从这些样例中我们可以清楚地看到,PixArt-Σ生成的图像在细节丰富度、色彩表现和整体质感上都有了显著提升。特别是在复杂场景和微小细节的呈现方面,PixArt-Σ展现出了卓越的能力。

使用指南

环境配置

要开始使用PixArt-Σ,首先需要配置适当的环境。建议使用Python 3.9或更高版本,并安装PyTorch 2.0.1+cu11.7。以下是基本的安装步骤:

conda create -n pixart python==3.9.0
conda activate pixart
conda install pytorch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 pytorch-cuda=11.7 -c pytorch -c nvidia

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

模型下载

PixArt-Σ提供了多个预训练模型,覆盖了从256px到4K的不同分辨率。您可以根据需求选择合适的模型:

# 下载SDXL-VAE和T5 checkpoints
git lfs install
git clone https://huggingface.co/PixArt-alpha/pixart_sigma_sdxlvae_T5_diffusers output/pretrained_models/pixart_sigma_sdxlvae_T5_diffusers

# 下载PixArt-Sigma checkpoints
python tools/download.py

快速开始

使用Gradio界面可以快速体验PixArt-Σ的图像生成能力:

python scripts/interface.py --model_path output/pretrained_models/PixArt-Sigma-XL-2-512-MS.pth --image_size 512 --port 11223

使用Diffusers库

PixArt-Σ已经集成到了流行的Diffusers库中,这使得在各种项目中使用该模型变得更加简单:

import torch
from diffusers import Transformer2DModel, PixArtSigmaPipeline

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
weight_dtype = torch.float16

transformer = Transformer2DModel.from_pretrained(
    "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS", 
    subfolder='transformer', 
    torch_dtype=weight_dtype,
    use_safetensors=True,
)
pipe = PixArtSigmaPipeline.from_pretrained(
    "PixArt-alpha/pixart_sigma_sdxlvae_T5_diffusers",
    transformer=transformer,
    torch_dtype=weight_dtype,
    use_safetensors=True,
)
pipe.to(device)

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