Diffusers 文档

PixArt-α

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

PixArt-α

PixArt-α:快速训练扩散Transformer以实现逼真文本到图像合成由Junsong Chen、Jincheng Yu、Chongjian Ge、Lewei Yao、Enze Xie、Yue Wu、Zhongdao Wang、James Kwok、Ping Luo、Huchuan Lu 和 Zhenguo Li 完成。

论文摘要如下:

最先进的文本到图像 (T2I) 模型需要巨大的训练成本(例如,数百万 GPU 小时),严重阻碍了 AIGC 社区的基础创新,同时增加了二氧化碳排放。本文介绍了 PIXART-α,一个基于 Transformer 的 T2I 扩散模型,其图像生成质量与最先进的图像生成器(例如 Imagen、SDXL 甚至 Midjourney)具有竞争力,达到了接近商业应用的水平。此外,它支持高达 1024 像素的高分辨率图像合成,且训练成本低,如图 1 和图 2 所示。为实现这一目标,提出了三个核心设计:(1)训练策略分解:我们设计了三个不同的训练步骤,分别优化像素依赖性、文本图像对齐和图像美学质量;(2)高效 T2I Transformer:我们将交叉注意力模块整合到扩散 Transformer (DiT) 中,以注入文本条件并简化计算密集型类别条件分支;(3)高信息量数据:我们强调文本图像对中概念密度的重要性,并利用大型视觉语言模型自动标记密集伪字幕以辅助文本图像对齐学习。结果,PIXART-α 的训练速度明显超过现有的大规模 T2I 模型,例如 PIXART-α 仅需 Stable Diffusion v1.5 训练时间的 10.8%(675 对比 6,250 A100 GPU 天),节省了近 30 万美元(26,000 美元对比 320,000 美元),并减少了 90% 的二氧化碳排放。此外,与更大的 SOTA 模型 RAPHAEL 相比,我们的训练成本仅为 1%。大量实验表明,PIXART-α 在图像质量、艺术性和语义控制方面表现出色。我们希望 PIXART-α 能为 AIGC 社区和初创公司提供新思路,以加速从零开始构建自己的高质量、低成本生成模型。

您可以在 PixArt-alpha/PixArt-alpha 找到原始代码库,在 PixArt-alpha 找到所有可用的检查点。

关于此管道的一些注意事项

  • 它使用 Transformer 主干(而不是 UNet)进行去噪。因此,它的架构与 DiT 相似。
  • 它使用从 T5 计算的文本条件进行训练。这使得该管道更擅长遵循具有复杂细节的复杂文本提示。
  • 它擅长生成不同纵横比的高分辨率图像。为了获得最佳结果,作者推荐了一些尺寸范围,可以在这里找到。
  • 它在图像生成质量方面可与当前最先进的文本到图像生成系统(截至撰写本文时)如 Stable Diffusion XL、Imagen 和 DALL-E 2 媲美,同时效率更高。

请务必查阅“调度器”指南,了解如何权衡调度器速度和质量,并参阅“跨管道复用组件”部分,了解如何高效地将相同组件加载到多个管道中。

使用低于 8GB GPU 显存进行推理

通过以 8 位精度加载文本编码器,使用低于 8GB GPU 显存运行 PixArtAlphaPipeline。让我们来看一个完整的例子。

首先,安装 bitsandbytes

pip install -U bitsandbytes

然后以 8 位加载文本编码器

from transformers import T5EncoderModel
from diffusers import PixArtAlphaPipeline
import torch

text_encoder = T5EncoderModel.from_pretrained(
    "PixArt-alpha/PixArt-XL-2-1024-MS",
    subfolder="text_encoder",
    load_in_8bit=True,
    device_map="auto",

)
pipe = PixArtAlphaPipeline.from_pretrained(
    "PixArt-alpha/PixArt-XL-2-1024-MS",
    text_encoder=text_encoder,
    transformer=None,
    device_map="auto"
)

现在,使用 `pipe` 编码提示

with torch.no_grad():
    prompt = "cute cat"
    prompt_embeds, prompt_attention_mask, negative_embeds, negative_prompt_attention_mask = pipe.encode_prompt(prompt)

由于文本嵌入已计算,将 `text_encoder` 和 `pipe` 从内存中删除,并释放一些 GPU 显存

import gc

def flush():
    gc.collect()
    torch.cuda.empty_cache()

del text_encoder
del pipe
flush()

然后以提示嵌入作为输入计算潜在变量

pipe = PixArtAlphaPipeline.from_pretrained(
    "PixArt-alpha/PixArt-XL-2-1024-MS",
    text_encoder=None,
    torch_dtype=torch.float16,
).to("cuda")

latents = pipe(
    negative_prompt=None,
    prompt_embeds=prompt_embeds,
    negative_prompt_embeds=negative_embeds,
    prompt_attention_mask=prompt_attention_mask,
    negative_prompt_attention_mask=negative_prompt_attention_mask,
    num_images_per_prompt=1,
    output_type="latent",
).images

del pipe.transformer
flush()

