Transformer 和扩散模型的生成式 AI 实用指南（预览版）

第一章：扩散模型

在 2020 年末，一个名为扩散模型的鲜为人知的模型类别开始在机器学习领域引起轰动。研究人员找出了如何使用这些模型生成比以前技术产生的合成图像质量更高的图像。随后出现了一系列论文，提出了改进和修改，进一步提高了质量。到 2021 年底，出现了像 GLIDE 这样的模型，展示了在文本到图像任务上令人难以置信的结果，几个月后，这些模型已经进入了主流，如 DALL-E 2 和 Stable Diffusion 等工具，使任何人都可以通过输入所需看到的文本描述来生成图像。

在本章中，我们将深入了解这些模型的工作原理。我们将概述使它们如此强大的关键见解，使用现有模型生成图像，以了解它们的工作方式，然后训练我们自己的模型，以进一步加深对此的理解。该领域仍在快速发展，但本章涵盖的主题应该为您打下坚实的基础。第五章将通过一个名为 Stable Diffusion 的模型，探索更高级的技术，第六章将探讨这些技术在简单图像生成之外的应用。

关键见解：迭代细化

那么，是什么使扩散模型如此强大呢？先前的技术，如 VAEs 或 GANs，通过模型的单次前向传递生成最终输出。这意味着模型必须在第一次尝试时就做到一切正确。如果它犯了一个错误，就无法返回并修复它。另一方面，扩散模型通过迭代多个步骤生成其输出。这种“迭代细化”允许模型纠正之前步骤中的错误，并逐渐改进输出。为了说明这一点，让我们看一个扩散模型的示例。

我们可以使用 Hugging Face 扩散器库加载预训练模型。该管道可用于直接创建图像，但这并不能显示出底层发生了什么：

# Load the pipeline
image_pipe = DDPMPipeline.from_pretrained("google/ddpm-celebahq-256")
image_pipe.to(device);

# Sample an image
image_pipe().images[0]

我们可以逐步重新创建采样过程，以更好地了解模型生成图像时发生了什么。我们使用随机噪声初始化我们的样本 x，然后通过模型运行 30 步。在右侧，您可以看到模型对特定步骤的最终图像预测-请注意，最初的预测并不特别好！我们不是直接跳到最终预测的图像，而是只在预测的方向上稍微修改 x（显示在左侧）。然后，我们再次将这个新的、稍微更好的 x 通过模型进行下一步的处理，希望能产生稍微改进的预测，这可以用来进一步更新 x，依此类推。经过足够的步骤，模型可以生成一些令人印象深刻的逼真图像。

# The random starting point for a batch of 4 images
x = torch.randn(4, 3, 256, 256).to(device)

# Set the number of timesteps lower
image_pipe.scheduler.set_timesteps(num_inference_steps=30)

# Loop through the sampling timesteps
for i, t in enumerate(image_pipe.scheduler.timesteps):

    # Get the prediction given the current sample x and the timestep t
    with torch.no_grad():
        noise_pred = image_pipe.unet(x, t)["sample"]

    # Calculate what the updated sample should look like with the scheduler
    scheduler_output = image_pipe.scheduler.step(noise_pred, t, x)

    # Update x
    x = scheduler_output.prev_sample

    # Occasionally display both x and the predicted denoised images
    if i % 10 == 0 or i == len(image_pipe.scheduler.timesteps) - 1:
        fig, axs = plt.subplots(1, 2, figsize=(12, 5))

        grid = torchvision.utils.make_grid(x, nrow=4).permute(1, 2, 0)
        axs[0].imshow(grid.cpu().clip(-1, 1) * 0.5 + 0.5)
        axs[0].set_title(f"Current x (step {i})")

        pred_x0 = scheduler_output.pred_original_sample
        grid = torchvision.utils.make_grid(pred_x0, nrow=4).permute(1, 2, 0)
        axs[1].imshow(grid.cpu().clip(-1, 1) * 0.5 + 0.5)
        axs[1].set_title(f"Predicted denoised images (step {i})")
        plt.show()

image

image

image

image

注意

如果那段代码看起来有点吓人-我们将在本章中解释这一切是如何工作的。现在，只需专注于结果。

学习如何逐渐改进“损坏”的输入的核心思想可以应用于广泛的任务范围。在本章中，我们将专注于无条件图像生成-也就是说，生成类似训练数据的图像，而不对这些生成的样本的外观进行额外的控制。扩散模型也已应用于音频、视频、文本等。虽然大多数实现使用我们将在这里介绍的某种“去噪”方法的变体，但正在出现利用不同类型的“损坏”以及迭代细化的新方法，这可能会使该领域超越目前专注于去噪扩散的焦点。令人兴奋的时刻！

训练扩散模型

在本节中，我们将从头开始训练一个扩散模型，以更好地了解它们的工作原理。我们将首先使用 Hugging Face diffusers 库中的组件。随着本章的进行，我们将逐渐揭开每个组件的工作原理。与其他类型的生成模型相比，训练扩散模型相对简单。我们反复进行：

• 从训练数据中加载一些图像。
• 以不同的量添加噪音。记住，我们希望模型能够很好地估计如何“修复”（去噪）极其嘈杂的图像和接近完美的图像。
• 将输入的嘈杂版本馈送到模型中。
• 评估模型在去噪这些输入时的表现。
• 使用这些信息来更新模型权重。

要使用经过训练的模型生成新图像，我们从完全随机的输入开始，并通过模型重复地将其馈送到模型中，在每次迭代中根据模型预测更新输入的小量。正如我们将看到的，有许多采样方法试图简化这个过程，以便我们可以尽可能少的步骤生成好的图像。

数据

对于这个例子，我们将使用来自 Hugging Face Hub 的图像数据集-具体来说，这个包含 1000 张蝴蝶图片的集合。在项目部分，您将看到如何使用自己的数据。

