pixart-alpha-pytorch以文生图AIGC算法模型

PixArt-alpha

论文

PixArt-α: Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis

• https://arxiv.org/abs/2310.00426

模型结构

该模型基于DiT（Diffusion Transformer）模型，添加了Multi-Head Cross-Attention用于对其文本与图像。

算法原理

模型中主要涉及Multi-Head Self-Attention和Multi-Head Cross-Attention，其中Multi-Head Self-Attention主要用于对图像建模，Multi-Head Cross-Attention用于对齐图像与文本。

环境配置

Docker（方法一）

docker pull image.sourcefind.cn:5000/dcu/admin/base/pytorch:2.1.0-centos7.6-dtk23.10.1-py38

docker run --shm-size 10g --network=host --name=opensora --privileged --device=/dev/kfd --device=/dev/dri --group-add video --cap-add=SYS_PTRACE --security-opt seccomp=unconfined -v 项目地址(绝对路径):/home/ -v /opt/hyhal:/opt/hyhal:ro -it <your IMAGE ID> bash

pip install flash_attn-2.0.4_torch2.1_dtk2310-cp38-cp38-linux_x86_64.whl  （whl.zip文件中）

pip install triton-2.1.0%2Bgit34f8189.abi0.dtk2310-cp38-cp38-manylinux2014_x86_64.whl （开发者社区下载）

cd xformers && pip install xformers==0.0.23 --no-deps && bash patch_xformers.rocm.sh  （whl.zip文件中）

pip install -r requirements.txt

pip install timm --no-deps

pip uninstall apex

# 安装diffusers
# 手动安装
git clone https://github.com/huggingface/diffusers.git

cd diffusers && python setup.py install

# 自动安装
pip install git+https://github.com/huggingface/diffusers

Dockerfile（方法二）

# 需要在对应的目录下
docker build -t <IMAGE_NAME>:<TAG> .

docker run --shm-size 10g --network=host --name=opensora --privileged --device=/dev/kfd --device=/dev/dri --group-add video --cap-add=SYS_PTRACE --security-opt seccomp=unconfined -v 项目地址(绝对路径):/home/ -v /opt/hyhal:/opt/hyhal:ro -it <your IMAGE ID> bash

pip install flash_attn-2.0.4_torch2.1_dtk2310-cp38-cp38-linux_x86_64.whl  （whl.zip文件中）

pip install triton-2.1.0%2Bgit34f8189.abi0.dtk2310-cp38-cp38-manylinux2014_x86_64.whl （开发者社区下载）

cd xformers && pip install xformers==0.0.23 --no-deps && bash patch_xformers.rocm.sh  （whl.zip文件中）

pip install -r requirements.txt

pip install timm --no-deps

cd diffusion && python setup.py install

pip uninstall apex

# 安装diffusers
# 手动安装
git clone https://github.com/huggingface/diffusers.git

cd diffusers && python setup.py install

# 自动安装
pip install git+https://github.com/huggingface/diffusers

Anaconda (方法三)

1、关于本项目DCU显卡所需的特殊深度学习库可从光合开发者社区下载安装： https://developer.hpccube.com/tool/

DTK驱动：dtk23.10.1
python：python3.8
torch:2.1.0
torchvision:0.16.0
triton:2.1.0

Tips：以上dtk驱动、python、torch等DCU相关工具版本需要严格一一对应

2、其它非特殊库参照requirements.txt安装

pip install flash_attn-2.0.4_torch2.1_dtk2310-cp38-cp38-linux_x86_64.whl  （whl.zip文件中）

cd xformers && pip install xformers==0.0.23 --no-deps && bash patch_xformers.rocm.sh  （whl.zip文件中）

pip install -r requirements.txt

pip install timm --no-deps

# 安装diffusers
# 手动安装
git clone https://github.com/huggingface/diffusers.git

cd diffusers && python setup.py install

# 自动安装
pip install git+https://github.com/huggingface/diffusers

数据集

注意：该数据集为训练数据集

完整数据：https://ai.meta.com/datasets/segment-anything/

测试数据：https://huggingface.co/datasets/PixArt-alpha/data_toy

数据下载完成后需要进行处理，可运行以下脚本：

# 使用LLava获取更加详细的图像描述
python tools/VLM_caption_lightning.py --output output/dir/ --data-root data/root/path --index path/to/data.json

# 提前生成训练需要的特征
python tools/extract_features.py --img_size=256 \
--json_path "data/data_toy/data_info.json" \
--t5_save_root "data/data_toy/caption_feature_wmask" \
--vae_save_root "data/data_toy/img_vae_features" \
--pretrained_models_dir "pretrained_models/hub/pixart_alpha" \
--dataset_root "data/data_toy/images/"

处理后获得下述数据结构

data/
└── data_toy
    ├── caption_feature_wmask
    │   ├── 0_1.npz
    │   └── 0_3.npz
    ├── captions
    │   ├── 0_1.txt
    │   └── 0_3.txt
    ├── data_info.json
    ├── images
    │   ├── 0_1.png
    │   └── 0_3.png
    ├── img_vae_features
    │   └── 256resolution
    │       └── noflip
    │           ├── 0_1.npy
    │           └── 0_3.npy
    └── partition
        └── part0.txt

训练

敬请期待!

推理

模型下载

Model+url	存放位置
T5	/path/to/save/models/pixart_alpha/t5_ckpts
sd-vae-ft-ema	/path/to/save/models/pixart_alpha/sd-vae-ft-ema

pixart_alpha/
├── sd-vae-ft-ema
│   ├── config.json
│   └── diffusion_pytorch_model.bin
└── t5_ckpts
    └── t5-v1_1-xxl
        ├── config.json
        ├── pytorch_model-00001-of-00002.bin
        ├── pytorch_model-00002-of-00002.bin
        ├── pytorch_model.bin.index.json
        ├── special_tokens_map.json
        ├── spiece.model
        └── tokenizer_config.json

注意：上述模型需手动下载，其余模型将在运行时自动下载。

export HF_ENDPOINT=https://hf-mirror.com
export HUB_HOME=/path/to/save/models

命令

# 快速测试
HIP_VISIBLE_DEVICES=0 python quick_inference_with_code.py <prompt>

WebUI

# diffusers version
DEMO_PORT=12345 python app/app.py

result

	prompt	output
	a dog is playing a basketball

精度

无

应用场景

算法类别

AIGC

热点应用行业

零售,广媒,教育

源码仓库及问题反馈

• ModelZoo / PixArt-alpha_pytorch · GitLab

参考资料

• GitHub - PixArt-alpha/PixArt-alpha: PixArt-α: Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis