PixArt-Σ

概述

PixArt-Σ: 用於 4K 文字到影像生成的弱到強擴散 Transformer 訓練由 Junsong Chen、Jincheng Yu、Chongjian Ge、Lewei Yao、Enze Xie、Yue Wu、Zhongdao Wang、James Kwok、Ping Luo、Huchuan Lu 和 Zhenguo Li 撰寫。

關於此管道的一些注意事項

它使用 Transformer 主幹（而不是 UNet）進行去噪。因此，它的架構與 DiT 類似。
它使用從 T5 計算的文字條件進行訓練。這使得該管道更擅長遵循具有複雜細節的複雜文字提示。
它擅長以不同的寬高比生成高解析度影像。為了獲得最佳結果，作者推薦了一些尺寸範圍，可以在此處找到。
它與最先進的文字到影像生成系統（截至本文撰寫之時）的質量相媲美，例如 PixArt-α、Stable Diffusion XL、Playground V2.0 和 DALL-E 3，同時比它們更高效。
它展示了生成超高解析度影像（如 2048px 甚至 4K）的能力。
它表明文字到影像模型可以透過多項改進（VAE、資料集等）從弱模型發展到強模型。

🤗 Optimum 擴充套件了 Diffusers 以支援在第二代 Neuron 裝置（支援 Trainium 和 Inferentia 2）上進行推理。它旨在繼承 Diffusers 在 Neuron 上的易用性。

匯出到 Neuron

要部署 PixArt-Σ 管道中的模型，您需要將它們編譯為針對 AWS Neuron 最佳化的 TorchScript。有四個元件需要匯出為 .neuron 格式以提高效能

文字編碼器
Transformer
VAE 編碼器
VAE 解碼器

您可以透過 CLI 或 NeuronPixArtSigmaPipeline 類編譯並匯出 PixArt-Σ 檢查點。

選項 1：CLI

optimum-cli export neuron --model Jingya/pixart_sigma_pipe_xl_2_512_ms --batch_size 1 --height 512 --width 512 --num_images_per_prompt 1 --torch_dtype bfloat16 --sequence_length 120 pixart_sigma_neuron_512/

我們建議使用 inf2.8xlarge 或更大的例項進行模型編譯。您還可以在僅使用 CPU 的例項（需要約 35 GB 記憶體）上使用 Optimum CLI 編譯模型，然後將預編譯的模型在 inf2.xlarge 上執行以降低成本。在這種情況下，請不要忘記透過新增 --disable-validation 引數來停用推理驗證。

選項 2：Python API

import torch
from optimum.neuron import NeuronPixArtSigmaPipeline

# Compile
compiler_args = {"auto_cast": "none"}
input_shapes = {"batch_size": 1, "height": 512, "width": 512, "sequence_length": 120}

neuron_model = NeuronPixArtSigmaPipeline.from_pretrained("Jingya/pixart_sigma_pipe_xl_2_512_ms", torch_dtype=torch.bfloat16, export=True, disable_neuron_cache=True, **compiler_args, **input_shapes)

# Save locally
neuron_model.save_pretrained("pixart_sigma_neuron_512/")

# Upload to the HuggingFace Hub
neuron_model.push_to_hub(
    "pixart_sigma_neuron_512/", repository_id="optimum/pixart_sigma_pipe_xl_2_512_ms_neuronx"  # Replace with your HF Hub repo id
)

文字到影像

NeuronPixArtSigmaPipeline 類允許您在 Neuron 裝置上根據文字提示生成影像，體驗類似於使用 Diffusers。

使用預編譯的 PixArt-Σ 模型，現在可以在 Neuron 上根據提示生成影像

from optimum.neuron import NeuronPixArtSigmaPipeline

neuron_model = NeuronPixArtSigmaPipeline.from_pretrained("pixart_sigma_neuron_512/")
prompt = "Oppenheimer sits on the beach on a chair, watching a nuclear exposition with a huge mushroom cloud, 120mm."
image = neuron_model(prompt=prompt).images[0]

NeuronPixArtSigmaPipeline

使用 PixArt-Σ 進行文字到影像生成的管道。

class optimum.neuron.NeuronPixArtSigmaPipeline

< 源 >

( **kwargs )

call

< 源 >

( *args **kwargs )

還有其他您希望我們在 🤗Optimum-neuron 中支援的擴散功能嗎？請在 Optimum-neuron Github 倉庫中提交問題或在 HuggingFace 社群論壇上與我們討論，謝謝 🤗！

AWS Trainium 和 Inferentia