dataset = load_dataset("huggan/smithsonian_butterflies_subset", split="train")

在使用这些数据训练模型之前，我们需要做一些准备。图像通常表示为一个“像素”网格，每个像素有三个颜色通道（红色、绿色和蓝色）的颜色值在 0 到 255 之间。为了处理这些图像并使它们准备好进行训练，我们需要： - 将它们调整为固定大小 - （可选）通过随机水平翻转来添加一些增强，有效地使我们的数据集大小加倍 - 将它们转换为 PyTorch 张量（表示颜色值为 0 到 1 之间的浮点数） - 将它们标准化为具有均值为 0 的值，值在-1 到 1 之间

我们可以使用torchvision.transforms来完成所有这些操作：

image_size = 64

# Define data augmentations
preprocess = transforms.Compose(
    [
        transforms.Resize((image_size, image_size)),  # Resize
        transforms.RandomHorizontalFlip(),  # Randomly flip (data augmentation)
        transforms.ToTensor(),  # Convert to tensor (0, 1)
        transforms.Normalize([0.5], [0.5]),  # Map to (-1, 1)
    ]
)

接下来，我们需要创建一个数据加载器，以便加载应用了这些转换的数据批次：

batch_size = 32

def transform(examples):
    images = [preprocess(image.convert("RGB")) for image in examples["image"]]
    return {"images": images}

dataset.set_transform(transform)

train_dataloader = torch.utils.data.DataLoader(
    dataset, batch_size=batch_size, shuffle=True
)

我们可以通过加载单个批次并检查图像来检查这是否有效。

batch = next(iter(train_dataloader))
print('Shape:', batch['images'].shape,
      '\nBounds:', batch['images'].min().item(), 'to', batch['images'].max().item())
show_images(batch['images'][:8]*0.5 + 0.5) # NB: we map back to (0, 1) for display
Shape: torch.Size([32, 3, 64, 64])
Bounds: -0.9921568632125854 to 1.0

image

添加噪音

我们如何逐渐破坏我们的数据？最常见的方法是向图像添加噪音。我们添加的噪音量由噪音时间表控制。不同的论文和方法以不同的方式处理这个问题，我们将在本章后面进行探讨。现在，让我们看看一种常见的方法，基于 Ho 等人的论文“去噪扩散概率模型”。在扩散器库中，添加噪音是由称为调度器的东西处理的，它接收一批图像和一个“时间步长”列表，并确定如何创建这些图像的嘈杂版本：

scheduler = DDPMScheduler(num_train_timesteps=1000, beta_start=0.001, beta_end=0.02)
timesteps = torch.linspace(0, 999, 8).long()

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))

image

在训练过程中，我们将随机选择时间步长。调度器接收一些参数（beta_start 和 beta_end），它用这些参数来确定给定时间步长应该存在多少噪音。我们将在第 X 节中更详细地介绍调度器。

UNet

UNet 是一种卷积神经网络，用于诸如图像分割之类的任务，其中期望的输出与输入具有相同的空间范围。它由一系列“下采样”层组成，用于减小输入的空间大小，然后是一系列“上采样”层，用于再次增加输入的空间范围。下采样层通常也后面跟着一个“跳跃连接”，将下采样层的输出连接到上采样层的输入。这允许上采样层“看到”网络早期的高分辨率表示，这对于具有图像样式输出的任务特别有用，其中这些高分辨率信息尤为重要。

扩散库中使用的 UNet 架构比 2015 年 Ronneberger 等人提出的原始 UNet更先进，增加了注意力和残差块等功能。我们稍后会更仔细地看一下，但这里的关键特点是它可以接受一个输入（嘈杂的图像）并产生一个形状相同的预测（预测的噪音）。对于扩散模型，UNet 通常还接受时间步作为额外的条件，我们将在 UNet 深入探讨部分再次探讨这一点。

这是我们如何创建一个 UNet 并将我们的一批嘈杂图像输入其中的方法：

# Create a UNet2DModel
model = UNet2DModel(
    in_channels=3,  # 3 channels for RGB images
    sample_size=64,  # Specify our input size
    block_out_channels=(64, 128, 256, 512), # N channels per layer
    down_block_types=("DownBlock2D", "DownBlock2D",
                      "AttnDownBlock2D", "AttnDownBlock2D"),
    up_block_types=("AttnUpBlock2D", "AttnUpBlock2D",
                    "UpBlock2D", "UpBlock2D"),
)

# Pass a batch of data through
with torch.no_grad():
    out = model(noised_x, timestep=timesteps).sample
out.shape
torch.Size([8, 3, 64, 64])

注意输出与输入的形状相同，这正是我们想要的。

训练

现在我们的模型和数据准备好了，我们可以开始训练。我们将使用学习率为 3e-4 的 AdamW 优化器。对于每个训练步骤，我们：

• 加载一批图像。
• 向图像添加噪音，选择随机时间步来确定添加多少噪音。
• 将嘈杂图像输入模型。
• 计算损失，即模型预测与目标之间的均方误差 - 在这种情况下是我们添加到图像中的噪音。这被称为噪音或“epsilon”目标。您可以在第 X 节中找到有关不同训练目标的更多信息。
• 反向传播损失并使用优化器更新模型权重。

在代码中，所有这些看起来是这样的：

num_epochs = 50 # How many runs through the data should we do?
lr = 1e-4 # What learning rate should we use
model = model.to(device) # The model we're training (defined in the previous section)
optimizer = torch.optim.AdamW(model.parameters(), lr=lr) # The optimizer
losses = [] # somewhere to store the loss values for later plotting

