【扩散模型】DDIM从原理到实战

DDPM过程的一个问题是在训练后生成图像的速度。当然，我们可能能够生成令人惊叹的图像，但生成一张图像需要1000次模型传递。在GPU上使图像通过模型1000次可能需要几秒钟，但在CPU上需要更长的时间。我们需要一种方法来加快生成过程。

1. DDIM简介

DDIM论文介绍了一种在图像质量几乎没有tradeoff的情况下加快图像生成的方法。它通过将扩散过程重新定义为非马尔可夫过程来做到这一点。

左图是原始的DDPM论文，它需要从时间$T$到时间$T-1$的所有过去的去噪步骤来获得时间T的下一个去噪图像。DDPM被建模为马尔可夫链，这意味着在生成$T$之前的整个链之前，不能生成时间$T$的图像。

DDIM论文提出了一种使过程非马尔可夫的方法（如右图所示），允许跳过去噪过程中的步骤，而不需要在当前状态之前访问所有过去的状态。DDIM最好的部分是，它们可以在训练模型后应用，因此DDPM模型可以很容易地转换为DDIM，而无需重新训练新模型。

首先，重新定义了single step的逆向扩散过程：

注：DDIM论文使用的是不带bar的alphas，但论文中的alphas值是DDPM论文中使用的alphas bar（cumulative alpha）值。这有点令人困惑，所以我将用alpha bars替换它们的alphas，以保持符号的一致性。

首先，这个reformalization相当于DDPM论文中的formalization，但仅当方差等于${\tilde{\beta }}_t$时。

作者没有明确指出他们的sigma公式只是${\tilde{\beta }}_t$, 但只要有一点代数，你就会发现情况确实如此。

当$\sigma =\theta$，得到DDIM：

请注意，数据中没有添加noise。这就是DDIM的诀窍。当$\sigma =\odot$时，去噪过程变得完全确定，唯一的噪声是$x_0$处的原始噪声，因为在去噪过程中没有添加新的噪声。

由于逆向过程中没有噪声，因此该过程是确定性的，并且我们不再需要使用马尔可夫链，因为马尔可夫链用于概率过程。我们可以使用非马尔可夫过程，它允许我们跳过步骤。

非马尔可夫正向和逆向过程

在上图中，我们从步骤$x_3$跳到$x_1$，跳过$x_2$。作者将新的扩散过程建模为子序列$\tau$，是原始扩散序列的子集。例如，我可以在扩散过程中每隔一个扩散步骤进行采样，得到$\tau =[0，2，4，\dots ，T-2，T]$的子序列。

最后，作者使用以下公式将模型的扩散模型方差确定为DDIM和DDPM之间的插值：

当$\eta =\odot$时，由于没有噪声，扩散模型是DDIM，当$\eta =1$时，扩散模型为原始DDPM。0和1之间的任何$\eta$都是DDIM和DDPM之间的插值。

当所采取的步骤数小于原始$T$步骤时，DDIM的性能比DDPM好得多。下表显示了从0到1的$\eta$插值以及10、20、50、100和1000个生成步骤的DDPM和DDIM FID分数（对多样性和图像质量进行评分）。请注意，原始模型是在$T=1000$步上训练的。

DDIM在不同的数据集上产生不同的$\eta$值和不同的步长。

FID分数越低越好。尽管DDPM在最初的1000个步骤中表现最好，但在生成具有更少生成步骤的图像时，DDIM紧随其后。在使用DDIM时，基本上需要在图像质量和生成时间之间进行权衡，而原始DDPM没有提供这种权衡。现在我们可以用更少的步骤生成高质量的图像！

2. 实战DDIM

安装依赖库：

!pip install -Uq transformers diffusers accelerate
1

加载依赖包：

import torch
import requests
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
from io import BytesIO
from tqdm.auto import tqdm
from matplotlib import pyplot as plt
from torchvision import transforms as tfms
from diffusers import StableDiffusionPipeline, DDIMScheduler

# 判断当前GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)
1
2
3
4
5
6
7
8
9
10
11
12
13
14

定义图像加载函数：

def load_image(url, size=None):
  response = requests.get(url,timeout=1)
  img = Image.open(BytesIO(response.content)).convert('RGB')
  if size is not None:
    img = img.resize(size)
  return img
1
2
3
4
5
6

2.1 载入一个预训练过的pipeline

使用StableDiffusionPipeline加载预训练模型并配置DDIM调度器，而后对预训练模型进行一次采样。

pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5").to(device)
# 配置DDIM调度器
pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
# 采样一次，保证代码正常
prompt = 'Beautiful DSLR Photograph of a penguin on the beach, golden hour'
negative_prompt = 'blurry, ugly, stock photo'

im = pipe(prompt, negative_prompt=negative_prompt).images[0]
im.resize((256, 256))  # 调整至有利于查看的尺寸
1
2
3
4
5
6
7
8
9

采样结果：

2.2 DDIM采样

在给定时刻$t$，带有噪声的图像$x_t$是通过对原始图像$x_0$加上高斯噪声$ϵ$得到的。DDIM论文给出了$x_t$的定义式：$$x_t=\sqrt{{\alpha }_t}{x}_0+\sqrt{1-{\alpha }_t}ϵ$$
其中$ϵ$是方差归一化后的高斯噪声，${\alpha }_t$在DDPM论文中被称为$\bar{\alpha }$，并被用于定义噪声调度器。在扩散模型中，$\alpha$被计算并排序存储在scheduler.alphas_cumprod中。

# 选择使用Diffusers中的alphas_cumprod函数来得到alphas
timesteps = pipe.scheduler.timesteps.cpu()
alphas = pipe.scheduler.alphas_cumprod[timesteps]
plt.plot(timesteps, alphas, label='alpha_t');
plt.legend();
1
2
3
4
5

从中可以看出，噪声曲线（在时间步0）是从一幅无噪的干净图像开始的，此时${\alpha }_t=1$。在到达更高的时间步后，便得到一幅几乎全是噪声的图像，${\alpha }_t$也几乎下降到0。

为了计算采样轨迹中下一个时刻的值$x_{t-1}$（因为是从后向前移动的），我们

• 首先需要得到预测噪声${ϵ}_{\theta }(x_t)$(这是模型的输出)，
• 然后用它预测不带噪声的图像$x_0$。
• 接下来，朝着“反转”的方向移动一步。
• 最后，可以加上一些带有${\sigma }_t$系数的额外噪声

在DDIM论文原文中与上述操作相关的内容是：

翻译如下：
根据公式(10)中的$p_{\theta }(x_{1:T})$，可以通过公式(12)从$x_t$推导出$x_{t-1}$，其中${ϵ}_t\sim \mathcal{N}(0,\mathrm{I})$是独立于$x_t$的标准高斯噪声，并且定义${\alpha }_0=1$，使用不同的$\alpha$值会导致不同的生成流程，因为同时使用了相同的模型${ϵ}_{\theta }$，所以不需要重新训练模型。对于所有时刻$t$，当${\theta }_t=\sqrt{(1-{\alpha }_{t-1})/(1-{\alpha }_t)}\sqrt{1-{\alpha }_t/{\alpha }_{t-1}}$时，前向过程将变成马尔可夫过程，生成过程变为DDPM。

另一个特殊情况是，即对于几乎所有时刻（t=1除外）的${\sigma }_t=0$，前向过程在给定$x_{t-1}$和$x_0$的情况下变得更加确定；在生成过程中，随机噪声${ϵ}_t$前面的系数变为0。得到的模型变成隐式概率模型（Mohamed & Lakshminarayanan， 2016），其中的样本是根据固定的过程从隐变量生成的（从$x_r$到$x_0$）。将这个模型命名为“去噪扩散隐式模型”（Denoising Diffusion Implicit Model, DDIM），因为它是一个使用DDPM目标进行训练的隐式概率模型（尽管前向过程不再是扩散过程）。

因此接下来的示例不需要再额外添加噪声，即可实现完全确定的DDIM采样：

