AWS Trainium & Inferentia 文档

PixArt-Σ

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

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]
PixArt-Σ generated image.

NeuronPixArtSigmaPipeline

使用 PixArt-Σ 进行文本到图像生成的管道。

class optimum.neuron.NeuronPixArtSigmaPipeline

< >

( **kwargs )

__call__

< >

( *args **kwargs )

还有其他您希望我们在 🤗Optimum-neuron 中支持的扩散功能吗?请在 Optimum-neuron Github 仓库中提交问题或在 HuggingFace 社区论坛上与我们讨论,谢谢 🤗!