Hugging Face's logo

Diffusers documentation

PixArt-Σ

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.42.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

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。

该论文的摘要如下：

在本文中，我们介绍了 PixArt-Σ，一种能够直接生成 4K 分辨率图像的 Diffusion Transformer 模型（DiT）。PixArt-Σ 相较于其前代模型 PixArt-α 实现了重大突破，提供了保真度显著更高且与文本提示对齐效果更好的图像。PixArt-Σ 的一个关键特性是其训练效率。借助 PixArt-α 的基础预训练，它通过引入更高质量的数据，从“较弱”的基线演变为“更强”的模型，我们将这一过程称为“弱到强训练”。PixArt-Σ 的进步体现在两个方面：(1) 高质量训练数据：PixArt-Σ 采用了更高质量的图像数据，并配以更精确、更详细的图像描述。(2) 高效的 Token 压缩：我们在 DiT 框架中提出了一种新颖的注意力模块，该模块同时压缩键（keys）和值（values），从而显著提升效率并促进超高分辨率图像的生成。得益于这些改进，PixArt-Σ 在图像质量和用户提示遵循能力方面均表现出色，且模型规模显著更小（0.6B 参数），优于现有的文本到图像扩散模型，例如 SDXL（2.6B 参数）和 SD Cascade（5.1B 参数）。此外，PixArt-Σ 生成 4K 图像的能力支持高分辨率海报和壁纸的创作，高效助力电影和游戏等行业的高质量视觉内容生产。

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

关于此流水线的一些说明：

  • 它使用 Transformer 主干网络（而非 UNet）进行去噪。因此，其架构与 DiT 类似。
  • 该模型使用从 T5 计算得出的文本条件进行训练。这一特性使该流水线在遵循包含复杂细节的复杂文本提示方面表现更佳。
  • 它擅长生成不同宽高比的高分辨率图像。为了获得最佳效果，作者推荐了一些尺寸范围，可在此处找到。
  • 其质量可与当前最先进的文本到图像生成系统（截至本文撰写时）相媲美，例如 PixArt-α、Stable Diffusion XL、Playground V2.0 和 DALL-E 3，同时效率更高。
  • 它展示了生成超高分辨率图像的能力，例如 2048px 甚至 4K。
  • 这表明，通过多项改进（如 VAE、数据集等），文本到图像模型可以从一个较弱的模型逐步演进为更强大的模型。

请务必参阅 调度器指南，了解如何探索调度器速度与质量之间的权衡，并查看 跨流水线复用组件 部分，了解如何高效地将相同组件加载到多个流水线中。

你可以通过将 PixArtSigmaPipeline 生成的图像传递给 SDXL refiner 模型，进一步提升生成质量。

在 GPU 显存低于 8GB 的情况下进行推理

通过以 8 位精度加载文本编码器，可以在 8GB 以下的 GPU 显存中运行 PixArtSigmaPipeline。下面通过一个完整的示例来说明。

首先，安装 bitsandbytes 库：

pip install -U bitsandbytes

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

from transformers import T5EncoderModel
from diffusers import PixArtSigmaPipeline
import torch

text_encoder = T5EncoderModel.from_pretrained(
    "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS",
    subfolder="text_encoder",
    load_in_8bit=True,
    device_map="auto",
)
pipe = PixArtSigmaPipeline.from_pretrained(
    "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS",
    text_encoder=text_encoder,
    transformer=None,
    device_map="balanced"
)

现在，使用 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 = PixArtSigmaPipeline.from_pretrained(
    "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS",
    text_encoder=None,
    dtype=torch.float16,
).to("cuda")  # or "mps", "xpu", "cpu"

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 显存中运行 PixArtSigmaPipeline。

如果您想要一份内存使用情况的报告，请运行此 脚本。

以 8 位精度计算的文本嵌入可能会因精度降低导致表示空间中的信息丢失，从而影响生成图像的质量。建议对比使用和不使用 8 位精度时的输出结果。

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

