👉 PixArt-Σ: 用于4K文本到图像生成的扩散变换器的弱到强训练

本仓库包含我们论文中探索的用于4K文本到图像生成的扩散变换器弱到强训练的PyTorch模型定义、预训练权重和推理/采样代码。您可以在我们的项目页面上找到更多可视化效果。

PixArt-Σ: 用于4K文本到图像生成的扩散变换器的弱到强训练
陈俊松*, 葛崇剑*, 谢恩泽*†, 吴越*, 姚乐为, 任晓哲, 王忠道, 骆平, 卢湖川, 李志国
华为诺亚方舟实验室, 大连理工大学, 香港大学, 香港科技大学

欢迎大家踊跃贡献🔥🔥!!

吸取了之前PixArt-α项目的经验，我们将尽量保持这个仓库的简洁，以便PixArt社区的每个人都能使用。

🆚 与PixArt-α比较

模型	T5令牌长度	VAE	2K/4K
PixArt-Σ	300	SDXL	✅
PixArt-α	120	SD1.5	❌
模型	样本-1	样本-2	样本-3
----------	---------------------------------------------------------------------------------------------------------------------------------------------------------------	---------------------------------------------------------------------------------------------------------------------------------------------------------------	---------------------------------------------------------------------------------------------------------------------------------------------------------------
PixArt-Σ
PixArt-α
提示词	特写，灰发，60多岁留胡子的男子，穿羊毛外套和棕色贝雷帽，戴眼镜，观察路人，电影感。	全身照，一位法国女性，摄影，法国街道背景，逆光，轮廓光，富士胶片。	两艘海盗船在一杯咖啡中航行作战的超写实特写视频。

提示词详情

样本-1完整提示词：60多岁灰发留胡子男子的极近特写，他坐在巴黎的一家咖啡馆里，陷入沉思，思考宇宙的历史，他的目光聚焦在画面外走过的人们身上，自己则几乎一动不动，他穿着羊毛外套和衬衫，戴着**棕色贝雷帽**和眼镜，看起来非常像教授，最后他露出一抹微妙的闭嘴微笑，仿佛找到了生命之谜的答案，灯光非常富有电影感，金色的光线和巴黎的街道与城市在背景中，景深，电影感35毫米胶片。

🔧 依赖和安装

Python >= 3.9（推荐使用Anaconda或Miniconda）
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

🔥 如何训练

1. PixArt训练

首先。

我们开始了一个新的仓库，以构建一个更加用户友好和兼容的代码库。主要模型结构与PixArt-α相同，您仍然可以基于原始仓库开发您的功能。 此仓库未来也将支持PixArt-alpha。

[!提示]
现在您可以在无需事先提取特征的情况下训练您的模型。我们重构了PixArt-α代码库中的数据结构，所以每个人都可以在一开始就轻松地训练、推理和可视化。

1.1 下载玩具数据集

首先下载玩具数据集。用于训练的数据集结构如下：

cd ./pixart-sigma-toy-dataset

数据集结构
├──InternImgs/  （图像保存在这里）
│  ├──000000000000.png
│  ├──000000000001.png
│  ├──......
├──InternData/
│  ├──data_info.json    （元数据）
可选（👇）
│  ├──img_sdxl_vae_features_1024resolution_ms_new    （运行tools/extract_caption_feature.py生成标题T5特征，与图像同名，扩展名为.npz）
│  │  ├──000000000000.npy
│  │  ├──000000000001.npy
│  │  ├──......
│  ├──caption_features_new
│  │  ├──000000000000.npz
│  │  ├──000000000001.npz
│  │  ├──......
│  ├──sharegpt4v_caption_features_new    （运行tools/extract_caption_feature.py生成标题T5特征，与图像同名，扩展名为.npz）
│  │  ├──000000000000.npz
│  │  ├──000000000001.npz
│  │  ├──......

1.2 下载预训练检查点

# SDXL-VAE, T5检查点
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检查点
python tools/download.py # 环境变量例如HF_ENDPOINT=https://hf-mirror.com可用于HuggingFace镜像

1.3 您已准备好开始训练！

从配置文件目录中选择您想要的配置文件。

