pixart_dmd.md

Summary

We combine PixArt-α and DMD to achieve one step image generation. This document will guide you how to train and test.

Important

Due to the difference between the DiT & Stable Diffusion,

We find that the setting of the start timestep of the student model is very important for DMD training.

refer to PixArt-Sigma paper's Supplementary for more details: https://arxiv.org/abs/2403.04692

---

How to Train

1. Environment

Refer to the PixArt-Sigma Environment Here. To use FSDP, you need to make sure:

• PyTorch >= 2.0.1+cu11.7

2. Data preparation (Image-Noise pairs)

Generate image-noise pairs

python tools/generate_dmd_data_noise_pairs.py  \
        --pipeline_load_from=PixArt-alpha/PixArt-XL-2-512x512 \
        --model_path=PixArt-alpha/PixArt-XL-2-512x512 \
        --save_img # (optinal)

Extract features in advance (Ignore if you already have)

python tools/extract_features.py --run_t5_feature_extract \
                                 --t5_models_dir=PixArt-alpha/PixArt-XL-2-512x512 \
                                 --t5_save_root=pixart-sigma-toy-dataset/InternData \
                                 --caption_label=prompt \
                                 --t5_json_path=pixart-sigma-toy-dataset/InternData/data_info.json

3. 🔥 Run

bash train_scripts/train_dmd.sh

---

How to Test

PixArt-DMD Demo

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

# PixArt-Sigma One step Sampler(DMD)
DEMO_PORT=12345 python app/app_pixart_dmd.py

Let's have a look at a simple example using the http://your-server-ip:12345.

---

Samples