【CV】稳定扩散模型（Stable Diffusion）

  🔎大家好，我是Sonhhxg_柒，希望你看完之后，能对你有所帮助，不足请指正！共同学习交流🔎

📝个人主页－Sonhhxg_柒的博客_CSDN博客 📃

🎁欢迎各位→点赞👍 + 收藏⭐️ + 留言📝​

📣系列专栏 - 机器学习【ML】 自然语言处理【NLP】  深度学习【DL】

​​

 🖍foreword

✔说明⇢本人讲解主要包括Python、机器学习（ML）、深度学习（DL）、自然语言处理（NLP）等内容。

如果你对这个系列感兴趣的话，可以关注订阅哟👋

文章目录

添加控制：条件扩散模型

准备数据

训练模型

采样

提高效率：潜在扩散

稳定扩散：深度组件

文本编码器

Classifier-free guidance

VAE

The UNet

将它们放在一起：带注释的采样循环

文本到图像模型 (TBD) 的训练数据

开放数据，开放模型

概括

我介绍了扩散模型和迭代优化的基本思想。到此，我们可以生成图像，但训练模型非常耗时，而且我们无法控制生成的图像。在本文中，我们将看到如何从这个模型过渡到可以根据文本描述有效生成图像的文本条件模型，并以称为稳定扩散 (SD) 的模型作为案例研究。不过，在介绍 SD 之前，我们将首先了解条件模型的工作原理，并回顾导致我们今天拥有的文本到图像模型的一些创新。

添加控制：条件扩散模型

在我们处理从文本描述生成图像的问题（一项非常具有挑战性的任务！）之前，让我们先关注一些稍微容易一些的事情。我们将看到如何将我们的模型输出导向特定类型或类别的图像。我们可以使用一种称为conditioning的方法，其想法是要求模型不仅生成任何图像，而且生成属于预定义类的图像。

模型调节是一个简单但有效的想法。我们将从使用的相同扩散模型开始，仅做一些更改。首先，我们将使用一个名为 Fashion MNIST 的新数据集来代替蝴蝶，以便我们可以轻松识别类别。然后，至关重要的是，我们将通过模型运行两个输入。我们不仅要向它展示真实图像的样子，还要告诉它每张图像所属的类别。我们希望该模型能够学会关联图像和标签，从而了解毛衣、靴子等的独特特征。

请注意，我们对解决分类问题不感兴趣——我们不希望模型在给定输入图像的情况下告诉我们类别——。我们仍然希望它执行相同的任务，即：请生成看起来像是来自该数据集的合理图像。唯一的区别是我们给它提供了关于这些图像的额外信息。我们将使用相同的损失函数和训练策略，因为它与之前的任务相同。

准备数据

我们需要一个包含不同图像组的数据集。用于计算机视觉分类任务的数据集非常适合此目的。我们可以从像 ImageNet 数据集这样的东西开始，它包含 1000 个类别的数百万张图像。然而，在这个数据集上训练模型需要非常长的时间。处理新问题时，通常最好先从较小的数据集开始，以确保一切按预期进行。这使得反馈循环很短，因此我们可以快速迭代并确保我们走在正确的轨道上。

对于这个例子，我们可以像那样选择 MNIST。为了让事情稍有不同，我们将选择 Fashion MNIST。Fashion MNIST 由 Zalando 开发并开源，是 MNIST 的替代品，具有一些相同的特征：紧凑的尺寸、黑白图像和 10 个类别。主要区别在于，类不是数字，而是对应于不同类型的衣服，并且图像包含比简单的手写数字更多的细节。

让我们看一些例子。

from datasets import load_dataset

fashion_mnist = load_dataset("fashion_mnist")

clothes = fashion_mnist["train"]["image"][:8]
classes = fashion_mnist["train"]["label"][:8]
show_images(clothes, titles=classes, figsize=(4,2.5))

所以类0表示 T 恤，2 类是毛衣，9 类表示靴子。以下是 Fashion MNIST 中 10 个类别的列表： fashion_mnist · Datasets at Hugging Face。我们准备数据集和数据加载器的方式与第 3 章中的方式类似，主要区别在于我们还将类信息作为输入。在这种情况下，我们不会调整大小，而是将我们的图像输入（大小为28 × 28像素）填充到32 × 32，就像我们在第 3 章中所做的那样。