# Train the model (this takes a while!)
for epoch in range(num_epochs):
    for step, batch in enumerate(train_dataloader):

        # Load the input images
        clean_images = batch["images"].to(device)

        # Sample noise to add to the images
        noise = torch.randn(clean_images.shape).to(clean_images.device)

        # Sample a random timestep for each image
        timesteps = torch.randint(
            0,
            scheduler.num_train_timesteps,
            (clean_images.shape[0],),
            device=clean_images.device,
        ).long()

        # Add noise to the clean images according timestep
        noisy_images = scheduler.add_noise(clean_images, noise, timesteps)

        # Get the model prediction for the noise
        noise_pred = model(noisy_images, timesteps, return_dict=False)[0]

        # Compare the prediction with the actual noise:
        loss = F.mse_loss(noise_pred, noise)

        # Store the loss for later plotting
        losses.append(loss.item())

        # Update the model parameters with the optimizer based on this loss
        loss.backward(loss)
        optimizer.step()
        optimizer.zero_grad()

在 GPU 上运行上述代码大约需要一个小时，所以在等待时喝杯茶或者减少时代的数量。训练后损失曲线如下所示：

# Plot the loss curve:
plt.plot(losses);

image

随着模型学习去噪图像，损失曲线呈下降趋势。由于根据每次迭代中随机时间步的随机抽样向图像添加不同数量的噪音，曲线相当嘈杂。仅通过观察噪音预测的均方误差很难判断这个模型是否能够很好地生成样本，所以让我们继续下一节，看看它的表现如何。

采样

扩散库使用了“管道”的概念，将生成扩散模型样本所需的所有组件捆绑在一起：

pipeline = DDPMPipeline(unet=model, scheduler=scheduler)
ims = pipeline(batch_size=4).images
show_images(ims, nrows=1)

image

当然，将创建样本的工作交给管道并不能真正展示出我们正在进行的工作。因此，这里有一个简单的采样循环，展示了模型如何逐渐改进输入图像：

# Random starting point (4 random images):
sample = torch.randn(4, 3, 64, 64).to(device)

for i, t in enumerate(scheduler.timesteps):

    # Get model pred
    with torch.no_grad():
        noise_pred = model(sample, t).sample

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

show_images(sample.clip(-1, 1)*0.5 + 0.5, nrows=1)

image

这与我们在本章开头使用的代码相同，用来说明迭代改进的概念，但希望现在你对这里发生的事情有了更好的理解。我们从一个完全随机的输入开始，然后在一系列步骤中由模型进行改进。每一步都是对输入的小更新，基于模型对该时间步的噪音的预测。我们仍然在pipeline.scheduler.step()的调用背后抽象了一些复杂性 - 在后面的章节中，我们将更深入地探讨不同的采样方法以及它们的工作原理。

评估

可以使用 FID 分数（Fréchet Inception Distance）来评估生成模型的性能。FID 分数通过比较从预训练的神经网络中提取的特征图之间的统计数据，衡量生成样本与真实样本的相似程度。分数越低，给定模型生成的图像的质量和逼真度就越好。FID 分数因其能够提供对不同类型生成网络的“客观”比较指标而受到欢迎，而无需依赖人类判断。

尽管 FID 分数很方便，但也有一些重要的注意事项需要注意：

• 给定模型的 FID 分数取决于用于计算它的样本数量，因此在模型之间进行比较时，我们需要确保报告的分数都是使用相同数量的样本计算的。通常做法是在这个目的上使用 50,000 个样本，尽管为了节省时间，您可能在开发过程中评估较少数量的样本，只有在准备发布结果时才进行完整评估。
• 在计算 FID 时，图像被调整为 299 像素的正方形图像。这使得它对于极低分辨率或高分辨率图像的度量变得不太有用。不同的深度学习框架处理调整大小的方式也有细微差异，这可能导致 FID 分数有小的差异！我们建议使用clean-fid这样的库来标准化 FID 计算。
• 用作 FID 特征提取器的网络通常是在 Imagenet 分类任务上训练的模型。在生成不同领域的图像时，这个模型学到的特征可能不太有用。更准确的方法是在特定领域的数据上训练一个分类网络，但这会使得在不同论文和方法之间比较分数变得更加困难，所以目前 Imagenet 模型是标准选择。
• 如果您保存生成的样本以供以后评估，格式和压缩也会影响 FID 分数。尽量避免低质量的 JPEG 图像。

即使考虑了所有这些警告，FID 分数只是质量的粗略度量，不能完美地捕捉使图像看起来更“真实”的微妙之处。因此，用它们来了解一个模型相对于另一个模型的表现，但也要看看每个模型生成的实际图像，以更好地了解它们的比较。人类偏好仍然是质量的黄金标准，最终这是一个相当主观的领域！

深入：噪音时间表

在上面的训练示例中，其中一步是“添加不同数量的噪音”。我们通过在 0 到 1000 之间选择一个随机时间步长来实现这一点，然后依靠调度程序添加适当数量的噪音。同样，在采样过程中，我们再次依靠调度程序告诉我们使用哪些时间步长以及如何根据模型预测从一个时间步长移动到下一个时间步长。事实证明，选择添加多少噪音是一个重要的设计决策，可以极大地影响给定模型的性能。在本节中，我们将看到为什么会这样，并探讨实践中使用的一些不同方法。

为什么要添加噪音？

在本章的开头，我们说扩散模型背后的关键思想是迭代改进。在训练期间，我们通过不同的方式“腐败”输入。在推断期间，我们从“最大程度上腐败”的输入开始，迭代地“去腐败”它，希望最终得到一个不错的最终结果。

