Diffusers 文档
PixArt-Σ
并获取增强的文档体验
开始
PixArt-Σ
PixArt-Σ: Weak-to-Strong Training of Diffusion Transformer for 4K Text-to-Image Generation 是 Junsong Chen, Jincheng Yu, Chongjian Ge, Lewei Yao, Enze Xie, Yue Wu, Zhongdao Wang, James Kwok, Ping Luo, Huchuan Lu, 和 Zhenguo Li。
论文摘要如下
在本文中,我们介绍了 PixArt-Σ,一种能够直接生成 4K 分辨率图像的扩散 Transformer 模型 (DiT)。PixArt-Σ 代表了对其前身 PixArt-α 的重大改进,提供了更高保真度的图像,并改进了与文本提示的对齐。PixArt-Σ 的一个关键特性是其训练效率。它利用 PixArt-α 的基础预训练,通过整合更高质量的数据,从“较弱”的基线演变为“较强”的模型,我们将此过程称为“弱到强训练”。PixArt-Σ 的改进是双重的:(1) 高质量训练数据:PixArt-Σ 整合了质量更高的图像数据,并配以更精确和详细的图像标题。(2) 高效令牌压缩:我们在 DiT 框架内提出了一种新颖的注意力模块,该模块压缩了键和值,显着提高了效率并促进了超高分辨率图像的生成。由于这些改进,PixArt-Σ 以明显更小的模型尺寸(0.6B 参数)实现了优异的图像质量和用户提示遵循能力,优于现有的文本到图像扩散模型,例如 SDXL(2.6B 参数)和 SD Cascade(5.1B 参数)。此外,PixArt-Σ 生成 4K 图像的能力支持创建高分辨率海报和壁纸,有效地促进了电影和游戏等行业中高质量视觉内容的制作。
你可以在 PixArt-alpha/PixArt-sigma 找到原始代码库,并在 PixArt-alpha 找到所有可用的检查点。
关于此 pipeline 的一些注意事项
- 它使用 Transformer 主干(而不是 UNet)进行去噪。因此,它具有与 DiT 相似的架构。
- 它使用从 T5 计算出的文本条件进行训练。这使得该 pipeline 更擅长遵循带有复杂细节的复杂文本提示。
- 它擅长生成不同宽高比的高分辨率图像。为了获得最佳效果,作者建议使用一些尺寸范围,可以在 这里 找到。
- 它的质量可与最先进的文本到图像生成系统(截至撰写本文时)相媲美,例如 PixArt-α、Stable Diffusion XL、Playground V2.0 和 DALL-E 3,同时比它们更高效。
- 它展示了生成超高分辨率图像的能力,例如 2048px 甚至 4K。
- 它表明,文本到图像模型可以通过多项改进(VAEs、数据集等)从弱模型成长为更强的模型。
请务必查看 Schedulers 指南,了解如何探索 scheduler 速度和质量之间的权衡,并查看跨 pipelines 重用组件部分 ,了解如何有效地将相同组件加载到多个 pipelines 中。
你可以通过将 PixArtSigmaPipeline 生成的图像传递给 SDXL refiner 模型来进一步提高生成质量。
在 8GB GPU VRAM 以下进行推理
通过以 8 位精度加载文本编码器,在 8GB GPU VRAM 以下运行 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
编码 prompt
with torch.no_grad():
prompt = "cute cat"
prompt_embeds, prompt_attention_mask, negative_embeds, negative_prompt_attention_mask = pipe.encode_prompt(prompt)
由于文本嵌入 (text embeddings) 已经计算完成,请从内存中移除 text_encoder
和 pipe
,并释放一些 GPU 显存 (VRAM)。
import gc
def flush():
gc.collect()
torch.cuda.empty_cache()
del text_encoder
del pipe
flush()
然后使用提示词嵌入 (prompt embeddings) 作为输入来计算潜在变量 (latents)。
pipe = PixArtSigmaPipeline.from_pretrained(
"PixArt-alpha/PixArt-Sigma-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
,这样就不会加载它。
一旦计算出潜在变量 (latents),就将其传递给 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 显存 (VRAM),您应该能够以低于 8GB GPU 显存 (VRAM) 运行 PixArtSigmaPipeline。
如果您想要内存使用报告,请运行这个脚本。
以 8 位计算的文本嵌入 (Text embeddings) 可能会影响生成图像的质量,因为降低的精度会导致表征空间中的信息丢失。建议比较使用和不使用 8 位时的输出结果。
在加载 text_encoder
时,您将 load_in_8bit
设置为 True
。您也可以指定 load_in_4bit
,以进一步降低内存需求至 7GB 以下。
PixArtSigmaPipeline
class diffusers.PixArtSigmaPipeline
< source >( tokenizer: T5Tokenizer text_encoder: T5EncoderModel vae: AutoencoderKL transformer: PixArtTransformer2DModel scheduler: KarrasDiffusionSchedulers )
用于使用 PixArt-Sigma 进行文本到图像生成的 Pipeline。
__call__
< source >( 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 = 300 **kwargs ) → ImagePipelineOutput or tuple
参数
- prompt (
str
或List[str]
, 可选) — 用于引导图像生成的提示词 (prompt)。如果未定义,则必须传入prompt_embeds
。 - negative_prompt (
str
或List[str]
, 可选) — 不用于引导图像生成的提示词 (prompt)。如果未定义,则必须传入negative_prompt_embeds
。当不使用 guidance 时(即,如果guidance_scale
小于1
时)将被忽略。 - num_inference_steps (
int
, 可选, 默认为 100) — 去噪步骤 (denoising steps) 的数量。更多的去噪步骤通常会带来更高质量的图像,但会牺牲推理速度。 - timesteps (
List[int]
, 可选) — 用于去噪过程的自定义时间步长 (timesteps),适用于在其set_timesteps
方法中支持timesteps
参数的调度器 (schedulers)。如果未定义,将使用传递num_inference_steps
时的默认行为。必须以降序排列。 - sigmas (
List[float]
, 可选) — 用于去噪过程的自定义 sigmas,适用于在其set_timesteps
方法中支持sigmas
参数的调度器 (schedulers)。如果未定义,将使用传递num_inference_steps
时的默认行为。 - guidance_scale (
float
, 可选, 默认为 4.5) — Guidance scale,定义见 Classifier-Free Diffusion Guidance。guidance_scale
定义为 Imagen Paper 公式 2 中的w
。通过设置guidance_scale > 1
启用 Guidance scale。较高的 guidance scale 鼓励生成与文本prompt
紧密相关的图像,但通常会以降低图像质量为代价。 - num_images_per_prompt (
int
, 可选, 默认为 1) — 每个 prompt 生成的图像数量。 - height (
int
, 可选, 默认为 self.unet.config.sample_size) — 生成图像的高度像素值。 - width (
int
, 可选, 默认为 self.unet.config.sample_size) — 生成图像的宽度像素值。 - eta (
float
, 可选, 默认为 0.0) — 对应于 DDIM 论文中的参数 eta (η): https://arxiv.org/abs/2010.02502。仅适用于 schedulers.DDIMScheduler,对其他调度器将被忽略。 - generator (
torch.Generator
或List[torch.Generator]
, 可选) — 用于使生成过程具有确定性的 torch generator(s) 或列表。 - latents (
torch.Tensor
, 可选) — 预生成的噪声潜在变量 (noisy latents),从高斯分布中采样,用作图像生成的输入。可用于通过不同的 prompt 微调相同的生成结果。如果未提供,将通过使用提供的随机generator
采样来生成 latents 张量。 - prompt_embeds (
torch.Tensor
, 可选) — 预生成的文本嵌入 (text embeddings)。可用于轻松调整文本输入,例如 prompt 权重。如果未提供,将从prompt
输入参数生成文本嵌入。 - prompt_attention_mask (
torch.Tensor
, 可选) — 预生成的文本嵌入 (text embeddings) 的注意力掩码 (attention mask)。 - negative_prompt_embeds (
torch.Tensor
, 可选) — 预生成的负面文本嵌入 (negative text embeddings)。对于 PixArt-Sigma,此负面 prompt 应为 ""。如果未提供,将从negative_prompt
输入参数生成 negative_prompt_embeds。 - negative_prompt_attention_mask (
torch.Tensor
, 可选) — 预生成的负面文本嵌入 (negative text embeddings) 的注意力掩码 (attention mask)。 - output_type (
str
, 可选, 默认为"pil"
) — 生成图像的输出格式。在 PIL:PIL.Image.Image
或np.array
之间选择。 - return_dict (
bool
, 可选, 默认为True
) — 是否返回~pipelines.stable_diffusion.IFPipelineOutput
而不是普通的元组 (tuple)。 - callback (
Callable
, 可选) — 一个函数,它将在推理 (inference) 期间每callback_steps
步被调用。该函数将使用以下参数调用:callback(step: int, timestep: int, latents: torch.Tensor)
。 - callback_steps (
int
, 可选, 默认为 1) —callback
函数将被调用的频率。如果未指定,则将在每个步骤调用回调。 - clean_caption (
bool
, 可选, 默认为True
) — 是否在创建 embeddings 之前清理 caption。需要安装beautifulsoup4
和ftfy
。如果未安装这些依赖项,将从原始 prompt 创建 embeddings。 - use_resolution_binning (
bool
默认为True
) — 如果设置为True
,请求的高度和宽度将首先使用ASPECT_RATIO_1024_BIN
映射到最接近的分辨率。在生成的 latents 被解码为图像后,它们将被调整回请求的分辨率。这对于生成非正方形图像很有用。 - 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
< source >( 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 = 300 **kwargs )
参数
- prompt (
str
或List[str]
, 可选) — 要编码的 prompt - negative_prompt (
str
或List[str]
, 可选) — 不用于引导图像生成的 prompt。如果未定义,则必须传递negative_prompt_embeds
。当不使用 guidance 时忽略(即,如果guidance_scale
小于1
则忽略)。对于 PixArt-Alpha,这应为""
。 - do_classifier_free_guidance (
bool
, 可选, 默认为True
) — 是否使用 classifier free guidance - num_images_per_prompt (
int
, 可选, 默认为 1) — 每个 prompt 应生成的图像数量 - device — (
torch.device
, 可选): 用于放置结果 embeddings 的 torch 设备 - prompt_embeds (
torch.Tensor
, 可选) — 预生成的文本 embeddings。可用于轻松调整文本输入,例如 prompt weighting。如果未提供,则将从prompt
输入参数生成文本 embeddings。 - negative_prompt_embeds (
torch.Tensor
, 可选) — 预生成的负文本 embeddings。对于 PixArt-Alpha,它应该是""
字符串的 embeddings。 - clean_caption (
bool
, 默认为False
) — 如果为True
,该函数将在编码之前预处理和清理提供的 caption。 - max_sequence_length (
int
, 默认为 300) — 用于 prompt 的最大序列长度。
将 prompt 编码为文本编码器隐藏状态。