稳定扩散与🧨扩散器

Steady diffusion and explosive diffuser

…使用🧨扩散器

稳定扩散(Stable Diffusion)是由CompVis、Stability AI和LAION的研究人员和工程师创建的一种文本到图像潜在扩散模型。它是在LAION-5B数据库的子集上训练的512×512图像。LAION-5B是目前存在的最大的、免费可访问的多模态数据集。

在本文中,我们将展示如何使用🧨扩散器库使用稳定扩散(Stable Diffusion),解释模型的工作原理,并深入了解diffusers如何允许定制图像生成流程。

注意:强烈建议您对扩散模型的工作原理有基本的了解。如果扩散模型对您来说完全是新的,请阅读以下博客文章之一:

  • 批注扩散模型
  • 开始使用🧨扩散器

现在,让我们通过生成一些图像🎨来开始吧。

运行稳定扩散

许可证

在使用该模型之前,您需要接受模型许可证以便下载和使用权重。注意:许可证不再需要通过用户界面显式接受。

该许可证旨在减轻这种强大机器学习系统的潜在有害影响。我们要求用户完整和仔细阅读许可证。这里我们提供一个摘要:

  1. 您不能使用该模型故意产生或分享非法或有害的输出或内容,
  2. 我们对您生成的输出不享有任何权利,您可以自由使用这些输出,并对其使用负责,不得违反许可证中的规定,以及
  3. 您可以重新分发权重,并将该模型用于商业目的和/或作为服务。如果您这样做,请注意您必须包含与许可证中相同的使用限制,并向所有用户共享CreativeML OpenRAIL-M的副本。

使用

首先,您应该安装diffusers==0.10.2以运行以下代码片段:

pip install diffusers==0.10.2 transformers scipy ftfy accelerate

在本文中,我们将使用模型版本v1-4,但您也可以使用其他版本的模型,如1.5、2和2.1,只需进行最小的代码更改即可。

稳定扩散模型可以通过使用StableDiffusionPipeline流水线进行推断,只需使用简单的from_pretrained函数调用即可完成设置。

from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")

如果有可用的GPU,让我们将其移到GPU上!

pipe.to("cuda")

注意:如果您的GPU内存有限且可用的GPU RAM少于10GB,请确保以float16精度而不是默认的float32精度加载StableDiffusionPipeline,如上述所做。

您可以通过从fp16分支加载权重,并告诉diffusers期望权重以float16精度进行加载来实现:

import torch
from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", revision="fp16", torch_dtype=torch.float16)

要运行流水线,只需定义提示并调用pipe

prompt = "一个骑马的宇航员的照片"

image = pipe(prompt).images[0]

# 您可以保存图像:
# image.save(f"astronaut_rides_horse.png")

结果如下所示

以前的代码每次运行都会给你一个不同的图像。

如果你在某个时刻得到了一张黑色的图像,可能是因为模型内置的内容过滤器检测到了不适宜内容。如果你认为这不应该是这样的,请尝试调整你的提示或使用不同的种子。事实上,模型的预测结果中包含了关于是否检测到不适宜内容的信息。让我们看看它们是什么样子:

result = pipe(prompt)
print(result)

{
    'images': [<PIL.Image.Image image mode=RGB size=512x512>],
    'nsfw_content_detected': [False]
}

如果你想要确定性的输出,你可以使用一个随机种子生成器来设置一个随机种子,并将其传递给流水线。每次使用相同种子的生成器时,你将得到相同的图像输出。

import torch

generator = torch.Generator("cuda").manual_seed(1024)
image = pipe(prompt, guidance_scale=7.5, generator=generator).images[0]

# 你可以保存图像
# image.save(f"astronaut_rides_horse.png")

结果将如下所示

你可以使用 num_inference_steps 参数来改变推理步骤的数量。

一般来说,使用更多的步骤会得到更好的结果,但是步骤越多,生成所需的时间就越长。稳定扩散在相对较少的步骤下工作得非常好,所以我们建议使用默认的推理步骤数 50。如果你想要更快的结果,可以使用较小的数字。如果你想要更高质量的结果,可以使用较大的数字。

让我们尝试使用较少的去噪步骤运行流水线。

import torch

generator = torch.Generator("cuda").manual_seed(1024)
image = pipe(prompt, guidance_scale=7.5, num_inference_steps=15, generator=generator).images[0]

# 你可以保存图像
# image.save(f"astronaut_rides_horse.png")