请注意,在初始化 `pipe` 时,您将 `text_encoder` 设置为 `None`,这样它就不会被加载。

潜在变量计算完成后,将其传递给 VAE 进行解码,生成真实图像

with torch.no_grad():
    image = pipe.vae.decode(latents / pipe.vae.config.scaling_factor, return_dict=False)[0]
image = pipe.image_processor.postprocess(image, output_type="pil")[0]
image.save("cat.png")

通过删除您不使用的组件并刷新 GPU 显存,您应该能够以低于 8GB GPU 显存运行 PixArtAlphaPipeline

如果您需要内存使用情况报告,请运行此脚本

以 8 位计算文本嵌入可能会影响生成图像的质量,因为精度降低导致表示空间中的信息丢失。建议比较 8 位和非 8 位下的输出。

加载 `text_encoder` 时,您将其 `load_in_8bit` 设置为 `True`。您也可以指定 `load_in_4bit`,以将内存需求进一步降低到 7GB 以下。

PixArtAlphaPipeline

diffusers.PixArtAlphaPipeline

< >

( tokenizer: T5Tokenizer text_encoder: T5EncoderModel vae: AutoencoderKL transformer: PixArtTransformer2DModel scheduler: DPMSolverMultistepScheduler )

参数

  • vae (AutoencoderKL) — 用于将图像编码和解码为潜在表示的变分自编码器 (VAE) 模型。
  • text_encoder (T5EncoderModel) — 冻结的文本编码器。PixArt-Alpha 使用 T5,特别是 t5-v1_1-xxl 变体。
  • 分词器 (T5Tokenizer) — T5Tokenizer 类的分词器。
  • transformer (PixArtTransformer2DModel) — 一个文本条件 PixArtTransformer2DModel,用于对编码的图像潜在变量进行去噪。
  • 调度器 (SchedulerMixin) — 用于与 transformer 结合对编码图像潜在变量进行去噪的调度器。

用于使用 PixArt-Alpha 生成文本到图像的管道。

此模型继承自 DiffusionPipeline。请查看超类文档,了解该库为所有管道实现的通用方法(例如下载或保存、在特定设备上运行等)。

__call__

< >

( prompt: typing.Union[str, typing.List[str]] = None negative_prompt: str = '' num_inference_steps: int = 20 timesteps: typing.List[int] = None sigmas: typing.List[float] = None guidance_scale: float = 4.5 num_images_per_prompt: typing.Optional[int] = 1 height: typing.Optional[int] = None width: typing.Optional[int] = None eta: float = 0.0 generator: typing.Union[torch._C.Generator, typing.List[torch._C.Generator], NoneType] = None latents: typing.Optional[torch.Tensor] = None prompt_embeds: typing.Optional[torch.Tensor] = None prompt_attention_mask: typing.Optional[torch.Tensor] = None negative_prompt_embeds: typing.Optional[torch.Tensor] = None negative_prompt_attention_mask: typing.Optional[torch.Tensor] = None output_type: typing.Optional[str] = 'pil' return_dict: bool = True callback: typing.Optional[typing.Callable[[int, int, torch.Tensor], NoneType]] = None callback_steps: int = 1 clean_caption: bool = True use_resolution_binning: bool = True max_sequence_length: int = 120 **kwargs ) ImagePipelineOutputtuple

