DDIM从原理到实战
- 1. DDIM简介
- 2. 实战DDIM
- 2.1 载入一个预训练过的pipeline
- 2.2 DDIM采样
- 2.3 反转(invert)
- 3. 组合封装
- 参考资料
DDPM过程的一个问题是在训练后生成图像的速度。当然,我们可能能够生成令人惊叹的图像,但生成一张图像需要1000次模型传递。在GPU上使图像通过模型1000次可能需要几秒钟,但在CPU上需要更长的时间。我们需要一种方法来加快生成过程。
1. DDIM简介
DDIM论文介绍了一种在图像质量几乎没有tradeoff的情况下加快图像生成的方法。它通过将扩散过程重新定义为非马尔可夫过程来做到这一点。
左图是原始的DDPM论文,它需要从时间 T T T到时间 T − 1 T-1 T−1的所有过去的去噪步骤来获得时间T的下一个去噪图像。DDPM被建模为马尔可夫链,这意味着在生成 T T T之前的整个链之前,不能生成时间 T T T的图像。
DDIM论文提出了一种使过程非马尔可夫的方法(如右图所示),允许跳过去噪过程中的步骤,而不需要在当前状态之前访问所有过去的状态。DDIM最好的部分是,它们可以在训练模型后应用,因此DDPM模型可以很容易地转换为DDIM,而无需重新训练新模型。
首先,重新定义了single step的逆向扩散过程:
注:DDIM论文使用的是不带bar的alphas,但论文中的alphas值是DDPM论文中使用的alphas bar(cumulative alpha)值。这有点令人困惑,所以我将用alpha bars替换它们的alphas,以保持符号的一致性。
首先,这个reformalization相当于DDPM论文中的formalization,但仅当方差等于 β ~ t \tilde{\beta}_t β~t时。
作者没有明确指出他们的sigma公式只是 β ~ t \tilde{\beta}_t β~t, 但只要有一点代数,你就会发现情况确实如此。
当 σ = θ \sigma=\theta σ=θ,得到DDIM:
请注意,数据中没有添加noise。这就是DDIM的诀窍。当 σ = ⊙ \sigma=\odot σ=⊙时,去噪过程变得完全确定,唯一的噪声是 x 0 x_0 x0处的原始噪声,因为在去噪过程中没有添加新的噪声。
由于逆向过程中没有噪声,因此该过程是确定性的,并且我们不再需要使用马尔可夫链,因为马尔可夫链用于概率过程。我们可以使用非马尔可夫过程,它允许我们跳过步骤。
非马尔可夫正向和逆向过程
在上图中,我们从步骤 x 3 x_3 x3跳到 x 1 x_1 x1,跳过 x 2 x_2 x2。作者将新的扩散过程建模为子序列 τ \tau τ,是原始扩散序列的子集。例如,我可以在扩散过程中每隔一个扩散步骤进行采样,得到 τ = [ 0 , 2 , 4 , … , T − 2 , T ] \tau=[0,2,4,…,T-2,T] τ=[0,2,4,…,T−2,T]的子序列。
最后,作者使用以下公式将模型的扩散模型方差确定为DDIM和DDPM之间的插值:
当 η = ⊙ \eta =\odot η=⊙时,由于没有噪声,扩散模型是DDIM,当 η = 1 \eta =1 η=1时,扩散模型为原始DDPM。0和1之间的任何 η \eta η都是DDIM和DDPM之间的插值。
当所采取的步骤数小于原始 T T T步骤时,DDIM的性能比DDPM好得多。下表显示了从0到1的 η \eta η插值以及10、20、50、100和1000个生成步骤的DDPM和DDIM FID分数(对多样性和图像质量进行评分)。请注意,原始模型是在 T = 1000 T=1000 T=1000步上训练的。
DDIM在不同的数据集上产生不同的 η \eta η值和不同的步长。
FID分数越低越好。尽管DDPM在最初的1000个步骤中表现最好,但在生成具有更少生成步骤的图像时,DDIM紧随其后。在使用DDIM时,基本上需要在图像质量和生成时间之间进行权衡,而原始DDPM没有提供这种权衡。现在我们可以用更少的步骤生成高质量的图像!
2. 实战DDIM
安装依赖库:
!pip install -Uq transformers diffusers accelerate
加载依赖包:
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)
定义图像加载函数:
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
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)) # 调整至有利于查看的尺寸
采样结果:
2.2 DDIM采样
在给定时刻 t t t,带有噪声的图像 x t x_t xt是通过对原始图像 x 0 x_0 x0加上高斯噪声 ϵ \epsilon ϵ得到的。DDIM论文给出了 x t x_t xt的定义式: x t = α t x 0 + 1 − α t ϵ x_t=\sqrt{\alpha_t}x_0+\sqrt{1-\alpha_t}ϵ xt=αtx0+1−αtϵ
其中 ϵ ϵ ϵ是方差归一化后的高斯噪声, α t \alpha_t α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();
从中可以看出,噪声曲线(在时间步0)是从一幅无噪的干净图像开始的,此时 α t = 1 \alpha_t=1 αt=1。在到达更高的时间步后,便得到一幅几乎全是噪声的图像, α t \alpha_t αt也几乎下降到0。
为了计算采样轨迹中下一个时刻的值 x t − 1 x_{t-1} xt−1(因为是从后向前移动的),我们
- 首先需要得到预测噪声 ϵ θ ( x t ) ϵ_θ(x_t) ϵθ(xt)(这是模型的输出),
- 然后用它预测不带噪声的图像 x 0 x_0 x0。
- 接下来,朝着“反转”的方向移动一步。
- 最后,可以加上一些带有 σ t σ_t σt系数的额外噪声
在DDIM论文原文中与上述操作相关的内容是:
翻译如下:
根据公式(10)中的 p θ ( x 1 : T ) p_{\theta}(x_{1:T}) pθ(x1:T),可以通过公式(12)从 x t x_t xt推导出 x t − 1 x_{t-1} xt−1,其中 ϵ t ∼ N ( 0 , I ) \epsilon _t\sim \mathcal{N}(0,\mathrm{I} ) ϵt∼N(0,I)是独立于 x t x_t xt的标准高斯噪声,并且定义 α 0 = 1 \alpha _0=1 α0=1,使用不同的 α \alpha α值会导致不同的生成流程,因为同时使用了相同的模型 ϵ θ \epsilon _\theta ϵθ,所以不需要重新训练模型。对于所有时刻 t t t,当 θ t = ( 1 − α t − 1 ) / ( 1 − α t ) 1 − α t / α t − 1 \theta_t=\sqrt{(1-\alpha _{t-1})/(1-\alpha _t)}\sqrt{1-\alpha _t/\alpha _{t-1}} θt=(1−αt−1)/(1−αt)1−αt/αt−1时,前向过程将变成马尔可夫过程,生成过程变为DDPM。
另一个特殊情况是,即对于几乎所有时刻(t=1除外)的 σ t = 0 \sigma_t=0 σt=0,前向过程在给定 x t − 1 x_{t-1} xt−1和 x 0 x_0 x0的情况下变得更加确定;在生成过程中,随机噪声 ϵ t \epsilon_t ϵt前面的系数变为0。得到的模型变成隐式概率模型(Mohamed & Lakshminarayanan, 2016),其中的样本是根据固定的过程从隐变量生成的(从 x r x_r xr到 x 0 x_0 x0)。将这个模型命名为“去噪扩散隐式模型”(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_sigmalatents = 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 latentslatent_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-1alpha_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_predlatents = alpha_t_prev.sqrt()*predicted_x0 + direction_pointing_to_xt# 后处理images = pipe.decode_latents(latents)images = pipe.numpy_to_pil(images)return images
生成一张图片:
prompt = 'Watercolor painting of a beach sunset'
sample(prompt, negative_prompt=negative_prompt, num_inference_steps=50)[0].resize((256, 256))
2.3 反转(invert)
反转的目标是“颠倒”采样的过程。 最终想得到“带噪”的隐式表示,如果将其用作正常采样过程的起点,那么生成的将是原始图像。
图片示例:
input_image = load_image('https://images.pexels.com/photos/8306128/pexels-photo-8306128.jpeg', size=(512, 512))
input_image
使用一个包含无分类器引导的文本提示语来进行反转操作。
定义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: continuet = timesteps[i]# 如果正在进行CFG,则对隐层进行扩展latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latentslatent_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)) # tnext_t = t # min(999, t.item() + (1000//num_inference_steps)) # t+1alpha_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)
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])
最终的隐式表达:
# 解码反转的最后一个隐层
with torch.no_grad():im = pipe.decode_latents(inverted_latents[-1].unsqueeze(0))
pipe.numpy_to_pil(im)[0]
将其作为起点噪声,通过常规调用方法call,将反转隐式地传递给pipeline:
# 可以通过常规调用方法,将反转隐层传递给管线
pipe(input_image_prompt, latents=inverted_latents[-1][None], num_inference_steps=50, guidance_scale=3.5).images[0]
显然,这并不是最初的那张照片。这是因为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]
显然,得到的结果与最初的图片很接近。
但是,为什么要这么做呢?因为现在想用一个新的文件提示语来生成图片。想要得到一张除了与提示语相关以外,其他内容都与原始图片大致相同的图片。例如:把小狗换成小猫
# 把小狗换成小猫,从第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]
还有一个问题是,为什么不直接使用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]
这时的草地发生了明显的变化。
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:
edit(input_image, 'A puppy on the grass', 'an old grey dog on the grass', num_steps=50, start_step=10)
示例2:
face = load_image('https://images.pexels.com/photos/1493111/pexels-photo-1493111.jpeg', size=(512, 512))
face
给原始图片增加一幅眼镜:
edit(face, 'A photograph of a face', 'A photograph of a face with sunglasses', num_steps=250, start_step=30, guidance_scale=3.5)
更多的迭代能够得到更好的表现,可以多试几次。
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)
推荐阅读:Null-text Inversion for Editing Real Images using Guided Diffusion Models——一个基于DDIM来优化空文本(无条件文本提示语)的反转过程。
参考资料
- Diffusion Models — DDPMs, DDIMs, and Classifier Free Guidance
- Denoising Diffusion Implicit Models
- Null-text Inversion:基于Null Prompt Finetuning的图像编辑技术