到目前为止，我们已经专注于一种特定的“腐败”：添加高斯噪声。这样做的一个原因是扩散模型的理论基础 - 如果我们使用不同的腐败方法，我们在技术上就不再做“扩散”了！然而，Bansal 等人的一篇题为Cold Diffusion的论文戏剧性地表明，我们不一定需要仅仅出于理论上的便利而限制自己使用这种方法。他们表明，类似扩散模型的方法适用于许多不同的“腐败”方法（见图 1-1）。最近，像MUSE、MaskGIT和PAELLA这样的模型已经使用了随机标记屏蔽或替换作为等效的“腐败”方法，用于量化数据 - 也就是说，用离散标记而不是连续值表示的数据。

image.png

图 1-1。Cold Diffusion Paper 中使用的不同退化的示意图

尽管如此，出于几个原因，添加噪音仍然是最受欢迎的方法：

• 我们可以轻松控制添加的噪音量，从“完美”到“完全损坏”平稳过渡。这对于像减少图像分辨率这样的事情并不适用，这可能会导致“离散”的过渡。
• 我们可以有许多有效的随机起始点进行推断，不像一些方法可能只有有限数量的可能的初始（完全损坏）状态，比如完全黑色的图像或单像素图像。

所以，至少目前，我们将坚持添加噪音作为我们的损坏方法。接下来，让我们更仔细地看看如何向图像添加噪音。

开始简单

我们有一些图像（x），我们想以某种方式将它们与一些随机噪音结合起来。

x = next(iter(train_dataloader))['images'][:8]
noise = torch.rand_like(x)

我们可以线性插值（lerp）它们之间的一些量。这给我们一个函数，它在“量”从 0 到 1 变化时，从原始图像 x 平稳过渡到纯噪音：

def corrupt(x, noise, amount):
  amount = amount.view(-1, 1, 1, 1) # make sure it's broadcastable
  return x*(1-amount) + noise*amount

让我们在一批数据上看看这个过程，噪音的量从 0 到 1 变化：

amount = torch.linspace(0, 1, 8)
noised_x = corrupt(x, noise, amount)
show_images(noised_x*0.5 + 0.5)

image

这似乎正是我们想要的，从原始图像平稳过渡到纯噪音。现在，我们在这里创建了一个噪音时间表，它接受从 0 到 1 的“量”值。这被称为“连续时间”方法，我们在时间尺度上表示从 0 到 1 的完整路径。其他方法使用离散时间方法，使用一些大整数的“时间步长”来定义噪音调度器。我们可以将我们的函数封装成一个类，将连续时间转换为离散时间步长，并适当添加噪音：

class SimpleScheduler():
  def __init__(self):
    self.num_train_timesteps = 1000
  def add_noise(self, x, noise, timesteps):
    amount = timesteps / self.num_train_timesteps
    return corrupt(x, noise, amount)

scheduler = SimpleScheduler()
timesteps = torch.linspace(0, 999, 8).long()
noised_x = scheduler.add_noise(x, noise, timesteps)
show_images(noised_x*0.5 + 0.5)

image

现在我们有了一些可以直接与扩散库中使用的调度器进行比较的东西，比如我们在训练中使用的 DDPMScheduler。让我们看看它是如何比较的：

scheduler = DDPMScheduler(beta_end=0.01)
timesteps = torch.linspace(0, 999, 8).long()
noised_x = scheduler.add_noise(x, noise, timesteps)
show_images((noised_x*0.5 + 0.5).clip(0, 1))

image

数学

文献中有许多竞争的符号和方法。例如，一些论文将噪音时间表参数化为“连续时间”，其中 t 从 0（无噪音）到 1（完全损坏）- 就像我们在上一节中的corrupt函数一样。其他人使用“离散时间”方法，其中整数时间步长从 0 到某个大数 T，通常为 1000。可以像我们的SimpleScheduler类一样在这两种方法之间进行转换-只需确保在比较不同模型时保持一致。我们将在这里坚持使用离散时间方法。

深入研究数学的一个很好的起点是之前提到的 DDPM 论文。您可以在这里找到注释实现，这是一个很好的额外资源，可以帮助理解这种方法。

论文开始时指定了从时间步 t-1 到时间步 t 的单个噪音步骤。这是他们的写法：

q(t|t-1)=(t;1-βtt-1,βt).

这里β t被定义为所有时间步长 t，并用于指定每个步骤添加多少噪声。这种符号可能有点令人生畏，但这个方程告诉我们的是，更嘈杂的 t是一个分布，其均值为1 - β t t-1，方差为β t。换句话说， t是 t-1（按1 - β t缩放）和一些随机噪声的混合，我们可以将其视为按β t缩放的单位方差噪声。给定x t-1和一些噪声ϵ，我们可以从这个分布中采样得到x t：

t=1-βtt-1+βtϵ

要在时间步 t 获得嘈杂的输入，我们可以从 t=0 开始，并反复应用这一步，但这将非常低效。相反，我们可以找到一个公式一次性移动到任何时间步 t。我们定义α t = 1 - β t，然后使用以下公式：

xt=α¯tx0+1-α¯tϵ

在这里 - ϵ 是一些方差为单位的高斯噪声 - α ¯（'alpha_bar'）是直到时间t的所有α值的累积乘积。

因此，x t是x 0（由α ¯ t缩放）和ϵ（由1 - α ¯ t缩放）的混合物。在 diffusers 库中，α ¯值存储在scheduler.alphas_cumprod中。知道这一点，我们可以绘制给定调度程序的不同时间步骤中原始图像x 0和噪音ϵ的缩放因子：

plot_scheduler(DDPMScheduler()) # The default scheduler

image

我们上面的 SimpleScheduler 只是在原始图像和噪音之间线性混合，我们可以看到如果我们绘制缩放因子（相当于α ¯ t和( 1 - α ¯ t )在 DDPM 情况下）：

plot_scheduler(SimpleScheduler())

image

一个良好的噪音调度将确保模型看到不同噪音水平的图像混合。最佳选择将根据训练数据而异。可视化一些更多选项，注意：