请注意,结构是相同的,但是宇航员服装和马的整体形态存在问题。这表明仅使用 15 个去噪步骤已经显著降低了生成结果的质量。正如之前所述,通常使用 50 个去噪步骤就足以生成高质量的图像。

除了 num_inference_steps,我们在所有之前的示例中都使用了另一个函数参数,称为 guidance_scaleguidance_scale 是一种增加对生成的条件信号(文本,在这种情况下)以及整体样本质量的粘合力的方法。它也被称为无分类器指导,简单地说就是强制生成更好地匹配提示的生成,可能以图像质量或多样性为代价。在稳定扩散中,通常选择 78.5 之间的值。默认情况下,流水线使用 guidance_scale 为 7.5。

如果你使用一个非常大的值,图像可能看起来很好,但是多样性会较小。你可以在本文的这一部分了解有关该参数的技术细节。

接下来,让我们看看如何一次生成多张相同提示的图像。首先,我们将创建一个 image_grid 函数来帮助我们将它们漂亮地显示在一个网格中。

from PIL import Image

def image_grid(imgs, rows, cols):
    assert len(imgs) == rows*cols

    w, h = imgs[0].size
    grid = Image.new('RGB', size=(cols*w, rows*h))
    grid_w, grid_h = grid.size
    
    for i, img in enumerate(imgs):
        grid.paste(img, box=(i%cols*w, i//cols*h))
    return grid

我们可以通过简单地使用一个重复多次的相同提示的列表来为同一个提示生成多个图像。我们将列表发送给流水线,而不是之前使用的字符串。

num_images = 3
prompt = ["一个骑着马的宇航员的照片"] * num_images

images = pipe(prompt).images

grid = image_grid(images, rows=1, cols=3)

# 你可以保存网格
# grid.save(f"astronaut_rides_horse.png")

默认情况下,稳定扩散生成的图像大小为512 × 512像素。你可以使用heightwidth参数来覆盖默认设置,以创建纵向或横向比例的矩形图像。

在选择图像大小时,请遵循以下建议:

  • 确保heightwidth都是8的倍数。
  • 降低到512以下可能会导致图像质量下降。
  • 在两个方向上超过512将会重复图像区域(全局一致性丢失)。
  • 创建非方形图像的最佳方法是在一个维度上使用512,在另一个维度上使用大于该值的值。

让我们运行一个示例:

prompt = "一个骑马的宇航员的照片"
image = pipe(prompt, height=512, width=768).images[0]

# 你可以保存图像
# image.save(f"astronaut_rides_horse.png")

稳定扩散是如何工作的?

通过观察稳定扩散可以生成的高质量图像,让我们更好地理解模型的运行方式。

稳定扩散基于一种特殊类型的扩散模型,称为潜在扩散,提出于《具有潜在扩散模型的高分辨率图像合成》。

一般来说,扩散模型是一种通过逐步去噪随机高斯噪声来训练的机器学习系统,以得到感兴趣的样本,例如图像。要了解它们的工作原理的更详细概述,请参阅此Colab。

扩散模型在生成图像数据方面已经展现出了最先进的结果。但是扩散模型的一个缺点是,由于其重复的顺序过程,反向去噪过程很慢。此外,这些模型在像素空间中操作,消耗大量内存,在生成高分辨率图像时变得非常庞大。因此,训练这些模型并将其用于推理是具有挑战性的。

潜在扩散可以通过在较低维度潜在空间上应用扩散过程,而不是使用实际的像素空间,来减少内存和计算复杂性。这是标准扩散和潜在扩散模型之间的关键区别:在潜在扩散中,模型被训练为生成图像的潜在(压缩)表示。

潜在扩散有三个主要组件。

  1. 自编码器(VAE)。
  2. U-Net。
  3. 文本编码器,例如CLIP的文本编码器。

1. 自编码器(VAE)

VAE模型由编码器和解码器两部分组成。编码器用于将图像转换为低维潜在表示,这将作为U-Net模型的输入。相反,解码器将潜在表示转换回图像。

在潜在扩散训练中,编码器用于获取图像的潜在表示(潜在变量),用于前向扩散过程,在每个步骤中逐渐应用更多噪声。在推理中,通过反向扩散过程生成的去噪潜在变量使用VAE解码器转换回图像。正如我们将在推理中看到的那样,我们只需要VAE解码器。

2. U-Net

U-Net由编码器部分和解码器部分组成,两者都包含ResNet块。编码器将图像表示压缩为较低分辨率的图像表示,解码器将较低分辨率的图像表示解码为原始较高分辨率的图像表示,假设该表示较少噪声。更具体地说,U-Net的输出预测了噪声残差,可以用于计算预测的去噪图像表示。

为了防止U-Net在下采样过程中丢失重要信息,通常在编码器的下采样ResNet和解码器的上采样ResNet之间添加了快捷连接。此外,稳定扩散U-Net能够通过交叉注意力层将其输出与文本嵌入进行条件化。交叉注意力层通常添加在U-Net的编码器和解码器部分之间的ResNet块之间。

3. 文本编码器

文本编码器负责将输入提示,例如”一个骑马的宇航员”,转换为可以被 U-Net 理解的嵌入空间。它通常是一个简单的基于 Transformer 的编码器,将一系列输入标记映射为一系列潜在的文本嵌入。

受 Imagen 的启发,稳定扩散在训练过程中不训练文本编码器,而是直接使用已经训练好的 CLIP 文本编码器 CLIPTextModel 。

为什么潜在扩散快速高效?

由于潜在扩散在低维空间上操作,相比像素空间扩散模型大大减少了内存和计算需求。例如,稳定扩散中使用的自动编码器的缩减因子为8。这意味着形状为 (3, 512, 512) 的图像在潜在空间中变为 (3, 64, 64),需要的内存减少了 8 × 8 = 64 倍。

这就是为什么可以在 16GB Colab GPU 上如此快速地生成 512 × 512 的图像!

推断过程中的稳定扩散

将这些内容整合在一起,现在让我们更详细地看一下模型在推断中的工作原理,通过说明逻辑流程。

稳定扩散模型同时接受潜在种子和文本提示作为输入。然后使用潜在种子生成大小为 64 × 64 的随机潜在图像表示,而文本提示被转换为大小为 77 × 768 的文本嵌入,通过 CLIP 的文本编码器。

接下来,U-Net 在被文本嵌入条件下迭代去噪随机潜在图像表示。U-Net 的输出是噪声残差,通过调度算法用于计算去噪后的潜在图像表示。可以使用许多不同的调度算法进行计算,每个算法都有其优缺点。对于稳定扩散,我们推荐使用以下之一:

  • PNDM 调度器(默认使用)
  • DDIM 调度器
  • K-LMS 调度器

调度算法的工作原理超出了本笔记本的范围,简单来说,它们根据前一个噪声表示和预测的噪声残差计算预测的去噪图像表示。要了解更多信息,建议参阅《阐明基于扩散的生成模型设计空间》。

去噪过程重复约50次,逐步获取更好的潜在图像表示。完成后,潜在图像表示通过变分自编码器的解码器部分进行解码。

在简要介绍了潜在扩散和稳定扩散之后,让我们看看如何更高级地使用 🤗 Hugging Face 的 diffusers 库!

编写自己的推断流程

最后,我们将展示如何使用 diffusers 创建自定义扩散流程。编写自定义推断流程是 diffusers 库的高级用法,可以用于替换某些组件,例如上面解释的 VAE 或调度器。

例如,我们将展示如何使用不同的调度器(Katherine Crowson 的 K-LMS 调度器)来使用稳定扩散,该调度器是在此 PR 中添加的。

预训练模型包含了设置完整扩散流程所需的所有组件。它们存储在以下文件夹中:

  • text_encoder:稳定扩散使用 CLIP,但其他扩散模型可能使用其他编码器,例如 BERT
  • tokenizer:必须与 text_encoder 模型使用的标记器相匹配。
  • scheduler:训练过程中用于逐步向图像添加噪声的调度算法。
  • unet:用于生成输入的潜在表示的模型。
  • vae:自编码器模块,用于将潜在表示解码为真实图像。

我们可以通过引用保存组件的文件夹来加载组件,使用from_pretrained函数的subfolder参数。

from transformers import CLIPTextModel, CLIPTokenizer
from diffusers import AutoencoderKL, UNet2DConditionModel, PNDMScheduler

# 1. 加载用于将潜在空间解码为图像空间的自编码器模型。
vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae")

# 2. 加载用于标记化和编码文本的分词器和文本编码器。
tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")

# 3. 生成潜在空间的UNet模型。
unet = UNet2DConditionModel.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="unet")

现在,我们加载具有一些拟合参数的K-LMS调度器,而不是加载预定义的调度器。

from diffusers import LMSDiscreteScheduler

scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)