preprocess = transforms.Compose([
    transforms.RandomHorizontalFlip(),   # Randomly flip (data augmentation)
    transforms.ToTensor(),               # Convert to tensor (0, 1)
    transforms.Pad(2),                   # Add 2 pixels on all sides
    transforms.Normalize([0.5], [0.5]),  # Map to (-1, 1)
])

要使用此模型进行预测，我们必须将类标签作为附加输入传递给forward方法：

x = torch.randn((1, 1, 32, 32))
with torch.no_grad():
    out = model(x, timestep=7, class_labels=torch.tensor([2])).sample
out.shape

torch.Size([1, 1, 32, 32])

笔记

您会注意到我们还将其他一些东西作为条件传递给模型：时间步长！没错，即使是第 3 章中的模型也可以认为是条件扩散模型！我们以时间步为条件对其进行调节，希望了解我们在扩散过程中的进展情况将有助于它生成更逼真的图像。

在内部，时间步长和类标签都变成了模型在前向传播过程中使用的嵌入。在整个 UNet 的多个阶段，这些嵌入被投影到与给定层中的通道数相匹配的维度上，然后被添加到该层的输出中。这意味着条件信息被馈送到 UNet 的每个块，为模型提供充分的机会来学习如何有效地使用它。

训练模型

添加噪点在灰度图像上的效果与在第 3 章中的蝴蝶上一样好。

scheduler = DDPMScheduler(num_train_timesteps=1000, beta_start=0.0001, beta_end=0.02)
timesteps = torch.linspace(0, 999, 8).long()
batch = next(iter(train_dataloader))
x = batch['images'][:8]
noise = torch.rand_like(x)
noised_x = scheduler.add_noise(x, noise, timesteps)
show_images((noised_x*0.5 + 0.5).clip(0, 1))

我们的训练循环也几乎与第 3 章中的完全相同，只是我们现在传递类标签以进行调节。请注意，这只是模型的附加信息，但它不会以任何方式影响我们的损失函数。

for step, batch in enumerate(train_dataloader):
        # Load the input images
        clean_images = batch["images"].to(device)
        class_labels = batch["labels"].to(device)

        # *Sample noise to add to the images*
        # *Sample a random timestep for each image*
        # *Add noise to the clean images according to the timestep*

        # Get the model prediction for the noise - note the use of class_labels
        noise_pred = model(noisy_images, timesteps, class_labels=class_labels, return_dict=False)[0]

        # *Calculate the loss and update the parameters as before*
        ...

在本例中，我们训练了 25 个 epoch——完整代码可以在补充材料中找到。

采样

现在我们有了一个模型，它在进行预测时需要两个输入：图像和类别标签。我们可以通过从随机噪声开始然后迭代去噪来创建样本，传入我们想要生成的任何类标签：

def generate_from_class(class_to_generate, n_samples=8):
    sample = torch.randn(n_samples, 1, 32, 32).to(device)
    class_labels = [class_to_generate] * n_samples
    class_labels = torch.tensor(class_labels).to(device)

    for i, t in tqdm(enumerate(scheduler.timesteps)):
        # Get model pred
        with torch.no_grad():
            noise_pred = model(sample, t, class_labels=class_labels).sample

        # Update sample with step
        sample = scheduler.step(noise_pred, t, sample).prev_sample

    return sample.clip(-1, 1)*0.5 + 0.5

# Generate t-shirts (class 0)
images = generate_from_class(0)
show_images(images, nrows=2)

1000it [00:21, 47.25it/s]

# Now generate some sneakers (class 7)
images = generate_from_class(7)
show_images(images, nrows=2)

1000it [00:21, 47.20it/s]

# ...or boots (class 9)images = generate_from_class(9)show_images(images,nrows=2)

1000it [00:21, 47.26it/s]

如您所见，生成的图像远非完美。如果我们探索架构并训练更长时间，它们可能会变得更好。但令人惊奇的是，该模型不仅学习了不同类型服装的形状，而且还意识到形状9看起来与形状0不同，只需将此信息与训练数据一起发送即可。换一种稍微不同的方式来说：模特习惯于看到 9 号球鞋。当我们要求它生成图像并提供9 时，它会以启动响应。我们已经成功构建了一个类条件模型，能够生成以来自 fasionMNIST 的类标签为条件的图像！