# 采样噪声（标准DDIM采样）
@torch.no_grad()
def sample(
    prompt, 
    start_step=0, 
    start_latents=None, 
    guidance_scale=3.5, 
    num_inference_steps=30, 
    num_images_per_prompt=1, 
    do_classifier_free_guidance=True,
    negative_prompt='',
    device=device
    ):
  # 对文本提示语进行编码
  text_embeddings = pipe._encode_prompt(prompt, device, num_images_per_prompt, do_classifier_free_guidance, negative_prompt)
  # 设置推理的步数
  pipe.scheduler.set_timesteps(num_inference_steps, device=device)
  # 创建随机起点
  if start_latents is None:
    start_latents = torch.randn(1, 4, 64, 64, device=device)
    start_latents *= pipe.scheduler.init_noise_sigma
  
  latents = start_latents.clone()
  for i in tqdm(range(start_step, num_inference_steps)):
    t = pipe.scheduler.timesteps[i]
  
    # 如果正在进行CFG，则对隐层进行扩展
    latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
    latent_model_input = pipe.scheduler.scale_model_input(latent_model_input, t)

    # 预测噪声
    noise_pred = pipe.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample

    # 进行引导
    if do_classifier_free_guidance:
      noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
      noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
    
    # 使用调度器更新步骤
    # Normally we'd rely on the scheduler to handle the update step:
    # latents = pipe.scheduler.step(noise_pred, t, latents).prev_sample

    # 现在不用调度器，而是自行实现
    prev_t = max(1, t.item()-(1000//num_inference_steps)) # t-1
    alpha_t = pipe.scheduler.alphas_cumprod[t.item()]
    alpha_t_prev = pipe.scheduler.alphas_cumprod[prev_t]
    predicted_x0 = (latents - (1-alpha_t).sqrt()*noise_pred) / alpha_t.sqrt()
    direction_pointing_to_xt = (1-alpha_t_prev).sqrt()*noise_pred
    latents = alpha_t_prev.sqrt()*predicted_x0 + direction_pointing_to_xt
  # 后处理
  images = pipe.decode_latents(latents)
  images = pipe.numpy_to_pil(images)

  return images
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54

生成一张图片：

prompt = 'Watercolor painting of a beach sunset'
sample(prompt, negative_prompt=negative_prompt, num_inference_steps=50)[0].resize((256, 256))
1
2

2.3 反转（invert）

反转的目标是“颠倒”采样的过程。 最终想得到“带噪”的隐式表示，如果将其用作正常采样过程的起点，那么生成的将是原始图像。
图片示例：

input_image = load_image('https://images.pexels.com/photos/8306128/pexels-photo-8306128.jpeg', size=(512, 512))
input_image
1
2

使用一个包含无分类器引导的文本提示语来进行反转操作。

定义invert()函数：

# 定义invert函数
@torch.no_grad()
def invert(
    start_latents, 
    prompt, 
    guidance_scale=3.5, 
    num_inference_steps=80,
    num_images_per_prompt=1,
    do_classifier_free_guidance=True,
    negative_prompt='',
    device=device):
  # 对提示文本进行编码
  text_embeddings = pipe._encode_prompt(prompt, device, num_images_per_prompt, do_classifier_free_guidance, negative_prompt)
  # 指定起点
  latents = start_latents.clone()
  # 用一个列表保存反转的隐层
  intermediate_latents = []
  # 设置推理的步数
  pipe.scheduler.set_timesteps(num_inference_steps, device=device)
  # 反转的时间步
  timesteps = reversed(pipe.scheduler.timesteps)

  for i in tqdm(range(1, num_inference_steps), total=num_inference_steps-1):
    # 跳过最后一次迭代
    if i >= num_inference_steps - 1: continue
    t = timesteps[i]
    # 如果正在进行CFG，则对隐层进行扩展
    latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
    latent_model_input = pipe.scheduler.scale_model_input(latent_model_input, t)
    # 预测残留的噪声
    noise_pred = pipe.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample
    # 引导
    if do_classifier_free_guidance:
      noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
      noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
    current_t = max(0, t.item() - (1000//num_inference_steps)) # t
    next_t = t # min(999, t.item() + (1000//num_inference_steps)) # t+1
    alpha_t = pipe.scheduler.alphas_cumprod[current_t]
    alpha_t_next = pipe.scheduler.alphas_cumprod[next_t]

    # 反转的更新步（重新排列更新步，利用当前隐层得到新的隐层）
    latents = (latents - (1-alpha_t).sqrt() * noise_pred) * (alpha_t_next.sqrt() / alpha_t.sqrt()) + (1-alpha_t_next).sqrt() * noise_pred
    # 保存隐层
    intermediate_latents.append(latents)
  return torch.cat(intermediate_latents)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45

invert函数与上文中的sample函数非常相似，但是在时间步上是朝着相反的方向移动的：从t=0开始，向噪声更多的方向移动，而不是在更新隐式层的过程中那样噪声越来越少。可以利用预测的噪声来撤回一步更新操作，并从t移动到t+1。

将invert函数应用于示例图片，得到在反转的过程中的一系列隐式表达：

inverted_latents = invert(l, input_image_prompt,num_inference_steps=50)
inverted_latents.shape  # torch.Size([48, 4, 64, 64])
1
2

最终的隐式表达：

# 解码反转的最后一个隐层
with torch.no_grad():
    im = pipe.decode_latents(inverted_latents[-1].unsqueeze(0))
pipe.numpy_to_pil(im)[0]
1
2
3
4

将其作为起点噪声，通过常规调用方法call，将反转隐式地传递给pipeline：

# 可以通过常规调用方法，将反转隐层传递给管线
pipe(input_image_prompt, latents=inverted_latents[-1][None], num_inference_steps=50, guidance_scale=3.5).images[0]
1
2

显然，这并不是最初的那张照片。这是因为DDIM反转需要一个重要的假设——在时刻 t 预测的噪声与在时刻 t+1 预测的噪声相同，但这个假设在反转50步或者100步时是不成立的。

当然，既可以使用更多的时间步来得到更准确的反转，也可以采取“作弊”的方式，直接从相应反转过程50步中的第20步的隐式表达开始，代码如下：

# 从第20步的隐式表示开始，得到的结果距离最初的图片很近了！
start_step=20
sample(input_image_prompt, start_latents=inverted_latents[-(start_step+1)][None], 
       start_step=start_step, num_inference_steps=50)[0]
1
2
3
4

显然，得到的结果与最初的图片很接近。

但是，为什么要这么做呢？因为现在想用一个新的文件提示语来生成图片。想要得到一张除了与提示语相关以外，其他内容都与原始图片大致相同的图片。例如：把小狗换成小猫

# 把小狗换成小猫，从第10步的隐式表示开始
start_step=10
new_prompt = input_image_prompt.replace('puppy', 'cat')
sample(new_prompt, start_latents=inverted_latents[-(start_step+1)][None], 
       start_step=start_step, num_inference_steps=50)[0]
1
2
3
4
5

还有一个问题是，为什么不直接使用Img2Img管线呢？或者是为什么要做反转？为什么不直接对输入图像添加噪声，然后用新的文本提示语直接“去噪”呢？当然可以，但是这会导致图片变化非常大或者图片没什么变化。如下所示：

start_step=10
num_inference_steps=50
pipe.scheduler.set_timesteps(num_inference_steps)
noisy_1 = pipe.scheduler.add_noise(l, torch.randn_like(l), pipe.scheduler.timesteps[start_step])
sample(new_prompt, start_latents=noisy_1, start_step=start_step, num_inference_steps=num_inference_steps)[0]
1
2
3
4
5

这时的草地发生了明显的变化。

3. 组合封装

将所有代码封装到一个函数中，输入一张图片和两个文本提示语，得到一张通过反转得到的修改后的图片：

def edit(input_image, input_image_prompt, edit_prompt, num_steps=100, start_step=30, guidance_scale=3.5):
  with torch.no_grad(): 
    latent = pipe.vae.encode(tfms.functional.to_tensor(input_image).unsqueeze(0).to(device)*2-1)
  l = 0.18215 * latent.latent_dist.sample()
  inverted_latents = invert(l, input_image_prompt,num_inference_steps=num_steps)
  final_im = sample(edit_prompt, start_latents=inverted_latents[-(start_step+1)][None], 
                    start_step=start_step, 
                    num_inference_steps=num_steps, 
                    guidance_scale=guidance_scale)[0]
  return final_im
1
2
3
4
5
6
7
8
9
10

示例1：

edit(input_image, 'A puppy on the grass', 'an old grey dog on the grass', num_steps=50, start_step=10)
1

示例2：

face = load_image('https://images.pexels.com/photos/1493111/pexels-photo-1493111.jpeg', size=(512, 512))
face
1
2

给原始图片增加一幅眼镜：

edit(face, 'A photograph of a face', 'A photograph of a face with sunglasses', num_steps=250, start_step=30, guidance_scale=3.5)
1

更多的迭代能够得到更好的表现，可以多试几次。

edit(face, 'A photograph of a face', 'Acrylic palette knife painting of a face, colorfull', num_steps=250, start_step=65, guidance_scale=5.5)
1

推荐阅读：Null-text Inversion for Editing Real Images using Guided Diffusion Models——一个基于DDIM来优化空文本（无条件文本提示语）的反转过程。

参考资料

1. Diffusion Models — DDPMs, DDIMs, and Classifier Free Guidance
2. Denoising Diffusion Implicit Models
3. Null-text Inversion：基于Null Prompt Finetuning的图像编辑技术