• 将 beta_end 设置得太低意味着我们永远不会完全擦除图像，因此模型永远不会看到任何类似于用作推理起点的随机噪音的东西。
• 将 beta_end 设置得非常高意味着大多数时间步骤都花在几乎完全的噪音上，这将导致训练性能不佳。
• 不同的 beta 调度给出不同的曲线。

“余弦”调度是一个受欢迎的选择，因为它可以使原始图像平稳过渡到噪音。

fig, (ax) = plt.subplots(1, 1, figsize=(8, 5))
plot_scheduler(DDPMScheduler(beta_schedule="linear"),
               label = 'default schedule', ax=ax, plot_both=False)
plot_scheduler(DDPMScheduler(beta_schedule="squaredcos_cap_v2"),
               label = 'cosine schedule', ax=ax, plot_both=False)
plot_scheduler(DDPMScheduler(beta_end=0.003, beta_schedule="linear"),
               label = 'Low beta_end', ax=ax, plot_both=False)
plot_scheduler(DDPMScheduler(beta_end=0.1, beta_schedule="linear"),
               label = 'High beta_end', ax=ax, plot_both=False)

image

注意

这里显示的所有调度都被称为“方差保持”（VP），这意味着模型输入的方差在整个调度过程中保持接近 1。您可能还会遇到“方差爆炸”（VE）公式，其中噪音只是以不同的量添加到原始图像（导致高方差输入）。我们将在采样章节中更详细地讨论这一点。我们的 SimpleScheduler 几乎是一个 VP 调度，但由于线性插值，方差并没有完全保持。

与许多与扩散相关的主题一样，不断有新的论文探讨噪音调度的主题，因此当您阅读本文时，可能会有大量的选项可供尝试！

输入分辨率和缩放的影响

直到最近，噪音调度的一个方面大多被忽视，即输入大小和缩放的影响。许多论文在小规模数据集和低分辨率上测试潜在的调度程序，然后使用表现最佳的调度程序来训练其最终模型的较大图像。这样做的问题在于，如果我们向两个不同大小的图像添加相同数量的噪音，就会看到问题。

image

图 1-2。比较向不同大小的图像添加噪音的效果

高分辨率的图像往往包含大量冗余信息。这意味着即使单个像素被噪音遮挡，周围的像素也包含足够的信息来重建原始图像。但对于低分辨率图像来说并非如此，其中单个像素可能包含大量信息。这意味着向低分辨率图像添加相同量的噪音将导致比向高分辨率图像添加等量噪音更加损坏的图像。

这种效应在两篇独立的论文中得到了彻底的调查，这两篇论文分别于 2023 年 1 月发表。每篇论文都利用新的见解来训练能够生成高分辨率输出的模型，而无需任何以前必需的技巧。Hoogeboom 等人的《简单扩散》介绍了一种根据输入大小调整噪音计划的方法，允许在低分辨率图像上优化的计划适当地修改为新的目标分辨率。陈婷的一篇名为“关于扩散模型噪音调度的重要性”的论文进行了类似的实验，并注意到另一个关键变量：输入缩放。也就是说，我们如何表示我们的图像？如果图像表示为 0 到 1 之间的浮点数，那么它们的方差将低于噪音（通常是单位方差），因此对于给定的噪音水平，信噪比将低于如果图像表示为-1 到 1 之间的浮点数（我们在上面的训练示例中使用的方式）或其他方式。缩放输入图像会改变信噪比，因此修改这种缩放是我们在训练更大图像时可以调整的另一种方式。

深入了解：UNets 和替代方案

现在让我们来讨论真正进行重要预测的模型！回顾一下，这个模型必须能够接收嘈杂的图像并估计如何去除噪音。这需要一个能够接收任意大小的图像并输出相同大小的图像的模型。此外，该模型应能够在像素级别进行精确预测，同时也捕捉关于整个图像的更高级别信息。一个流行的方法是使用一种称为 UNet 的架构。UNet 是在 2015 年为医学图像分割而发明的，并且后来成为各种与图像相关的任务的流行选择。就像我们在上一章中看到的 AutoEncoders 和 VAEs 一样，UNets 由一系列“下采样”和“上采样”块组成。下采样块负责减小图像的大小，而上采样块负责增加图像的大小。下采样块通常由一系列卷积层组成，然后是池化或下采样层。上采样块通常由一系列卷积层组成，然后是上采样或“转置卷积”层。转置卷积层是一种特殊类型的卷积层，它增加图像的大小，而不是减小图像的大小。

常规的 AutoEncoder 或 VAE 之所以不适合这个任务，是因为它们在像素级别进行精确预测的能力较弱，因为输出必须完全从低维潜在空间重新构建。在 UNet 中，通过“跳跃连接”连接下采样和上采样块，允许信息直接从下采样块流向上采样块。这使得模型能够在像素级别进行精确预测，同时也捕捉关于整个图像的更高级别信息。

一个简单的 UNet

为了更好地理解 UNet 的结构，让我们从头开始构建一个简单的 UNet。

image.png

图 1-3。我们简单的 UNet 架构

这个 UNet 以 32px 分辨率接收单通道输入，并以 32px 分辨率输出单通道输出，我们可以用它来构建 MNIST 数据集的扩散模型。编码路径中有三层，解码路径中也有三层。每一层由卷积后跟激活函数和上采样或下采样步骤（取决于我们是否在编码或解码路径中）组成。跳过连接允许信息直接从下采样块流向上采样块，并通过将下采样块的输出添加到相应上采样块的输入来实现。一些 UNet 将下采样块的输出连接到相应上采样块的输入，并可能还在跳过连接中包含额外的层。以下是这个网络的代码：

from torch import nn