提高效率：潜在扩散

现在我们可以训练一个条件模型，我们需要做的就是放大它并以文本而不是类标签为条件，对吧？好吧，不完全是。随着图像大小的增加，处理这些图像所需的计算能力也会增加。这在称为自注意力的操作中尤为明显，其中操作量随输入数量呈二次方增长。一个 128px 的正方形图像的像素是 64px 正方形图像的 4 倍，因此需要 16x（即 4个2个) 在自我注意层中的内存和计算。对于任何想要生成高分辨率图像的人来说，这都是一个问题！

图 2-1。Latent Diffusion Models 论文中介绍的架构。请注意左侧的 VAE 编码器和解码器，用于在像素空间和潜在空间之间进行转换

潜在扩散试图通过使用称为变分自动编码器 (VAE) 的单独模型来缓解此问题。正如我们在第 2 章中看到的，VAE 可以将图像压缩到更小的空间维度。这背后的基本原理是图像往往包含大量冗余信息——给定足够的训练数据，VAE 有望学习生成输入图像的更小表示，然后基于这个小的潜在表示重建图像高保真度。SD 中使用的 VAE 接收 3 通道图像并生成 4 通道潜在表示，每个空间维度的缩减因子为 8。也就是说，512 像素的正方形输入图像将被压缩为 4x64x64 的潜在图像。

通过在这些较小的潜在表示而不是全分辨率图像上应用扩散过程，我们可以获得使用较小图像带来的许多好处（更低的内存使用量、UNet 中需要的层数更少、生成时间更快……）和一旦我们准备好查看它，仍然将结果解码回高分辨率图像。这项创新大大降低了训练和运行这些模型的成本。介绍这个想法的论文（High-Resolution Image Synthesis with Latent Diffusion ModelsRombach 等人）通过训练以分割图、类标签和文本为条件的模型证明了该技术的强大功能。令人印象深刻的结果促使作者与 RunwayML、LAION 和 EleutherAI 等合作伙伴进一步合作，以训练该模型的更强大版本，该模型后来成为 Stable Diffusion。

稳定扩散：深度组件

Stable Diffusion 是一种基于文本的潜在扩散模型。由于它的流行，有数百个网站和应用程序可以让您使用它来创建图像，而无需任何技术知识。它也得到像 之类的库的很好支持diffusers，这让我们可以使用用户友好的管道对 SD 图像进行采样：

pipe("Watercolor illustration of a rose").images[0]

  0%| | 0/50 [00:00<?, ?it/s]

在本节中，我们将探讨使这成为可能的所有组件。

文本编码器

那么 Stable Diffusion 是如何理解文本的呢？早些时候，我们展示了如何向 UNet 提供额外的信息，使我们能够对生成的图像类型进行一些额外的控制。给定图像的噪声版本，该模型的任务是根据其他线索（例如类别标签）预测去噪版本。在 SD 的情况下，额外的线索是文本提示。在推理时，我们可以输入我们希望看到的图像的描述和一些纯噪声作为起点，模型会尽力将随机输入降噪为与说明相匹配的内容。

图 2-2。文本编码器将输入字符串转换为文本嵌入，并将其与时间步长和噪声潜伏一起馈送到 UNet 中。

为此，我们需要创建文本的数字表示，以捕获有关其描述内容的相关信息。为此，SD 利用基于 CLIP 的预训练转换器模型，该模型也在第 2 章中介绍。文本编码器是一个转换器模型，它接收一系列标记并为每个标记生成一个 1024 维向量（0r 768 维（在我们用于本节演示的 SD 版本 1 的情况下）。我们没有将这些向量组合成一个表示，而是将它们分开并将它们用作 UNet 的条件。这允许 UNet 单独使用每个标记中的信息，而不仅仅是整个提示的整体含义。因为我们从 CLIP 模型的内部表示中提取这些文本嵌入，它们通常被称为“编码器隐藏状态”。图 3 显示了文本编码器架构。

图 2-3。显示文本编码过程的图表，该过程将输入提示转换为一组文本嵌入（encoder_hidden_​​states），然后可以将其作为条件输入到 UNet。

