用 Diffusers 运行 Stable Diffusion 并自定义推理 pipeline

Stable Diffusion 🎨 ...使用 🧨 Diffusers

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

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

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

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

运行 Stable Diffusion

许可证

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

该许可证旨在减轻这种强大的机器学习系统可能带来的有害影响。我们要求用户 完整且仔细地阅读许可证。以下是我们提供的摘要：

你不能使用该模型故意生成或分享非法或有害的输出或内容，
我们对你生成的输出不主张任何权利，你可以自由使用它们，并对它们的使用负责，且不应违反许可证中的规定，并且
你可以重新分发权重，并将模型用于商业用途和/或作为服务。如果你这样做，请注意你必须包含与许可证中相同的使用限制，并向所有用户分享 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 pipeline，只需几行代码就可以运行 Stable Diffusion 模型的推理。该 pipeline 设置了你需要的一切，通过简单的 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 精度加载 StableDiffusionPipeline，而不是像上面那样默认的 float32 精度。

你可以通过从 fp16 分支加载权重，并告诉 diffusers 预期权重为 float16 精度来做到这一点：

import torch
from diffusers import StableDiffusionPipeline

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

要运行 pipeline，只需定义提示词并调用 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]
}

如果你想要确定性的输出，可以设置随机种子，并向 pipeline 传入一个 generator。每次使用相同种子的 generator 时，你都会得到相同的图像输出。

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。如果你想要更快的结果，可以使用更小的数值。如果你想要可能更高质量的结果，可以使用更大的数值。

让我们尝试用更少的去噪步数来运行 pipeline。

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 之间的值通常是很好的选择。默认情况下，pipeline 使用的 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

我们可以通过简单地使用一个将同一提示词重复多次的列表，来为同一个提示词生成多张图像。我们将把这个列表发送给 pipeline，而不是之前使用的字符串。

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 基于一种特定类型的扩散模型，称为潜空间扩散（Latent Diffusion），在High-Resolution Image Synthesis with Latent Diffusion Models 中提出。

一般来说，扩散模型是经过训练以逐步去噪随机高斯噪声的机器学习系统，从而得到感兴趣的样本，例如图像。有关其工作原理的更详细概述，请查看这个 colab。

扩散模型已被证明在生成图像数据方面能够达到最先进的结果。但扩散模型的一个缺点是，反向去噪过程由于其重复的、顺序的性质而较慢。此外，这些模型消耗大量内存，因为它们在像素空间中操作，当生成高分辨率图像时，像素空间变得非常大。因此，训练这些模型并将它们用于推理是具有挑战性的。

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

潜在扩散有三个主要组成部分。

一个自动编码器（VAE）。
一个 U-Net。
一个文本编码器，例如 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. 文本编码器

文本编码器负责将输入提示，例如 “An astronaut riding a horse” 转换为 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

综合以上所述，现在让我们通过展示逻辑流程来更仔细地了解模型在推理时是如何工作的。

sd-pipeline

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

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

有关调度器算法工作原理的理论超出了本笔记本的范围，但简而言之，需要记住它们根据先前的噪声表示和预测的噪声残差来计算预测的去噪图像表示。欲了解更多信息，我们建议查阅 Elucidating the Design Space of Diffusion-Based Generative Models

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

在简要介绍了潜在扩散和 Stable Diffusion 之后，让我们看看如何高级使用 🤗 HF Mirror diffusers 库！

编写你自己的推理管道

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

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

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

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

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

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]

我们已经从使用 🤗 HF Mirror 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 = {HF Mirror Blog},
  year = {2022},
  note = {[https://hf-mirror.com/blog/rlhf](https://hf-mirror.com/blog/stable_diffusion)},
}