class BasicUNet(nn.Module):
    """A minimal UNet implementation."""
    def __init__(self, in_channels=1, out_channels=1):
        super().__init__()
        self.down_layers = torch.nn.ModuleList([
            nn.Conv2d(in_channels, 32, kernel_size=5, padding=2),
            nn.Conv2d(32, 64, kernel_size=5, padding=2),
            nn.Conv2d(64, 64, kernel_size=5, padding=2),
        ])
        self.up_layers = torch.nn.ModuleList([
            nn.Conv2d(64, 64, kernel_size=5, padding=2),
            nn.Conv2d(64, 32, kernel_size=5, padding=2),
            nn.Conv2d(32, out_channels, kernel_size=5, padding=2),
        ])
        self.act = nn.SiLU() # The activation function
        self.downscale = nn.MaxPool2d(2)
        self.upscale = nn.Upsample(scale_factor=2)

    def forward(self, x):
        h = []
        for i, l in enumerate(self.down_layers):
            x = self.act(l(x)) # Through the layer and the activation function
            if i < 2: # For all but the third (final) down layer:
              h.append(x) # Storing output for skip connection
              x = self.downscale(x) # Downscale ready for the next layer

        for i, l in enumerate(self.up_layers):
            if i > 0: # For all except the first up layer
              x = self.upscale(x) # Upscale
              x += h.pop() # Fetching stored output (skip connection)
            x = self.act(l(x)) # Through the layer and the activation function

        return x

在 MNIST 上使用这种架构训练的扩散模型产生以下样本（代码包含在补充材料中，这里为了简洁起见而省略）：

改进 UNet

这个简单的 UNet 适用于这个相对简单的任务，但远非理想。那么，我们可以做些什么来改进它呢？

• 添加更多参数。可以通过在每个块中使用多个卷积层，通过在每个卷积层中使用更多的滤波器，或者通过使网络更深来实现。
• 添加残差连接。使用 ResBlocks 而不是常规卷积层可以帮助模型学习更复杂的功能，同时保持训练稳定。
• 添加归一化，如批归一化。批归一化可以帮助模型更快、更可靠地学习，确保每一层的输出都围绕 0 中心，并具有标准差为 1。
• 添加正则化，如 dropout。Dropout 有助于防止模型过度拟合训练数据，这在处理较小的数据集时非常重要。
• 添加注意力。通过引入自注意力层，我们允许模型在不同时间集中关注图像的不同部分，这有助于学习更复杂的功能。类似变压器的注意力层的添加也可以增加可学习参数的数量，这有助于模型学习更复杂的功能。缺点是注意力层在更高分辨率时计算成本要比常规卷积层高得多，因此我们通常只在较低分辨率（即 UNet 中的较低分辨率块）使用它们。
• 添加一个额外的输入用于时间步长，这样模型可以根据噪音水平调整其预测。这称为时间步长调节，几乎所有最近的扩散模型都在使用。我们将在下一章中更多地了解有条件的模型。

作为对比，在 diffusers 库中使用 UNet 实现时，在 MNIST 上的结果如下，该库包含了上述所有改进：

警告

这一部分可能会在未来通过结果和更多细节进行扩展。我们还没有开始训练具有不同改进的变体！

替代架构

最近，一些替代架构已被提出用于扩散模型。这些包括：

• 变压器。Peebles 和 Xie 的 DiT 论文（“具有变压器的可扩展扩散模型”）显示了基于变压器的架构可以用于训练扩散模型，并取得了很好的结果。然而，变压器架构的计算和内存需求仍然是非常高分辨率的挑战。
• Simple Diffusion paper链接中的UViT架构旨在兼顾两者的优点，通过用一大堆变压器块替换 UNet 的中间层。该论文的一个关键见解是，将大部分计算集中在 UNet 的较低分辨率块上，可以更有效地训练高分辨率扩散模型。对于非常高的分辨率，他们使用称为小波变换的东西进行一些额外的预处理，以减少输入图像的空间分辨率，同时通过使用额外的通道尽可能保留更多信息，再次减少在更高空间分辨率上花费的计算量。
• 循环接口网络。RIN paper（Jabri 等）采用类似的方法，首先将高分辨率输入映射到更易处理和低维的“潜在”表示，然后通过一堆变压器块进行处理，然后解码回到图像。此外，RIN 论文引入了“循环”概念，其中信息从上一个处理步骤传递给模型，这对于扩散模型旨在执行的迭代改进可能是有益的。

尚不清楚基于变压器的方法是否会完全取代 UNet 作为扩散模型的首选架构，还是像 UViT 和 RIN 架构这样的混合方法将被证明是最有效的。

深入：目标和预处理

我们已经谈到扩散模型接受嘈杂的输入并“学会去噪”它。乍一看，你可能会认为网络的自然预测目标是图像的去噪版本，我们将其称为x0。然而，在代码中，我们将模型预测与用于创建嘈杂版本的单位方差噪声进行了比较（通常称为 epsilon 目标，eps）。这两者在数学上看起来是相同的，因为如果我们知道噪声和时间步长，我们可以推导出x0，反之亦然。虽然这是真的，但目标的选择对不同时间步的损失有一些微妙的影响，因此模型学习最佳去噪哪个噪声水平。为了获得一些直觉，让我们在不同的时间步上可视化一些不同的目标：

image

在极低的噪声水平下，x0目标是非常容易的，而准确预测噪声几乎是不可能的。同样，在极高的噪声水平下，eps目标是容易的，而准确预测去噪图像几乎是不可能的。两种情况都不理想，因此引入了其他目标，使模型在不同的时间步预测x0和eps的混合。“用于快速采样扩散模型的渐进蒸馏。”用于快速采样扩散模型的渐进蒸馏。“用于快速采样扩散模型的渐进蒸馏。”中引入的v目标是其中之一，它被定义为v = α ¯ · ϵ + 1 - α ¯ · x 0。Karras 等人在EDM paper中通过一个称为c_skip的参数引入了类似的想法，并将不同的扩散模型公式统一到一个一致的框架中。如果您对了解不同目标、缩放和其他不同扩散模型公式的微妙之处感兴趣，我们建议阅读他们的论文以获得更深入的讨论。