PixArtSigmaPipeline

class diffusers.PixArtSigmaPipeline

< >

( tokenizer: T5Tokenizertext_encoder: T5EncoderModelvae: AutoencoderKLtransformer: PixArtTransformer2DModelscheduler: KarrasDiffusionSchedulers )

参数

  • vae (AutoencoderKL) — 变分自编码器（VAE）模型，用于将图像编码为潜在表示，以及从潜在表示解码回图像。
  • text_encoder (T5EncoderModel) — 冻结的文本编码器。PixArt-Alpha 使用 T5，具体为 t5-v1_1-xxl 变体。
  • tokenizer (T5Tokenizer) — T5Tokenizer 类的分词器。
  • transformer (PixArtTransformer2DModel) — 一个以文本为条件的 PixArtTransformer2DModel，用于对编码后的图像潜变量进行去噪。在配置中最初发布为 Transformer2DModel ，但此不一致可以忽略。
  • scheduler (SchedulerMixin) — 用于与 transformer 配合，对编码后的图像潜变量进行去噪的调度器。

使用 PixArt-Sigma 进行文本到图像生成的 Pipeline。

该模型继承自 DiffusionPipeline。有关该库为所有流水线实现的通用方法（例如下载或保存、在特定设备上运行等），请参阅父类文档。

__call__

< >

( prompt: str | list[str] = Nonenegative_prompt: str = ''num_inference_steps: int = 20timesteps: list = Nonesigmas: list = Noneguidance_scale: float = 4.5num_images_per_prompt: int | None = 1height: int | None = Nonewidth: int | None = Noneeta: float = 0.0generator: typing.Union[torch.Generator, list[torch.Generator], NoneType] = Nonelatents: typing.Optional[torch.Tensor] = Noneprompt_embeds: typing.Optional[torch.Tensor] = Noneprompt_attention_mask: typing.Optional[torch.Tensor] = Nonenegative_prompt_embeds: typing.Optional[torch.Tensor] = Nonenegative_prompt_attention_mask: typing.Optional[torch.Tensor] = Noneoutput_type: str | None = 'pil'return_dict: bool = Truecallback: typing.Optional[typing.Callable[[int, int, torch.Tensor], NoneType]] = Nonecallback_steps: int = 1clean_caption: bool = Trueuse_resolution_binning: bool = Truemax_sequence_length: int = 300**kwargs ) → ImagePipelineOutput or tuple

