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

准备数据

我们需要一个包含不同图像组的数据集。用于计算机视觉分类任务的数据集非常适合此目的。我们可以从像 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])

在内部，时间步长和类标签都变成了模型在前向传播过程中使用的嵌入。在整个 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 的类标签为条件的图像！

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

文本编码器

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

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

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

没有单一的“最佳”分词器；相反，每种语言模型都有自己的模型。差异在于支持的标记数量和标记化策略——我们是使用单个字符，正如我们刚刚描述的，还是我们应该考虑不同的原始单位。在以下示例中，我们将看到短语的标记化如何与 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

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

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

VAE

VAE 的任务是将图像压缩成更小的潜在表示，然后再压缩回来。与 Stable Diffusion 一起使用的 VAE 是一个真正令人印象深刻的模型。我们不会在这里详细介绍训练细节，但除了第 2 章中描述的常见重建损失和 KL 散度之外，他们还使用了额外的基于补丁的鉴别器损失来帮助模型学习输出合理的细节和纹理。这在训练中添加了类似 GAN 的组件，并有助于避免以前 VAE 中常见的轻微模糊输出。与文本编码器一样，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 中的图表显示了如何在不同点输入此文本条件（以及基于时间步长的条件）。

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

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

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