编码文本的第一步是遵循称为标记化的过程。这会将字符序列转换为数字序列，其中每个数字代表一组不同的字符。通常一起出现的字符（如最常见的单词）可以分配一个表示整个单词或组的标记。长的或复杂的词，或有许多变形的词，可能会被翻译成多个标记，其中每个标记通常代表单词的一个有意义的部分。

没有单一的“最佳”分词器；相反，每种语言模型都有自己的模型。差异在于支持的标记数量和标记化策略——我们是使用单个字符，正如我们刚刚描述的，还是我们应该考虑不同的原始单位。在以下示例中，我们将看到短语的标记化如何与 Stable Diffusion 的标记器一起工作。我们句子中的每个单词都被分配了一个唯一的标记号（例如，photograph恰好8853在标记器的词汇表中）。还有一些额外的标记用于提供额外的上下文，例如句子结束的地方。

prompt = 'A photograph of a puppy'

# Turn the text into a sequence of tokens:
text_input = pipe.tokenizer(prompt, padding="max_length",
                            max_length=pipe.tokenizer.model_max_length,
                            truncation=True, return_tensors="pt")

# See the individual tokens
for t in text_input['input_ids'][0][:8]: # We'll just look at the first 7
    print(t, pipe.tokenizer.decoder.get(int(t)))

tensor(49406) <|startoftext|>
tensor(320) a</w>
tensor(8853) photograph</w>
tensor(539) of</w>
tensor(320) a</w>
tensor(6829) puppy</w>
tensor(49407) <|endoftext|>
tensor(49407) <|endoftext|>

一旦文本被标记化，我们就可以将其传递给文本编码器，以获得将被送入 UNet 的最终文本嵌入：

# Grab the output embeddings
text_embeddings = pipe.text_encoder(text_input.input_ids.to(device))[0]
print('Text embeddings shape:', text_embeddings.shape)

Text embeddings shape: torch.Size([1, 77, 768])

我们将在关注转换器模型的章节中更详细地介绍转换器模型如何处理一串标记。

Classifier-free guidance

事实证明，即使付出了所有努力使文本条件尽可能有用，模型在进行预测时仍然倾向于默认主要依赖嘈杂的输入图像而不是提示。在某种程度上，这是有道理的——许多标题与其关联的图像只是松散相关，因此模型学会了不要过分依赖描述！然而，当需要生成新图像时，这是不可取的——如果模型不遵循提示，那么我们可能会得到与我们的描述完全无关的图像。

图 2-4。从提示“An oil painting of a collie in a top hat”生成的图像，CFG 比例为 0、1、2 和 10（从左到右）

为了解决这个问题，我们使用了一种称为无分类器指导 (CGF) 的技巧。在训练期间，文本条件有时会保持空白，迫使模型学习对没有任何文本信息的图像进行去噪（无条件生成）。然后在推理时，我们做出两个独立的预测：一个以文本提示作为条件，一个没有。然后，我们可以使用这两个预测之间的差异来创建最终的组合预测，根据某个比例因子（指导比例）进一步推动文本条件预测所指示的方向，希望产生更好匹配的图像迅速的。图 4 显示了不同指导比例下提示的输出 - 如您所见，较高的值会生成与描述更匹配的图像。

VAE

VAE 的任务是将图像压缩成更小的潜在表示，然后再压缩回来。与 Stable Diffusion 一起使用的 VAE 是一个真正令人印象深刻的模型。我们不会在这里详细介绍训练细节，但除了第 2 章中描述的常见重建损失和 KL 散度之外，他们还使用了额外的基于补丁的鉴别器损失来帮助模型学习输出合理的细节和纹理。这在训练中添加了类似 GAN 的组件，并有助于避免以前 VAE 中常见的轻微模糊输出。与文本编码器一样，VAE 通常单独训练，并在扩散模型训练和采样过程中用作冻结组件。

图 2-5。用VAE编解码和图像

让我们加载一张图片，看看它被 VAE 压缩和解压后的样子：

# NB, this will be our own image as part of the supplementary material to avoid external URLs
im = load_image('https://images.pexels.com/photos/14588602/pexels-photo-14588602.jpeg', size=(512, 512))
show_image(im);

