使用🧨 Diffusers 进行 Stable Diffusion

Stable Diffusion 🎨 ...使用 🧨 Diffusers

Stable Diffusion 是一个文本到图像的潜在扩散模型，由 CompVis、Stability AI 和 LAION 的研究人员和工程师创建。它在 LAION-5B 数据库的一个子集中训练了 512x512 图像。LAION-5B 是目前最大、免费访问的多模态数据集。

在这篇文章中，我们将展示如何使用 🧨 Diffusers 库 来使用 Stable Diffusion，解释模型的工作原理，并深入探讨 diffusers 如何允许自定义图像生成管道。

注意：强烈建议对扩散模型的工作原理有基本的了解。如果您是扩散模型的新手，我们建议阅读以下博客文章之一

• 图解扩散模型
• Diffusers 入门

现在，让我们开始生成一些图像 🎨。

运行 Stable Diffusion

许可证

在使用模型之前，您需要接受模型许可证才能下载和使用权重。注意：不再需要通过 UI 明确接受许可证。

该许可证旨在减轻如此强大的机器学习系统可能产生的有害影响。我们请求用户完整而仔细地阅读许可证。这里我们提供一个摘要

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，只需少量代码更改。

Stable Diffusion 模型只需几行代码即可通过 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 = "a photograph of an astronaut riding a horse"

image = pipe(prompt).images[0]

# you can save the image with
# image.save(f"astronaut_rides_horse.png")

结果如下所示

之前的代码每次运行时都会生成不同的图像。

如果您在某个时候获得了黑色图像，这可能是因为模型中内置的内容过滤器检测到 NSFW 结果。如果您认为不应该出现这种情况，请尝试调整您的提示或使用不同的种子。实际上，模型预测中包含了有关特定结果是否检测到 NSFW 的信息。让我们看看它们是什么样子

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]

# you can save the image with
# image.save(f"astronaut_rides_horse.png")

结果如下所示

您可以使用 num_inference_steps 参数更改推理步数。

一般来说，步数越多结果越好，但步数越多，生成时间越长。Stable Diffusion 在相对较少的步数下也能很好地工作，因此我们建议使用默认的 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]

# you can save the image with
# image.save(f"astronaut_rides_horse.png")

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

除了 num_inference_steps，我们在之前所有示例中都使用了另一个函数参数，称为 guidance_scale。guidance_scale 是一种提高对引导生成（在本例中是文本）的条件信号的遵循程度以及整体样本质量的方法。它也被称为无分类器引导，简单来说就是强制生成结果更好地匹配提示，这可能会以图像质量或多样性为代价。对于 Stable Diffusion，介于 7 和 8.5 之间的值通常是不错的选择。默认情况下，管道使用 7.5 的 guidance_scale。

如果您使用一个非常大的值，图像可能看起来不错，但多样性会降低。您可以在本文的此部分了解此参数的技术细节。

接下来，我们看看如何一次生成多个相同提示的图像。首先，我们将创建一个 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 = ["a photograph of an astronaut riding a horse"] * num_images

images = pipe(prompt).images

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

# you can save the grid with
# grid.save(f"astronaut_rides_horse.png")

默认情况下，Stable Diffusion 生成 512 × 512 像素的图像。使用 height 和 width 参数很容易覆盖默认值，以创建纵向或横向比例的矩形图像。

选择图像尺寸时，我们建议如下

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

让我们运行一个例子

prompt = "a photograph of an astronaut riding a horse"
image = pipe(prompt, height=512, width=768).images[0]

# you can save the image with
# image.save(f"astronaut_rides_horse.png")

Stable Diffusion 工作原理

了解了 Stable Diffusion 可以产生的高质量图像后，让我们尝试更好地理解模型的功能。

Stable Diffusion 基于一种特殊类型的扩散模型，称为潜在扩散，在《使用潜在扩散模型进行高分辨率图像合成》中提出。

