PixArt-alpha

论文

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

https://arxiv.org/abs/2310.00426

模型结构

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

算法原理

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

环境配置

Docker（方法一）

<code>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

训练

敬请期待!

推理

模型下载

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

精度

无

应用场景

算法类别

<code>AIGC

热点应用行业

零售,广媒,教育

源码仓库及问题反馈

ModelZoo / PixArt-alpha_pytorch · GitLab

参考资料

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