# Encode the image
with torch.no_grad():
    tensor_im = transforms.ToTensor()(im).unsqueeze(0).to(device)*2-1
    latent = vae.encode(tensor_im.half()) # Encode the image to a distribution
    latents = latent.latent_dist.sample() # Sampling from the distribution
    latents = latents * 0.18215 # This scaling factor was introduced by the SD authors to reduce the variance of the latents

latents.shape

torch.Size([1, 4, 64, 64])

可视化低分辨率的潜在表示，我们可以看到输入图像的一些粗略结构在不同的通道中仍然可见：

# Plot the individual channels of the latent representation
show_images([l for l in latents[0]], titles=[f'Channel {i}' for i in range(latents.shape[1])], ncols=4)

并解码回图像空间，我们得到与原始图像几乎相同的输出图像。您看得出来差别吗？

# Decode the image
with torch.no_grad():
    image = vae.decode(latents / 0.18215).sample
image = (image / 2 + 0.5).clamp(0, 1)
show_image(image[0].float());

从头开始生成图像时，我们会创建一组随机的潜在变量作为起点。我们迭代地细化这些嘈杂的潜伏以生成样本，然后使用 VAE 解码器将这些最终潜伏解码为我们可以查看的图像。只有当我们想从现有图像开始处理时才使用编码器，我们将在第 5 章中探讨这一点。

The UNet

稳定扩散中使用的 UNet 与我们在第 3 章中用于生成图像的 UNet 有点相似。我们没有采用 3 通道图像作为输入，而是采用 4 通道潜在图像。时间步嵌入的输入方式与本章开头示例中的类条件相同。但是这个 UNet 还需要接受文本嵌入作为附加条件。分散在整个 UNet 中的是交叉注意力层。UNet 中的每个空间位置都可以处理 文本条件中的不同标记，从提示中引入相关信息。图 7 中的图表显示了如何在不同点输入此文本条件（以及基于时间步长的条件）。

图 2-6。稳定扩散 UNet

UNet for Stable Diffusion 版本 1 和 2 具有大约 8.6 亿个参数。最近的 SD XL 甚至更多，大约（细节待定），大部分附加参数是通过残差块中的附加通道（N 与原始中的 1280）和附加变换器块在较低分辨率阶段添加的。

注意：Stable Diffusion XL 尚未公开发布，因此此部分将在更多信息公开时更新。

将它们放在一起：带注释的采样循环

现在我们知道每个组件的作用，让我们将它们放在一起以生成图像而不依赖管道。以下是我们将使用的设置：

# Some settings
prompt = ["Acrylic palette knife painting of a flower"] # What we want to generate
height = 512                        # default height of Stable Diffusion
width = 512                         # default width of Stable Diffusion
num_inference_steps = 30            # Number of denoising steps
guidance_scale = 7.5                # Scale for classifier-free guidance
seed = 42                           # Seed for random number generator

第一步是对文本提示进行编码。因为我们计划进行无分类器指导，所以我们实际上将创建两组文本嵌入：一组带有提示，另一组表示空字符串。您还可以对否定提示进行编码以代替空字符串，或者将具有不同权重的多个提示组合在一起，但这是最常见的用法：

# Tokenize the input
text_input = pipe.tokenizer(prompt, padding="max_length", max_length=pipe.tokenizer.model_max_length, truncation=True, return_tensors="pt")

# Feed through the text encoder
with torch.no_grad():
    text_embeddings = pipe.text_encoder(text_input.input_ids.to(device))[0]

# Do the same for the unconditional input (a blank string)
uncond_input = pipe.tokenizer("", padding="max_length", max_length=pipe.tokenizer.model_max_length, return_tensors="pt")
with torch.no_grad():
    uncond_embeddings = pipe.text_encoder(uncond_input.input_ids.to(device))[0]

# Concatenate the two sets of text embeddings embeddings
text_embeddings = torch.cat([uncond_embeddings, text_embeddings])

接下来我们创建随机初始潜伏并设置调度程序以使用所需的推理步骤数：

# Prepare the Scheduler
pipe.scheduler.set_timesteps(num_inference_steps)