接下来,让我们将模型移动到GPU上。

torch_device = "cuda"
vae.to(torch_device)
text_encoder.to(torch_device)
unet.to(torch_device) 

现在,我们定义将用于生成图像的参数。

请注意,guidance_scale的定义类似于Imagen论文中方程(2)的guidance weight wguidance_scale == 1表示不进行无分类器引导。在这里,我们将其设置为7.5,与之前的设置相同。

与之前的示例不同,我们将num_inference_steps设置为100,以获得更明确的图像。

prompt = ["一张骑马的宇航员的照片"]

height = 512                        # Stable Diffusion的默认高度
width = 512                         # Stable Diffusion的默认宽度

num_inference_steps = 100           # 去噪步骤数

guidance_scale = 7.5                # 无分类器引导的缩放比例

generator = torch.manual_seed(0)    # 设置生成器的种子以创建初始潜在噪声

batch_size = len(prompt)

首先,我们获取传递的提示的text_embeddings。这些嵌入将用于对UNet模型进行条件设置,并将图像生成引导至应该类似于输入提示的内容。

text_input = tokenizer(prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")

text_embeddings = text_encoder(text_input.input_ids.to(torch_device))[0]

我们还会获取用于无分类器引导的无条件文本嵌入,这些嵌入仅为填充令牌(空文本)的嵌入。它们需要与条件的text_embeddingsbatch_sizeseq_length)具有相同的形状。