项目时间：训练您自己的扩散模型

现在您已经了解了扩散模型的基础知识，现在是时候自己训练一些了！本章的补充材料包括一个笔记本，指导您如何在自己的数据集上训练扩散模型的过程。在您进行操作时，请回顾本章，看看不同部分是如何相互配合的。笔记本还包括许多建议的更改，以更好地探索不同的模型架构和训练策略如何影响结果。

总结

在本章中，我们看到了迭代改进的想法如何应用于训练扩散模型，使其能够将噪音转化为美丽的图像。您已经看到了一些设计选择，这些选择是创建成功的扩散模型所必需的，并希望通过训练自己的模型来实践这些选择。在下一章中，我们将看看一些更先进的技术，这些技术已经被开发出来，以改进扩散模型的性能，并赋予它们非凡的新能力！

参考文献

Ho, Jonathan, Ajay Jain, and Pieter Abbeel. “Denoising diffusion probabilistic models.” Advances in Neural Information Processing Systems 33 (2020): 6840-6851.

Ronneberger, O., Fischer, P. and Brox, T., 2015. U-net: Convolutional networks for biomedical image segmentation. In Medical Image Computing and Computer-Assisted Intervention–MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18 (pp. 234-241). Springer International Publishing.

Bansal, Arpit, Eitan Borgnia, Hong-Min Chu, Jie S. Li, Hamid Kazemi, Furong Huang, Micah Goldblum, Jonas Geiping, and Tom Goldstein. “Cold diffusion: Inverting arbitrary image transforms without noise.” arXiv preprint arXiv:2208.09392 (2022).

Hoogeboom, Emiel, Jonathan Heek, and Tim Salimans. “simple diffusion: End-to-end diffusion for high resolution images.” arXiv preprint arXiv:2301.11093 (2023).

Chang, Huiwen, Han Zhang, Jarred Barber, A. J. Maschinot, Jose Lezama, Lu Jiang, Ming-Hsuan Yang et al. “Muse: Text-To-Image Generation via Masked Generative Transformers.” arXiv preprint arXiv:2301.00704 (2023).

Chang, Huiwen, Han Zhang, Lu Jiang, Ce Liu, and William T. Freeman. “Maskgit: Masked generative image transformer.” In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 11315-11325. 2022.

Rampas, Dominic, Pablo Pernias, Elea Zhong, and Marc Aubreville. “Fast Text-Conditional Discrete Denoising on Vector-Quantized Latent Spaces.” arXiv preprint arXiv:2211.07292 (2022).

Chen, Ting “On the Importance of Noise Scheduling for Diffusion Models.” arXiv preprint arXiv:2301.10972 (2023).

Peebles, William, and Saining Xie. “Scalable Diffusion Models with Transformers.” arXiv preprint arXiv:2212.09748 (2022).

Jabri, Allan, David Fleet, and Ting Chen. “Scalable Adaptive Computation for Iterative Generation.” arXiv preprint arXiv:2212.11972 (2022).

Salimans, Tim, and Jonathan Ho. “Progressive distillation for fast sampling of diffusion models.” arXiv preprint arXiv:2202.00512 (2022).)

Karras, Tero, Miika Aittala, Timo Aila, and Samuli Laine. “Elucidating the design space of diffusion-based generative models.” arXiv preprint arXiv:2206.00364 (2022).

第二章：逐步构建稳定扩散

在上一章中，我们介绍了扩散模型和迭代细化的基本思想。在章节结束时，我们可以生成图像，但训练模型非常耗时，而且我们无法控制生成的图像。在本章中，我们将看到如何从这一点过渡到基于文本描述有效生成图像的条件化模型，以 Stable Diffusion（SD）模型作为案例研究。不过，在我们深入研究 SD 之前，我们将首先看一下条件化模型的工作原理，并回顾一些导致我们今天拥有文本到图像模型的创新。

添加控制：条件化扩散模型

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

模型条件化是一个简单但有效的想法。我们将从第三章中使用的相同扩散模型开始，只做了一些改变。首先，我们将使用一个叫做 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))

image

所以类别0表示 T 恤，2是毛衣，9表示靴子。以下是 Fashion MNIST 中的 10 个类别列表：https://huggingface.co/datasets/fashion_mnist#data-fields。我们准备我们的数据集和数据加载器，类似于我们在第三章中所做的，主要的区别是我们还将类别信息作为输入包含进去。在这种情况下，我们不是调整大小，而是将我们的图像输入（大小为28×28像素）填充到32×32，就像我们在第三章中所做的那样。

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)
])
batch_size = 256

def transform(examples):
    images = [preprocess(image.convert("L")) for image in examples["image"]]
    return {"images": images, "labels": examples["label"]}

train_dataset = fashion_mnist["train"].with_transform(transform)

train_dataloader = torch.utils.data.DataLoader(
    train_dataset, batch_size=batch_size, shuffle=True
)

创建一个类别条件化模型

如果我们使用 diffusers 库中的UNet模型，我们可以提供自定义的条件信息。在这里，我们创建了一个类似于第三章中使用的模型，但我们在UNet构造函数中添加了一个num_class_embeds参数。该参数告诉模型我们想要使用类别标签作为额外的条件。我们将使用 10，因为 Fashion MNIST 中有 10 个类别。

model = UNet2DModel(
    in_channels=1,   # 1 channel for grayscale images
    out_channels=1,  # output channels must also be 1
    sample_size=32,
    block_out_channels=(32, 64, 128, 256),
    norm_num_groups=8,
    num_class_embeds=10, # Enable class conditioning
)

