Hugging Face 教你用 🧨 Diffusers 运行 Stable Diffusion 并自定义生成管线

Stable Diffusion 🎨 ...使用 🧨 Diffusers

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

在这篇文章中，我们想展示如何将 Stable Diffusion 与 🧨 Diffusers 库一起使用，解释模型的工作原理，最后更深入地探讨 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 管道，只需几行代码即可运行 Stable Diffusion 模型进行推理。该管道设置了从文本生成图像所需的一切，只需一个简单的 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)

要运行管道，只需定义提示并调用 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. 文本编码器

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

受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 之后，让我们看看如何高级使用 🤗 Hugging Face diffusers 库！

编写自己的推理管道

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

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

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

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]

我们已经从使用 🤗 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)},
}