該儲存庫包含 PyTorch 模型定義、預訓練權重和推理/採樣程式碼,用於我們探索 4K 文字到圖像生成的擴散變壓器的弱到強訓練的論文。您可以在我們的專案頁面上找到更多視覺化效果。
PixArt-Σ:用於 4K 文字到影像產生的擴散變壓器的弱到強訓練
陳俊松*、葛崇建*、謝恩澤*†、吳悅*、姚樂偉、任小哲、王忠道、羅平、陸虎川、李振國
華為諾亞方舟實驗室、大連理工大學、香港大學、香港科技大學
借鑒之前的 PixArt-α 項目,我們會盡量保持這個 repo 盡可能簡單,以便 PixArt 社區的每個人都可以使用它。
? diffusers
使用補丁的? diffusers
可實現快速體驗!-主要的
-指導
-其他的
模型 | T5令牌長度 | VAE | 2K/4K |
---|---|---|---|
原相-Σ | 300 | SDXL | ✅ |
原相-α | 120 | SD1.5 |
模型 | 樣品1 | 樣品2 | 樣品3 |
---|---|---|---|
原相-Σ | |||
原相-α | |||
迅速的 | 特寫鏡頭,60 多歲的白髮、留著鬍子的男人,穿著羊毛外套和棕色貝雷帽,戴著眼鏡,觀察路人,電影般。 | 身體拍攝,法國女人,攝影,法國街道背景,逆光,邊緣光,富士膠片。 | 兩艘海盜船在一杯咖啡內航行時互相戰鬥的逼真特寫影片。 |
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-α相同,您仍然可以在原始儲存庫的基礎上開發您的功能。另外,這個repo將來會支持PixArt-alpha 。
提示
現在您無需事先提取特徵即可訓練模型。我們對PixArt-α程式碼庫中的資料結構進行了改造,讓每個人都可以從一開始就開始訓練、推理和視覺化,沒有任何痛苦。
首先下載玩具資料集。訓練的資料集結構為:
cd ./pixart-sigma-toy-dataset
Dataset Structure
├──InternImgs/ (images are saved here)
│ ├──000000000000.png
│ ├──000000000001.png
│ ├──......
├──InternData/
│ ├──data_info.json (meta data)
Optional(?)
│ ├──img_sdxl_vae_features_1024resolution_ms_new (run tools/extract_caption_feature.py to generate caption T5 features, same name as images except .npz extension)
│ │ ├──000000000000.npy
│ │ ├──000000000001.npy
│ │ ├──......
│ ├──caption_features_new
│ │ ├──000000000000.npz
│ │ ├──000000000001.npz
│ │ ├──......
│ ├──sharegpt4v_caption_features_new (run tools/extract_caption_feature.py to generate caption T5 features, same name as images except .npz extension)
│ │ ├──000000000000.npz
│ │ ├──000000000001.npz
│ │ ├──......
# 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 # environment eg. HF_ENDPOINT=https://hf-mirror.com can use for HuggingFace mirror
從設定檔目錄中選擇所需的設定檔。
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
首先,先安裝所需的依賴項。確保您已將模型(即將推出)中的檢查點檔案下載到output/pretrained_models
資料夾,然後在本機電腦上執行:
# SDXL-VAE, T5 checkpoints
git lfs install
git clone https://huggingface.co/PixArt-alpha/pixart_sigma_sdxlvae_T5_diffusers output/pixart_sigma_sdxlvae_T5_diffusers
# PixArt-Sigma checkpoints
python tools/download.py
# demo launch
python scripts/interface.py --model_path output/pretrained_models/PixArt-Sigma-XL-2-512-MS.pth --image_size 512 --port 11223
重要的
升級您的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 )
# Enable memory optimizations.
# pipe.enable_model_cpu_offload()
prompt = "A small cactus with a happy face in the Sahara desert."
image = pipe ( prompt ). images [ 0 ]
image . save ( "./catcus.png" )
pip install git+https://github.com/huggingface/diffusers
# PixArt-Sigma 1024px
DEMO_PORT=12345 python app/app_pixart_sigma.py
# PixArt-Sigma One step Sampler(DMD)
DEMO_PORT=12345 python app/app_pixart_dmd.py
讓我們來看一個使用http://your-server-ip:12345
的簡單範例。
直接從抱臉下載
或運行:
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 | 擴散器:pixart_sigma_sdxlvae_T5_diffusers | 即將推出 |
原相-Σ-256 | 0.6B | pth:PixArt-Sigma-XL-2-256x256.pth 擴散器:PixArt-Sigma-XL-2-256x256 | 即將推出 |
原相-Σ-512 | 0.6B | pth:PixArt-Sigma-XL-2-512-MS.pth 擴散器:PixArt-Sigma-XL-2-512-MS | 即將推出 |
原相-α-512-DMD | 0.6B | 擴散器:PixArt-Alpha-DMD-XL-2-512x512 | 即將推出 |
原相-Σ-1024 | 0.6B | pth:PixArt-Sigma-XL-2-1024-MS.pth 擴散器:PixArt-Sigma-XL-2-1024-MS | 即將推出 |
原相-Σ-2K | 0.6B | pth:PixArt-Sigma-XL-2-2K-MS.pth 擴散器:PixArt-Sigma-XL-2-2K-MS | 即將推出 |
我們會盡力發布
@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}