# Prepare the random starting latents
latents = torch.randn(
    (1, pipe.unet.in_channels, height // 8, width // 8), # Shape of the latent representation
    generator=torch.manual_seed(32),  # Seed the random number generator
).to(device).half()
latents = latents * pipe.scheduler.init_noise_sigma

现在我们遍历采样步骤，在每个阶段获得模型预测并使用它来更新潜在：

# Sampling loop
for i, t in enumerate(pipe.scheduler.timesteps):

    # Create two copies of the latents to match the two text embeddings (unconditional and conditional)
    latent_model_input = torch.cat([latents] * 2)
    latent_model_input = pipe.scheduler.scale_model_input(latent_model_input, t)

    # predict the noise residual for both sets of inputs
    with torch.no_grad():
        noise_pred = pipe.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample

    # Split the prediction into unconditional and conditional versions:
    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)

    # perform classifier-free guidance
    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

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

注意无分类器指导步骤。我们最终的噪声预测是 noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)，将预测从无条件预测推向基于提示做出的预测。尝试更改制导比例以查看这对输出有何影响。

在循环结束时，潜伏现在应该代表一个与提示相匹配的合理图像。最后一步是使用 VAE 将 latents 解码为图像，以便我们可以看到结果：

# scale and decode the image latents with vae
latents = 1 / 0.18215 * latents
with torch.no_grad():
    image = vae.decode(latents).sample
image = (image / 2 + 0.5).clamp(0, 1)

# Display
show_image(image[0].float());

如果您浏览StableDiffusionPipeline的源代码，您会发现上面的代码与管道使用的调用方法非常匹配。希望这个带注释的版本表明幕后没有什么太神奇的事情发生！当我们遇到为这个基础添加额外技巧的额外管道时，使用它作为参考。

文本到图像模型 (TBD) 的训练数据

注意：我们可能会在这里添加一个更深入的部分，其中包含 LAION 如何聚集在一起的历史和技术细节，以及围绕从互联网上收集的公共数据进行培训的一些细微差别和争论。

开放数据，开放模型

LAION-5B 数据集包括从互联网上收集的超过 50 亿个图像说明对。该数据集是由开源社区创建并为开源社区创建的，他们看到了对此类可公开访问的数据集的需求。在 LAION 计划之前，只有少数大公司的研究实验室可以访问此类数据。这些组织将其私有数据集的详细信息保密，这使得他们的结果无法验证或复制。通过创建一个公开可用的训练数据源，LAION 使一波较小的社区和组织能够训练模型并进行研究，否则这些是不可能的。

图 2-7。“艺术创造力的爆发”——作者使用稳定扩散生成的图像

Stable Diffusion 就是这样一种模型，它是在 LAION 的一个子集上训练的，这是发明了潜在扩散模型的研究人员与一个名为 Stability AI 的组织之间合作的一部分。训练像 SD 这样的模型需要大量的 GPU 时间。即使有免费提供的 LAION 数据集，也没有多少人能负担得起这项投资。这就是为什么模型权重和代码的公开发布如此重要——它标志着一个强大的文本到图像模型首次向所有人提供，该模型具有与最佳闭源替代方案相似的功能。Stable Diffusion 的公开可用性使其成为过去一年寻求探索该技术的研究人员和开发人员的首选。数百篇论文建立在基本模型之上，添加新功能或寻找创新方法来提高其速度和质量。无数的初创公司已经找到了将这些快速改进的工具集成到他们的产品中的方法，从而催生了一个完整的新应用程序生态系统。

Stable Diffusion 推出后的几个月展示了公开分享这些技术的影响。SD 不是最好的文本到图像模型，但它是我们大多数人可以访问的最好的模型，因此成千上万的人花时间让它变得更好，并在这个开放的基础上进行构建。希望这个例子能鼓励其他人效仿并在未来与开源社区分享他们的工作！

概括

在本文中，我们了解了调节如何为我们提供新的方法来控制扩散模型生成的图像。我们已经看到潜在扩散如何让我们更有效地训练扩散模型。我们已经了解了如何使用文本编码器根据文本提示来调节扩散模型，从而实现强大的文本到图像功能。我们通过深入研究采样循环并查看不同组件如何协同工作，探索了所有这些如何在稳定扩散模型中结合在一起。