max_length = text_input.input_ids.shape[-1]
uncond_input = tokenizer(
    [""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt"
)
uncond_embeddings = text_encoder(uncond_input.input_ids.to(torch_device))[0]   

对于无分类器引导,我们需要进行两次前向传播:一次使用有条件的输入(text_embeddings),另一次使用无条件嵌入(uncond_embeddings)。实际上,我们可以将两者连接成一个批次,以避免进行两次前向传播。

text_embeddings = torch.cat([uncond_embeddings, text_embeddings])

接下来,我们生成初始的随机噪声。

latents = torch.randn(
    (batch_size, unet.in_channels, height // 8, width // 8),
    generator=generator,
)
latents = latents.to(torch_device)

如果我们在此阶段检查latents,我们会看到它们的形状是torch.Size([1, 4, 64, 64]),比我们要生成的图像要小得多。该模型将会将这种潜在表示(纯噪声)转换为512 × 512的图像。

接下来,我们使用选择的 num_inference_steps 初始化调度器。这将计算在去噪过程中使用的 sigmas 和精确的时间步值。

scheduler.set_timesteps(num_inference_steps)

K-LMS 调度器需要将 latents 乘以其 sigma 值。让我们在这里做一下:

latents = latents * scheduler.init_noise_sigma

我们准备好编写去噪循环了。

from tqdm.auto import tqdm

scheduler.set_timesteps(num_inference_steps)

for t in tqdm(scheduler.timesteps):
    # 如果我们使用无分类器指导,则扩展 latents 以避免进行两次前向传递。
    latent_model_input = torch.cat([latents] * 2)

    latent_model_input = scheduler.scale_model_input(latent_model_input, timestep=t)

    # 预测噪声残差
    with torch.no_grad():
        noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample

    # 进行指导
    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

    # 计算先前的有噪声样本 x_t -> x_t-1
    latents = scheduler.step(noise_pred, t, latents).prev_sample

现在我们使用 vae 将生成的 latents 解码回图像。

# 使用 vae 对图像 latents 进行缩放和解码
latents = 1 / 0.18215 * latents
with torch.no_grad():
    image = vae.decode(latents).sample

最后,让我们将图像转换为 PIL 格式,以便显示或保存。

image = (image / 2 + 0.5).clamp(0, 1)
image = image.detach().cpu().permute(0, 2, 3, 1).numpy()
images = (image * 255).round().astype("uint8")
pil_images = [Image.fromarray(image) for image in images]
pil_images[0]

我们从使用🤗 Hugging Face Diffusers的基本 Stable Diffusion 的用法过渡到了库的更高级用法,并尝试在现代扩散系统中介绍了所有的组成部分。如果您喜欢这个主题并想了解更多,我们推荐以下资源:

  • 我们的 Colab 笔记本。
  • 入门 Diffusers 笔记本,提供了关于 Diffusion 系统的更广泛概述。
  • 批注 Diffusion Model 博文。
  • 我们在 GitHub 上的代码,如果 diffusers 对您有用,我们将非常高兴您给个 ⭐!

引用:

@article{patil2022stable,
  author = {Patil, Suraj and Cuenca, Pedro and Lambert, Nathan and von Platen, Patrick},
  title = {Stable Diffusion with 🧨 Diffusers},
  journal = {Hugging Face Blog},
  year = {2022},
  note = {[https://huggingface.co/blog/rlhf](https://huggingface.co/blog/stable_diffusion)},
}