python -m torch.distributed.launch --nproc_per_node=1 --master_port=12345 \
          train_scripts/train.py \
          configs/pixart_sigma_config/PixArt_sigma_xl2_img512_internalms.py \
          --load-from output/pretrained_models/PixArt-Sigma-XL-2-512-MS.pth \
          --work-dir output/your_first_pixart-exp \
          --debug

💻 如何测试

1. 使用Gradio快速开始

首先安装所需的依赖项。确保您已从模型（即将推出）下载检查点文件到output/pretrained_models文件夹，然后在本地机器上运行：

# SDXL-VAE, T5检查点
git lfs install
git clone https://huggingface.co/PixArt-alpha/pixart_sigma_sdxlvae_T5_diffusers output/pixart_sigma_sdxlvae_T5_diffusers

# PixArt-Sigma检查点
python tools/download.py

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

2. 集成到diffusers

[!重要]
升级您的diffusers以使PixArtSigmaPipeline可用！
pip install git+https://github.com/huggingface/diffusers
对于diffusers<0.28.0，请查看此脚本以获取帮助。

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)

# 启用内存优化。
# pipe.enable_model_cpu_offload()

prompt = "撒哈拉沙漠中一个带有开心表情的小仙人掌。"
image = pipe(prompt).images[0]
image.save("./catcus.png")

3. PixArt演示

pip install git+https://github.com/huggingface/diffusers

PixArt-Sigma 1024像素

DEMO_PORT=12345 python app/app_pixart_sigma.py

PixArt-Sigma 单步采样器(DMD)

DEMO_PORT=12345 python app/app_pixart_dmd.py

让我们来看一个使用`http://your-server-ip:12345`的简单示例。

## 4. 将.pth检查点转换为diffusers版本
直接从[Hugging Face](https://huggingface.co/PixArt-alpha/PixArt-Sigma-XL-2-1024-MS)下载

或运行以下命令：
```bash
pip install git+https://github.com/huggingface/diffusers

python tools/convert_pixart_to_diffusers.py --orig_ckpt_path output/pretrained_models/PixArt-Sigma-XL-2-1024-MS.pth --dump_path output/pretrained_models/PixArt-Sigma-XL-2-1024-MS --only_transformer=True --image_size=1024 --version sigma

⏬ 可用模型

所有模型将自动下载此处。您也可以选择从此链接手动下载。

模型	参数数量	检查点路径	在OpenXLab下载
T5 & SDXL-VAE	4.5B	Diffusers: pixart_sigma_sdxlvae_T5_diffusers	即将推出
PixArt-Σ-256	0.6B	pth: PixArt-Sigma-XL-2-256x256.pth Diffusers: PixArt-Sigma-XL-2-256x256	即将推出
PixArt-Σ-512	0.6B	pth: PixArt-Sigma-XL-2-512-MS.pth Diffusers: PixArt-Sigma-XL-2-512-MS	即将推出
PixArt-α-512-DMD	0.6B	Diffusers: PixArt-Alpha-DMD-XL-2-512x512	即将推出
PixArt-Σ-1024	0.6B	pth: PixArt-Sigma-XL-2-1024-MS.pth Diffusers: PixArt-Sigma-XL-2-1024-MS	即将推出
PixArt-Σ-2K	0.6B	pth: PixArt-Sigma-XL-2-2K-MS.pth Diffusers: PixArt-Sigma-XL-2-2K-MS	即将推出

💪待办事项

我们将尽最大努力发布

🤗致谢

感谢PixArt-α、DiT和OpenDMD的出色工作和代码库！
感谢Diffusers提供的出色技术支持和精彩合作！
感谢Hugging Face赞助了精美的演示！

📖引用

@misc{chen2024pixartsigma,
  title={PixArt-\Sigma: Weak-to-Strong Training of Diffusion Transformer for 4K Text-to-Image Generation},
  author={Junsong Chen and Chongjian Ge and Enze Xie and Yue Wu and Lewei Yao and Xiaozhe Ren and Zhongdao Wang and Ping Luo and Huchuan Lu and Zhenguo Li},
  year={2024},
  eprint={2403.04692},
  archivePrefix={arXiv},
  primaryClass={cs.CV}