参数

  • 提示 (strList[str], 可选) — 用于引导图像生成的提示。如果未定义,则必须传入 prompt_embeds
  • 负面提示 (strList[str], 可选) — 不用于引导图像生成的提示。如果未定义,则必须传入 negative_prompt_embeds。当不使用引导时(即,如果 guidance_scale 小于 1 时被忽略),则忽略。
  • num_inference_steps (int, 可选, 默认为 100) — 去噪步数。更多去噪步数通常会带来更高质量的图像,但推理速度会变慢。
  • timesteps (List[int], 可选) — 用于去噪过程的自定义时间步,适用于其 set_timesteps 方法支持 timesteps 参数的调度器。如果未定义,将使用传递 num_inference_steps 时的默认行为。必须按降序排列。
  • sigmas (List[float], 可选) — 用于去噪过程的自定义 sigmas,适用于其 set_timesteps 方法支持 sigmas 参数的调度器。如果未定义,将使用传递 num_inference_steps 时的默认行为。
  • 引导比例 (float, 可选, 默认为 4.5) — Classifier-Free Diffusion Guidance 中定义的引导比例。guidance_scale 定义为 Imagen Paper 方程 2 中的 w。通过将 guidance_scale > 1 启用引导比例。更高的引导比例鼓励生成与文本 prompt 紧密相关的图像,通常以牺牲较低图像质量为代价。
  • num_images_per_prompt (int, 可选, 默认为 1) — 每个提示要生成的图像数量。
  • height (int, 可选, 默认为 self.unet.config.sample_size) — 生成图像的高度(像素)。
  • width (int, 可选, 默认为 self.unet.config.sample_size) — 生成图像的宽度(像素)。
  • eta (float, 可选, 默认为 0.0) — 对应于 DDIM 论文中的参数 eta (η): https://huggingface.co/papers/2010.02502。仅适用于 schedulers.DDIMScheduler,对其他调度器将被忽略。
  • 生成器 (torch.GeneratorList[torch.Generator], 可选) — 一个或多个 torch 生成器,用于使生成确定性。
  • 潜在变量 (torch.Tensor, 可选) — 预生成的噪声潜在变量,从高斯分布中采样,用作图像生成的输入。可用于通过不同提示调整相同生成。如果未提供,将通过使用提供的随机 generator 采样生成一个潜在变量张量。
  • prompt_embeds (torch.Tensor, 可选) — 预生成的文本嵌入。可用于轻松调整文本输入,例如 提示加权。如果未提供,将从 prompt 输入参数生成文本嵌入。
  • prompt_attention_mask (torch.Tensor, 可选) — 预生成的文本嵌入的注意力掩码。
  • negative_prompt_embeds (torch.Tensor, 可选) — 预生成的负面文本嵌入。对于 PixArt-Alpha,此负面提示应为 ""。如果未提供,负面提示嵌入将从 negative_prompt 输入参数生成。
  • negative_prompt_attention_mask (torch.Tensor, 可选) — 预生成的负面文本嵌入的注意力掩码。
  • output_type (str, 可选, 默认为 "pil") — 生成图像的输出格式。选择 PIL: PIL.Image.Imagenp.array
  • return_dict (bool, 可选, 默认为 True) — 是否返回 ~pipelines.stable_diffusion.IFPipelineOutput 而非普通元组。
  • 回调 (Callable, 可选) — 在推理过程中每 callback_steps 步都会调用的函数。该函数将使用以下参数调用:callback(step: int, timestep: int, latents: torch.Tensor)
  • callback_steps (int, 可选, 默认为 1) — 调用 callback 函数的频率。如果未指定,则在每个步骤都会调用回调函数。
  • clean_caption (bool, 可选, 默认为 True) — 是否在创建嵌入之前清理标题。需要安装 beautifulsoup4ftfy。如果未安装依赖项,则将从原始提示创建嵌入。
  • use_resolution_binning (bool 默认为 True) — 如果设置为 True,则首先使用 ASPECT_RATIO_1024_BIN 将请求的高度和宽度映射到最接近的分辨率。生成的潜在图像解码为图像后,会将其大小调整回请求的分辨率。对于生成非方形图像非常有用。
  • max_sequence_length (int 默认为 120) — 与 prompt 一起使用的最大序列长度。

返回

ImagePipelineOutputtuple

如果 return_dictTrue,则返回 ImagePipelineOutput,否则返回一个 tuple,其中第一个元素是生成的图像列表

调用管道进行生成时调用的函数。

示例

>>> import torch
>>> from diffusers import PixArtAlphaPipeline

>>> # You can replace the checkpoint id with "PixArt-alpha/PixArt-XL-2-512x512" too.
>>> pipe = PixArtAlphaPipeline.from_pretrained("PixArt-alpha/PixArt-XL-2-1024-MS", torch_dtype=torch.float16)
>>> # 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]

encode_prompt

< >

( prompt: typing.Union[str, typing.List[str]] do_classifier_free_guidance: bool = True negative_prompt: str = '' num_images_per_prompt: int = 1 device: typing.Optional[torch.device] = None prompt_embeds: typing.Optional[torch.Tensor] = None negative_prompt_embeds: typing.Optional[torch.Tensor] = None prompt_attention_mask: typing.Optional[torch.Tensor] = None negative_prompt_attention_mask: typing.Optional[torch.Tensor] = None clean_caption: bool = False max_sequence_length: int = 120 **kwargs )

参数

  • prompt (strList[str], 可选) — 待编码的提示词
  • negative_prompt (strList[str], 可选) — 不用于引导图像生成的提示词。如果未定义,则必须传递 negative_prompt_embeds。在使用无分类器引导时(即,如果 guidance_scale 小于 1),此参数将被忽略。对于 PixArt-Alpha,此参数应为“”。
  • do_classifier_free_guidance (bool, 可选, 默认为 True) — 是否使用无分类器引导
  • num_images_per_prompt (int, 可选, 默认为 1) — 每个提示词应生成的图像数量
  • device — (torch.device, 可选): 放置结果嵌入的 torch 设备
  • prompt_embeds (torch.Tensor, 可选) — 预先生成的文本嵌入。可用于轻松调整文本输入,例如提示词权重。如果未提供,将从 prompt 输入参数生成文本嵌入。
  • negative_prompt_embeds (torch.Tensor, 可选) — 预先生成的负面文本嵌入。对于 PixArt-Alpha,它应该是“ ”字符串的嵌入。
  • clean_caption (bool, 默认为 False) — 如果为 True,函数将在编码前对提供的标题进行预处理和清理。
  • max_sequence_length (int, 默认为 120) — 用于提示词的最大序列长度。

将提示编码为文本编码器隐藏状态。

< > 在 GitHub 上更新