要使用该模型进行预测，我们必须将类别标签作为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 的每个块中，使模型有充分的机会学习如何有效地使用它。

训练模型

添加噪音对灰度图像的效果与第三章中的蝴蝶一样好。

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))

image

我们的训练循环几乎与第三章中的相同，只是现在我们传递类别标签进行条件处理。请注意，这只是模型的附加信息，但并不以任何方式影响我们的损失函数。

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 个时期的训练-完整的代码可以在补充材料中找到。

采样

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

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]

image

# Now generate some sneakers (class 7)
images = generate_from_class(7)
show_images(images, nrows=2)
1000it [00:21, 47.20it/s]

image

# ...or boots (class 9)
images = generate_from_class(9)
show_images(images, nrows=2)
1000it [00:21, 47.26it/s]

image

正如你所看到的，生成的图像远非完美。如果我们探索架构并进行更长时间的训练，它们可能会变得更好。但令人惊讶的是，该模型不仅学会了不同类型的服装形状，还意识到形状9看起来与形状0不同，只是通过在训练数据旁边发送这些信息。换句话说：该模型习惯于看到数字9伴随着靴子。当我们要求它生成一张图像并提供9时，它会回应一个靴子。我们成功地构建了一个能够生成基于 fashionMNIST 类别标签的图像的有条件模型！

提高效率：潜在扩散

现在我们可以训练一个有条件的模型，我们所需要做的就是将其扩展并以文本而不是类别标签进行条件处理，对吗？嗯，并不完全是这样。随着图像尺寸的增加，处理这些图像所需的计算能力也会增加。这在一个称为自注意力的操作中尤为明显，其中操作的数量随着输入数量的增加呈二次增长。一个 128 像素的正方形图像有 4 倍于 64 像素正方形图像的像素数量，因此在自注意力层中需要 16 倍（即4 2）的内存和计算。这对于希望生成高分辨率图像的人来说是一个问题！

image.png

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

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

通过在这些较小的潜在表示上应用扩散过程，而不是在全分辨率图像上应用，我们可以获得许多使用较小图像所带来的好处（更低的内存使用、UNet 中需要更少的层、更快的生成时间...），并且一旦准备查看结果，仍然可以将结果解码回高分辨率图像。这种创新大大降低了训练和运行这些模型的成本。引入这一想法的论文（Rombach 等人的《使用潜在扩散模型进行高分辨率图像合成》）通过训练模型展示了这一技术的威力，这些模型是基于分割地图、类标签和文本的。令人印象深刻的结果导致作者与 RunwayML、LAION 和 EleutherAI 等合作伙伴进一步合作，以训练模型的更强大版本，后来成为 Stable Diffusion。

Stable Diffusion：深入组件

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

pipe("Watercolor illustration of a rose").images[0]
0%|          | 0/50 [00:00<?, ?it/s]

image

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

文本编码器

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

simplified_unet.png

图 2-2。文本编码器将输入字符串转换为文本嵌入，这些嵌入与时间步长和嘈杂的潜在一起输入 UNet。

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

text encoder

图 2-3。显示文本编码过程的图表，将输入提示转换为一组文本嵌入（编码器隐藏状态），然后可以将其作为条件输入到 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])

我们将在关于变压器模型处理标记字符串的章节中详细介绍。

无分类器引导

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

image.png

图 2-4。从提示“戴礼帽的牧羊犬的油画”生成的图像，CFG 比例为 0、1、2 和 10（从左到右）

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

VAE

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

vae.drawio.png

图 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 解码器将这些最终潜变量解码成我们可以查看的图像。编码器仅在我们想要从现有图像开始该过程时使用，这是我们将在第五章中探讨的内容。

UNet

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

SD 图表

图 2-6. 稳定扩散 UNet

稳定扩散版本 1 和 2 的 UNet 大约有 8.6 亿个参数。更新更近期的 SD XL 拥有更多参数，大约为（详细信息待定），大部分额外参数是通过在残差块中添加额外通道（原始版本中的 N 对 1280）和添加变压器块来增加低分辨率阶段的。

注：稳定扩散 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 将潜变量解码成图像，以便我们可以看到结果：

# 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的源代码，您会发现上面的代码与管道使用的call方法非常相似。希望这个带注释的版本显示了幕后并没有太神奇的事情发生！当我们遇到添加额外技巧的其他管道时，请将其用作参考。

文本到图像模型的训练数据（待定）

注：我们可能会在这里添加一个更深入的部分，介绍 LAION 是如何形成的历史和技术细节，以及一些关于在从互联网上抓取的公共数据上进行训练的细微差别和争论。

开放数据，开放模型

LAION-5B 数据集包括从互联网上抓取的 50 亿多个图像标题对。这个数据集是由开源社区创建的，他们认为有必要有这样一个公开可访问的数据集。在 LAION 计划之前，只有少数大公司的研究实验室才能访问这样的数据。这些组织将他们的私有数据集的细节保留给自己，这使得他们的结果无法验证或复制。通过创建一个公开可用的训练数据来源，LAION 使得一波较小的社区和组织能够进行训练模型和进行研究，否则这是不可能的。

image

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

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

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

总结

在本章中，我们已经看到条件如何为我们提供了控制扩散模型生成图像的新方法。我们已经看到潜在扩散如何让我们更有效地训练扩散模型。我们已经看到文本编码器如何被用来在文本提示上对扩散模型进行条件设定，实现强大的文本到图像的能力。我们已经探讨了所有这些是如何在稳定扩散模型中结合在一起的，通过深入研究采样循环并看到不同组件如何一起工作。在下一章中，我们将展示可以添加到扩散模型（如 SD）中的许多额外功能，使它们超越简单的图像生成。在本书的第二部分中，您将学习如何微调 SD，以向模型添加新的知识或功能。