参数

  • prompt (str 或 list[str]，可选) — 用于引导图像生成的提示词或提示词列表。如果未定义，则必须改为传入 prompt_embeds。
  • negative_prompt (str 或 list[str]，可选) — 用于不引导图像生成的提示词。如果未定义，则必须传入 negative_prompt_embeds。在未使用引导时会被忽略（即当 guidance_scale 小于 1 时忽略）。
  • num_inference_steps (int，可选，默认为 100) — 去噪步数。更多的去噪步数通常能生成更高质量的图像，但代价是推理速度变慢。
  • timesteps (list[int], 可选) — 用于支持 timesteps 参数的调度器去噪过程的自定义时间步，该参数位于其 set_timesteps 方法中。如果未定义，则使用传入 num_inference_steps 时的默认行为。必须按降序排列。
  • sigmas (list[float]，可选) — 用于在支持 sigmas 参数的调度器的 set_timesteps 方法中，为去噪过程指定自定义的 sigmas。如果未定义，则使用传入 num_inference_steps 时的默认行为。
  • guidance_scale (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，对其他调度器将被忽略。
  • generator (torch.Generator 或 list[torch.Generator]，可选) — 一个或一组 torch 生成器，用于使生成过程具有确定性。
  • latents (torch.Tensor, 可选) — 预先生成的噪声潜变量，从高斯分布中采样得到，用作图像生成的输入。可用于使用不同的提示词对同一生成过程进行微调。如果未提供，将通过使用提供的随机generator进行采样来生成一个潜变量张量。
  • prompt_embeds (torch.Tensor, 可选) — 预生成的文本嵌入。可用于轻松调整文本输入，例如 提示词加权。如果未提供，文本嵌入将从 prompt 输入参数生成。
  • prompt_attention_mask (torch.Tensor，可选) — 为文本嵌入预生成的注意力掩码。
  • negative_prompt_embeds (torch.Tensor, 可选) — 预生成的负向文本嵌入。对于 PixArt-Sigma，此负向提示应为 ""。如果未提供，negative_prompt_embeds 将根据 negative_prompt 输入参数生成。
  • negative_prompt_attention_mask (torch.Tensor, 可选) — 为负向文本嵌入预生成的注意力掩码。
  • output_type (str，可选，默认为 "pil") — 生成图像的输出格式。可选值为 PIL：PIL.Image.Image 或 np.array。
  • return_dict (bool，可选，默认为 True) — 是否返回 ~pipelines.stable_diffusion.IFPipelineOutput 而不是普通元组。
  • callback (Callable, 可选) — 一个函数，将在推理过程中每 callback_steps 步被调用一次。该函数将以以下参数被调用：callback(step: int, timestep: int, latents: torch.Tensor)。
  • callback_steps (int，可选，默认为 1) — callback 函数的调用频率。如果未指定，则回调将在每一步被调用。
  • clean_caption (bool，可选，默认为 True) — 是否在创建嵌入之前清理字幕。需要安装 beautifulsoup4 和 ftfy。如果未安装这些依赖项，嵌入将直接从原始提示创建。
  • use_resolution_binning (bool，默认为 True) — 如果设置为 True，请求的高度和宽度会先通过 ASPECT_RATIO_1024_BIN 映射到最接近的分辨率。生成的潜变量解码为图像后，会被重新调整回请求的分辨率。适用于生成非正方形图像。
  • max_sequence_length（int，默认为 300）— 与 prompt 配合使用的最大序列长度。

返回值

ImagePipelineOutput 或 tuple

如果 return_dict 为 True，则返回 ImagePipelineOutput；否则返回一个 tuple，其中第一个元素是包含生成图像的列表

调用生成 pipeline 时调用的函数。

示例：

>>> import torch
>>> from diffusers import PixArtSigmaPipeline

>>> # You can replace the checkpoint id with "PixArt-alpha/PixArt-Sigma-XL-2-512-MS" too.
>>> pipe = PixArtSigmaPipeline.from_pretrained(
...     "PixArt-alpha/PixArt-Sigma-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: str | list[str]do_classifier_free_guidance: bool = Truenegative_prompt: str = ''num_images_per_prompt: int = 1device: typing.Optional[torch.device] = Noneprompt_embeds: typing.Optional[torch.Tensor] = Nonenegative_prompt_embeds: typing.Optional[torch.Tensor] = Noneprompt_attention_mask: typing.Optional[torch.Tensor] = Nonenegative_prompt_attention_mask: typing.Optional[torch.Tensor] = Noneclean_caption: bool = Falsemax_sequence_length: int = 300**kwargs )

参数

  • prompt (str 或 list[str]，可选) — 要编码的提示
  • negative_prompt (str 或 list[str]，可选) — 用于不引导图像生成的提示词。如果未定义，则必须改为传入 negative_prompt_embeds。在未使用引导时会被忽略（即当 guidance_scale 小于 1 时会被忽略）。对于 PixArt-Alpha，此值应为 ""。
  • do_classifier_free_guidance (bool，可选，默认为 True) — 是否使用 classifier free guidance
  • 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，该函数将在编码之前对提供的 caption 进行预处理和清理。
  • max_sequence_length (int，默认为 300) — 用于提示词的最大序列长度。

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

Update on GitHub

←PixArt-α PRX→