一般来说，扩散模型是机器学习系统，它们经过训练，可以一步步地对随机高斯噪声进行去噪，以获得感兴趣的样本，例如图像。有关它们如何工作的更详细概述，请查看此 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 之间添加快捷连接。此外，Stable Diffusion U-Net 能够通过交叉注意力层以文本嵌入为条件输出。交叉注意力层通常添加到 U-Net 的编码器和解码器部分，通常在 ResNet 块之间。

3. 文本编码器

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

受Imagen的启发，Stable Diffusion在训练期间不训练文本编码器，而是简单地使用CLIP的已训练文本编码器，即CLIPTextModel。

为什么潜在扩散既快速又高效？

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

这就是为什么即使在 16GB Colab GPU 上，也能如此快速地生成 512 × 512 图像！

推理过程中的 Stable Diffusion

综上所述，现在让我们通过说明逻辑流程来仔细研究模型如何在推理中工作。

Stable Diffusion 模型同时接收潜在种子和文本提示作为输入。然后，潜在种子用于生成大小为 $64 \times 64$ 的随机潜在图像表示，而文本提示则通过 CLIP 的文本编码器转换为大小为 $77 \times 768$ 的文本嵌入。

接下来，U-Net 在文本嵌入的条件下迭代地去噪随机潜在图像表示。U-Net 的输出，即噪声残差，用于通过调度器算法计算去噪后的潜在图像表示。许多不同的调度器算法可用于此计算，每种算法都有其优点和缺点。对于 Stable Diffusion，我们建议使用以下算法之一

• PNDM 调度器（默认使用）
• DDIM 调度器
• K-LMS 调度器

调度器算法函数的工作原理超出了本笔记本的范围，但简而言之，应该记住它们从先前的噪声表示和预测的噪声残差计算预测的去噪图像表示。有关更多信息，我们建议查阅《阐明基于扩散的生成模型的设计空间》

去噪过程重复约 50 次，以逐步检索更好的潜在图像表示。完成后，潜在图像表示由变分自编码器的解码器部分解码。

在对潜在和 Stable Diffusion 进行了简要介绍之后，让我们看看如何高级使用 🤗 Hugging Face diffusers 库！

编写自己的推理管道

最后，我们展示了如何使用 diffusers 创建自定义扩散管道。编写自定义推理管道是 diffusers 库的高级用法，可用于切换某些组件，例如上面解释的 VAE 或调度器。

例如，我们将展示如何使用 Stable Diffusion 和不同的调度器，即Katherine Crowson的 K-LMS 调度器，该调度器在此 PR中添加。

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

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

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

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

# 1. Load the autoencoder model which will be used to decode the latents into image space. 
vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae")

# 2. Load the tokenizer and text encoder to tokenize and encode the text. 
tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")

# 3. The UNet model for generating the latents.
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) 的引导权重 w。guidance_scale == 1 对应于不进行无分类器引导。这里我们将其设置为 7.5，如前所述。

与前面的示例不同，我们将 num_inference_steps 设置为 100，以获得更清晰的图像。

prompt = ["a photograph of an astronaut riding a horse"]

height = 512                        # default height of Stable Diffusion
width = 512                         # default width of Stable Diffusion

num_inference_steps = 100           # Number of denoising steps

guidance_scale = 7.5                # Scale for classifier-free guidance

generator = torch.manual_seed(0)    # Seed generator to create the initial latent noise

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_embeddings 具有相同的形状（batch_size 和 seq_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):
    # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes.
    latent_model_input = torch.cat([latents] * 2)

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

    # predict the noise residual
    with torch.no_grad():
        noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample

    # perform guidance
    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

    # compute the previous noisy sample x_t -> x_t-1
    latents = scheduler.step(noise_pred, t, latents).prev_sample

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

# scale and decode the image latents with vae
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 入门》笔记本，它对扩散系统进行了更广泛的概述。
• 《注释扩散模型》博客文章。
• 我们在 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)},
}