stable diffusion到底是如何工作的

stable diffusion简单入门

stable diffusion是一个文生图模型，主要由CompVis、Stability AI和LAION的研究者们创建。这个模型主要是在512X512分辨率的图像上训练的，训练数据集是LAION-5B，该数据集是目前可访问的最大的多模态数据集。

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

如果你是一个stable diffusion的新人，可以先了解一下sd的原理及代码相关的博客，这里推荐两篇比较好的博客

The Annotated Diffusion Model

Getting started with 🧨 Diffusers

用法

首先，应安装 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 可用，让我们将其移至cuda！

pipe.to("cuda")

如果您受到 GPU 内存的限制并且可用的 GPU RAM 少于 10GB，请确保加载 StableDiffusionPipeline float16 精度，而不是上面默认的 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 的信息。让我们看看它们是什么样子的：

可以打印一下结果，会有一个nsfw内容的boolean值返回，如果为True，则表示内容包含NSFW。

如果需要确定性输出，可以设定随机种子的种子，并将生成器传递到管道。每次使用具有相同种子的生成器时，您都会获得相同的图像输出。

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 是一种提高对生成的条件信号的影响（在本例中为文本）以及整体样本质量的依从性的方法。它也被称为无分类器引导，简单来说，它迫使生成者更好地匹配提示，可能会以牺牲图像质量或多样性为代价。介于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")

默认情况下，稳定扩散会生成 512 × 512 像素图像。使用 height 和 width 参数以纵向或横向比例创建矩形图像非常容易覆盖默认值。

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

• 确保 height 和 width 都是 8 的倍数。
• 低于 512 可能会导致图像质量降低。
• 在两个方向上超过 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")

在了解了稳定扩散可以产生的高质量图像之后，让我们尝试更好地理解模型的功能。

稳定扩散基于一种称为潜在扩散的特定类型的扩散模型，该模型由基于latent diffusion models合成高分辨率图像这一论文提出。

一般来说，扩散模型是经过训练的机器学习系统，可以逐步对随机高斯噪声进行降噪，以获得感兴趣的样本，例如图像。有关它们如何工作的更详细概述，请查看此colab。

扩散模型已被证明可以实现生成sota的图像数据。但扩散模型的一个缺点是，由于其重复的、连续的性质，反向去噪过程很慢。此外，这些模型会消耗大量内存，因为它们在像素空间中运行，这在生成高分辨率图像时会变得巨大。因此，训练这些模型并将它们用于推理是具有挑战性的。

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

潜在空间扩散三大组件

潜伏扩散有三个主要成分。

1. 一个自编码器(VAE)
2. 一个U-Net
3. 一个文本编码器，例如CLIP文本编码器

自动编码器

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

在潜在扩散训练期间，编码器用于获取前向扩散过程的图像的潜在表示（潜伏），该过程用于前向的噪声扩散过程，每一步会在前一步的基础上生成更多的噪声。在推理过程中，反向扩散过程产生的去噪潜伏使用VAE解码器转换回图像。正如我们将在推理过程中看到的那样，我们只需要 VAE 解码器。

UNet网络

U-Net 有一个编码器部分和一个解码器部分，两者都由 ResNet 模块组成。编码器将图像表示压缩为较低分辨率的图像表示，解码器将较低分辨率的图像表示解码回原始的高分辨率图像表示，该图像表示应该噪声较小。更具体地说，U-Net 输出预测噪声残差，可用于计算预测的去噪图像表示。

为了防止U-Net在下采样时丢失重要信息，通常在编码器的下采样ResNet和解码器的上采样ResNet之间添加捷径连接。此外，稳定的扩散 U-Net 能够通过交叉注意力层在文本嵌入上调节其输出。交叉注意力层被添加到U-Net的编码器和解码器部分，通常在ResNet模块之间。

文本编码器

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

受 Imagen 的启发，Stable Diffusion 不会在训练期间训练文本编码器，而只是使用 CLIP 已经训练好的文本编码器 CLIPTextModel。

为什么潜伏扩散快速高效？

由于潜在扩散在低维空间上运行，因此与像素空间扩散模型相比，它大大降低了内存和计算要求。例如，Stable Diffusion 中使用的自动编码器的折减系数为 8。这意味着形状 (3, 512, 512) 的图像会进入 (3, 64, 64) 潜在空间，这意味着 可以减少8 × 8 = 64倍的内存开销。

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

推理过程中的稳定扩散

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

稳定扩散模型将潜在种子和文本提示作为输入。然后，潜伏种子用于生成大小 64×6464×64 的随机潜在图像表示，其中文本提示通过 CLIP 的文本编码器转换为大小 77×76877×768 的文本嵌入。

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

• PNDM scheduler
• DDIM scheduler
• K-LMS scheduler

关于调度程序算法函数如何超出本笔记本范围的理论，但简而言之，人们应该记住，它们从先前的噪声表示和预测的噪声残差中计算预测的去噪图像表示。有关更多信息，我们建议研究阐明基于扩散的生成模型的设计空间

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

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

编写自己的推理管道

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

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

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

• 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 inital latent noise

batch_size = len(prompt)

首先，我们得到 text_embeddings for the passed 提示。这些嵌入将用于调节 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]

参考链接：

https://huggingface.co/blog/stable_diffusion#how-does-stable-diffusion-work