PixArt-α
PixArt-α: 用于照片级逼真文本到图像合成的扩散Transformer快速训练 的作者是陈俊松、余锦成、葛崇建、姚乐伟、谢恩泽、吴悦、王忠道、郭志华、罗平、陆虎川和李正国。
论文中的摘要如下:
最先进的文本到图像(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与6250个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位精度加载文本编码器,使用PixArtAlphaPipeline运行低于8GB GPU显存。让我们逐步完成一个完整的示例。
首先,安装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位的输出。
在加载text_encoder
时,您将load_in_8bit
设置为True
。您还可以指定load_in_4bit
,将内存需求进一步降低到低于7GB。
PixArtAlphaPipeline
class 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 的分词器。 - 变换器 (PixArtTransformer2DModel) — 一个文本条件的
PixArtTransformer2DModel
,用于对编码后的图像潜在变量进行去噪。 - 调度器 (SchedulerMixin) — 与
transformer
结合使用,对编码后的图像潜在变量进行去噪的调度器。
使用 PixArt-Alpha 进行文本到图像生成的管道。
此模型继承自 DiffusionPipeline。请查看超类文档,了解库为所有管道实现的通用方法(例如下载或保存、在特定设备上运行等)。
__call__
< 源代码 >( 提示: Union = None 负面提示: str = '' 推理步数: int = 20 时间步长: List = None 方差: List = None 引导尺度: float = 4.5 每个提示的图像数量: Optional = 1 高度: Optional = None 宽度: Optional = None eta: float = 0.0 生成器: Union = None 潜在变量: Optional = None 提示嵌入: Optional = None 提示注意力掩码: Optional = None 负面提示嵌入: Optional = None 负面提示注意力掩码: Optional = None 输出类型: Optional = 'pil' 返回字典: bool = True 回调函数: Optional = None 回调步数: int = 1 清理标题: bool = True 使用分辨率分箱: bool = True 最大序列长度: int = 120 **kwargs ) → ImagePipelineOutput 或 元组
参数
- 提示 (
str
或List[str]
, 可选) — 指导图像生成的提示或提示列表。如果未定义,则必须传递prompt_embeds
。代替。 - 负面提示 (
str
或List[str]
, 可选) — 不用于指导图像生成的提示或提示列表。如果未定义,则必须传递negative_prompt_embeds
代替。在不使用引导时忽略(即,如果guidance_scale
小于1
,则忽略)。 - 推理步数 (
int
, 可选, 默认为 100) — 去噪步骤的数量。更多的去噪步骤通常会导致更高的图像质量,但推理速度会变慢。 - timesteps (
List[int]
, 可选) — 用于降噪过程的自定义时间步长,适用于在其set_timesteps
方法中支持timesteps
参数的调度器。如果未定义,则将使用传递num_inference_steps
时的默认行为。必须按降序排列。 - sigmas (
List[float]
, 可选) — 用于降噪过程的自定义西格玛,适用于在其set_timesteps
方法中支持sigmas
参数的调度器。如果未定义,则将使用传递num_inference_steps
时的默认行为。 - guidance_scale (
float
, 可选, defaults to 4.5) — 指导尺度,如 Classifier-Free Diffusion Guidance 中所定义。guidance_scale
定义为 Imagen 论文 中公式 2 的w
。通过设置guidance_scale > 1
来启用指导尺度。较高的指导尺度鼓励生成与文本prompt
密切相关的图像,通常以降低图像质量为代价。 - num_images_per_prompt (
int
, 可选, defaults to 1) — 每个提示生成图像的数量。 - height (
int
, 可选, defaults to self.unet.config.sample_size) — 生成的图像的高度(以像素为单位)。 - width (
int
, 可选, defaults to self.unet.config.sample_size) — 生成的图像的宽度(以像素为单位)。 - eta (
float
, 可选, defaults to 0.0) — 对应于 DDIM 论文中的参数 eta (η):https://arxiv.org/abs/2010.02502。仅适用于 schedulers.DDIMScheduler,其他情况将被忽略。 - generator (
torch.Generator
或List[torch.Generator]
, 可选) — 一个或一组 torch 生成器,用于使生成确定性。 - latents (
torch.Tensor
, 可选) — 预生成的噪声潜在变量,从高斯分布中采样,用作图像生成的输入。可用于使用不同的提示调整相同的生成。如果未提供,则将通过使用提供的随机generator
采样生成潜在变量张量。 - prompt_embeds (
torch.Tensor
, 可选) — 预生成的文本嵌入。可用于轻松调整文本输入,例如 提示加权。如果未提供,则将从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
默认为 120) — 与prompt
一起使用的最大序列长度。
返回
ImagePipelineOutput 或 tuple
如果return_dict
为True
,则返回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: Union do_classifier_free_guidance: bool = True negative_prompt: str = '' num_images_per_prompt: int = 1 device: Optional = None prompt_embeds: Optional = None negative_prompt_embeds: Optional = None prompt_attention_mask: Optional = None negative_prompt_attention_mask: Optional = None clean_caption: bool = False max_sequence_length: int = 120 **kwargs )
参数
- prompt(
str
或List[str]
,可选) — 要编码的提示 - negative_prompt(
str
或List[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) — 用于提示的最大序列长度。
将提示编码为文本编码器隐藏状态。