Diffusion打卡笔记3-StableDiffusion

---

前言

学习了扩散模型第三章关于StableDiffusion的内容，下面做一个简单的记录。

---

一、介绍

上面的图像是通过StableDiffusion生成的，可以看出生成效果还是很好的。
从定义上来看，StableDiffusion是一个强大的文本条件潜在扩散模型，即可以从文本描述中生成极佳的图像。所以，要想了解StableDiffusion具体原理，就要逐一进行单独剖析，下面将分别介绍。

1.1潜在扩散（Latent Diffusion)

之所以要进行这项操作的原因是因为要处理大规模图像数据的需要。当图像大小增加时，所需要处理图像的计算操作也在不断增加。这尤其在Self-attention模块比较明显，因为是自己和自己做计算，所以计算操作量随着输入数量的二次增长而增长。128px 方形图像的像素数是 64px 方形图像的 4 倍，因此需要 16 倍的内存和自注意力层的计算。对于任何想要生成高分辨率图像的人来说，这都是一个问题！
所以这时候针对此问题就提出了潜在扩散的方法。潜在扩散通过使用称为变分自动编码器 （VAE） 的单独模型将图像压缩到较小的空间维度来帮助缓解此问题。这背后的基本原理是，在给定足够的训练数据下，图像往往包含大量冗余信息，VAE有望学会生成输入图像的更小的表示，然后基于这种小的潜在表示以高度保真度重建图像。SD 中使用的 VAE 接收 3 通道图像并生成 4 通道潜在表示，每个空间维度的折减系数为 8。也就是说，一个 512px 的正方形输入图像将被压缩为 4x64x64 的潜在图像。以下是论文中的潜在扩散模型图。

这样如果将扩散过程即加噪处理过程应用于潜在表示向量而不是全分辨率的图像，就可以降低内存使用量、减少Unet的层数以及更快生成图像等。在准备好查看最终结果时再解码可以生成高分辨率图像，这大大降低了训练的成本。

1.2文本条件反射(Text Conditioning)

上图也显示了如何添加条件反射的过程，即打卡笔记2提到的如何向UNet提供附加信息，对生成图像类型进行额外的控制。给定图像的嘈杂版本，该模型的任务是根据其他线索（例如类标签）或图像的文本描述（在稳定扩散的情况下）预测去噪版本。在推理时，我们可以输入我们希望看到的图像的描述和一些纯噪声作为起点，模型会尽最大努力将随机输入“去噪”为与标题匹配的内容。
要将所谓的文本描述提示输入到网络中，首先要做的是将输入提示转换为一组文本嵌入（encoder_hidden_states）的文本编码，然后可以将其作为条件输入到 UNet，下图展示了这一过程。具体过程是:1、将输入序列进行词元化。2、进行词嵌入，降低特征维度。3、进行位置编码，防止自注意力机制输入顺序调换后生成结果不变。4、输入网络进行注意力计算，最后输出当前隐藏层的状态向量。（Transformer每一层都进行）

而在SD里面其采用了基于CLIP的预训练transformer。CLIP的文本编码器旨在将图像标题处理成可用于比较图像和文本的形式，因此它非常适合从图像描述中创建有用的表示。
首先还是对输入提示进行词元化（基于大型词汇表，其中每个单词或子单词都被分配了一个特定的词元），然后通过 CLIP 文本编码器馈送，为每个标记生成一个 768 维（在 SD 1.X 的情况下）或 1024 维（SD 2.X）向量。为了保持一致，提示总是被填充/截断为 77 个标记长，因此我们用作条件的最终表示是每个提示形状为 77x1024 的张量。下图显示了这一具体过程。

可以看出条件信息是以交叉注意力的形式输入到Unet当中的。UNet 中的每个空间位置都可以“参与”文本条件中的不同标记，从而从提示中引入相关信息。上图显示了这种文本条件反射（以及基于时间步长的条件反射）是如何在不同点输入的。

1.3无分类器指导

事实证明，即使付出了所有努力使文本条件尽可能有用，模型在进行预测时仍然倾向于默认主要依赖嘈杂的输入图像而不是提示。在某种程度上，这是因为许多提示词只与其关联的图像松散相关，因此模型学会了不要过分依赖描述！然而，当需要生成新图像时，这是不可取的 - 如果模型不遵循提示，那么我们可能会得到与我们的描述完全无关的图像。
为了解决这个问题，可以使用一种称为无分类器指导 （CFG） 的技巧。在训练过程中，文本条件有时保持空白，迫使模型学习对没有任何文本信息的图像进行降噪（无条件生成）。
在推理时，我们做出两个单独的预测：一个使用文本提示作为条件，另一个没有。然后，我们可以利用这两个预测之间的差异来创建一个最终的组合预测，该预测根据某个比例因子（指导比例）进一步朝着文本条件预测指示的方向推进，希望产生与提示更匹配的图像。下图显示了不同指导比例下的提示输出 - 可见，引导值越高，图像越符合描述。

1.4其他类型的调节（超分辨率、修复和深化图像）

除了支持文本条件反射以外，StableDiffusion还支持其他其他条件反射，例如，Depth-to-Image 模型具有额外的输入通道，可以获取有关去噪图像的深入信息，在推理时，我们可以输入目标图像的深度图（使用单独的模型估计），以期生成具有相似整体结构的图像。
下图即为通过深度条件反射的SD生成的具有相同整体结构的不同图像。

以类似的方式，我们可以输入低分辨率图像作为条件反射，并让模型生成高分辨率版本（如StableDiffusion Upscaler）。最后，我们可以输入一个蒙版，显示图像的一个区域，作为“绘画中”任务的一部分重新生成，其中非蒙版区域需要保持完整，同时为蒙版区域生成新内容。
此外还可以运用微调技术。DreamBooth 是一种微调文本到图像模型以“教”它一个新概念的技术，例如特定的对象或样式。该技术最初是为谷歌的 Imagen 模型开发的，但很快就被调整为StableDiffusion使用。结果可能非常令人印象深刻（如果你最近在社交媒体上看到任何人有 AI 个人资料图片，那么它很可能来自基于 dreambooth 的服务）。下面是基于Imagen 模型的 dreambooth 项目页面。

二、StableDiffusion实践

下面是关于StableDiffusion文生图等方面的一些实践。

2.1模块安装

%pip install -Uq diffusers ftfy accelerate

# Installing transformers from source for now since we need the latest version for Depth2Img
%pip install -Uq git+https://github.com/huggingface/transformers 

import torch
import requests
from PIL import Image
from io import BytesIO
from matplotlib import pyplot as plt

# We'll be exploring a number of pipelines today!
from diffusers import (
    StableDiffusionPipeline, 
    StableDiffusionImg2ImgPipeline,
    StableDiffusionInpaintPipeline, 
    StableDiffusionDepth2ImgPipeline
    )       

# We'll use a couple of demo images later in the notebook
def download_image(url):
    response = requests.get(url)
    return Image.open(BytesIO(response.content)).convert("RGB")

# Download images for inpainting example
img_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png"
mask_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png"

init_image = download_image(img_url).resize((512, 512))
mask_image = download_image(mask_url).resize((512, 512))
     

# Set device
device = (
    "mps"
    if torch.backends.mps.is_available()
    else "cuda"
    if torch.cuda.is_available()
    else "cpu"
)

2.2 文本生成图像

首先加载一个StableDiffusion pipeline观察下作用。

# Load the pipeline
model_id = "stabilityai/stable-diffusion-2-1-base"
pipe = StableDiffusionPipeline.from_pretrained(model_id).to(device)

一些减少GPU中RAM使用量的方法，防止内存不足：
1、加载FP16版本，与此同时调整张量为torch.float16:

pipe = StableDiffusionPipeline.from_pretrained(model_id, revision="fp16", torch_dtype=torch.float16).to(device)

2、启用注意力切片，减少GPU内存使用量，但速度会降低。

pipe.enable_attention_slicing()

3、减少生成图像的大小。
加载完毕后，使用包含提示词的代码生成图像：

# Set up a generator for reproducibility
generator = torch.Generator(device=device).manual_seed(42)

# Run the pipeline, showing some of the available arguments
pipe_output = pipe(
    prompt="Palette knife painting of an autumn cityscape", # What to generate
    negative_prompt="Oversaturated, blurry, low quality", # What NOT to generate
    height=480, width=640,     # Specify the image size
    guidance_scale=8,          # How strongly to follow the prompt
    num_inference_steps=35,    # How many steps to take
    generator=generator        # Fixed random seed
)

# View the resulting image
pipe_output.images[0]

可以进行不同参数的调整生成不同风格的图像，关键注意点如下：
1、宽度和高度指定生成图像的大小。它们必须能被 8 整除才能使 VAE 正常工作
2、步数会影响生成质量。默认值 （50） 效果很好，但在某些情况下，只需 20 个步骤即可完成，
3、否定提示在无分类器指导过程中使用，可能是添加其他控件的有用方法。
4、参数 guidance_scale 确定无分类器引导 （CFG） 的强度。更高的比例会推动生成的图像更好地匹配提示，但如果比例过高，结果可能会变得过度饱和和令人不快。
下面通过调整guidance scale生成三幅不同的狗狗图片进行比较：

#@markdown comparing guidance scales:
cfg_scales = [1.1, 8, 12] #@param
prompt = "A collie with a pink hat" #@param
fig, axs = plt.subplots(1, len(cfg_scales), figsize=(16, 5))
for i, ax in enumerate(axs):
  im = pipe(prompt, height=480, width=480,
    guidance_scale=cfg_scales[i], num_inference_steps=35,
    generator=torch.Generator(device=device).manual_seed(42)).images[0]
  ax.imshow(im); ax.set_title(f'CFG Scale {cfg_scales[i]}');

效果如下：

直观可以看出当把CFG scale在8-12范围内会得到更好的效果。

2.3 管道组件(Pipeline Components)

下面将具体介绍StableDiffusionPipeline,这比之前打卡笔记2使用的DDPMPipeline要相对复杂一些。除了Unet和scheduler，还包含了其他组件：

print(list(pipe.components.keys())) # List components

['vae', 'text_encoder', 'tokenizer', 'unet', 'scheduler', 'safety_checker', 'feature_extractor']

为了加深对pipeline的理解，将逐一拆分进行理解，最后再组合在一起实现相应的功能。
1、VAE（变分自动编码器）是一种模型，可以将其输入编码为压缩表示，然后将此“潜在”表示解码为接近原始输入的内容。当使用稳定扩散生成图像时，我们首先通过在VAE的“潜在空间”中应用扩散过程来生成潜变量，然后在最后对其进行解码以查看生成的图像。过程如下。

下面是一些代码，它接受输入图像，将其编码为潜在表示，然后使用 VAE 再次解码：

# Create some fake data (a random image, range (-1, 1))
images = torch.rand(1, 3, 512, 512).to(device) * 2 - 1 
print("Input images shape:", images.shape)

# Encode to latent space
with torch.no_grad():
  latents = 0.18215 * pipe.vae.encode(images).latent_dist.mean
print("Encoded latents shape:", latents.shape)

# Decode again
with torch.no_grad():
  decoded_images = pipe.vae.decode(latents / 0.18215).sample
print("Decoded images shape:", decoded_images.shape)

Input images shape: torch.Size([1, 3, 512, 512])
Encoded latents shape: torch.Size([1, 4, 64, 64])
Decoded images shape: torch.Size([1, 3, 512, 512])

可以看到图像的高和宽被压缩了8倍，通道数为4，这些潜变量图像包含信息更加丰富，比用原图像作为输入可以实现更快的扩散模型，且内存占用量较小。而0.1825是为了匹配SD训练期间所需的比例因子。
2、分词器和文本编码器(The Tokenizer and Text Encoder)
文本编码器的目标是将输入字符串（提示）转换为数字表示形式，该表示形式可以作为条件提供给 UNet。首先使用管道的分词器将文本转换为一系列词元。文本编码器的词汇表大约有 50k 个词元 - 任何不在此词汇表中的单词都会被拆分为更小的子词。然后，词元输入到文本编码器模型去生成相应的嵌入表示，这个编码器本质是 一个transformer，最初被训练为 CLIP 的文本编码器。希望这个预训练的transformer已经学习到了丰富的文本表示，这对扩散任务也很有用。
下面通过对一个例子的提示进行编码来测试此过程，首先手动词元化并通过文本编码器将其馈送，然后使用 pipelines encode_prompt 方法显示整个过程，包括填充/截断长度到 77 个标记的最大长度：

# Tokenizing and encoding an example prompt manually

# Tokenize
input_ids = pipe.tokenizer(["A painting of a flooble"])['input_ids']
print("Input ID -> decoded token")
for input_id in input_ids[0]:
  print(f"{input_id} -> {pipe.tokenizer.decode(input_id)}")

# Feed through CLIP text encoder
input_ids = torch.tensor(input_ids).to(device)
with torch.no_grad():
  text_embeddings = pipe.text_encoder(input_ids)['last_hidden_state']
print("Text embeddings shape:", text_embeddings.shape)

Input ID -> decoded token
49406 -> <|startoftext|>
320 -> a
3086 -> painting
539 -> of
320 -> a
4062 -> floo
1059 -> ble
49407 -> <|endoftext|>
Text embeddings shape: torch.Size([1, 8, 1024])

# Get the final text embeddings using the pipeline's encode_prompt function
text_embeddings = pipe._encode_prompt("A painting of a flooble", device, 1, False, '')
text_embeddings.shape

torch.Size([1, 77, 1024])

这些文本嵌入（即文本编码器模型最后一个transformer块输出的隐藏状态）将作为forward附加参数提供给Unet.
3、Unet
UNet 接受噪声输入并预测噪声。与前面的打卡笔记不同的是，这次输入不是图像，而是图像的潜变量表示。除了时间步长条件控制之外，此 UNet 还接受提示的文本嵌入作为附加输入。

这里自己生成一些虚拟数据查看效果：

# Dummy inputs
timestep = pipe.scheduler.timesteps[0]
latents = torch.randn(1, 4, 64, 64).to(device)
text_embeddings = torch.randn(1, 77, 1024).to(device)

# Model prediction
with torch.no_grad():
  unet_output = pipe.unet(latents, timestep, text_embeddings).sample
print('UNet output shape:', unet_output.shape) # Same shape as the input latents

UNet output shape: torch.Size([1, 4, 64, 64])

4、Scheduler
Scheduler用来控制噪声的添加，并根据模型预测管理噪声样本的更新。默认的是 PNDMScheduler ，也有其他类型（例如 LMSDiscreteScheduler ）。
下面绘制一副噪声时间表，来观察随时间变化的噪声水平：

plt.plot(pipe.scheduler.alphas_cumprod, label=r'
')
plt.xlabel('Timestep (high noise to low noise ->)');
plt.title('Noise schedule');plt.legend();

可以看到时间步越大，噪声所占比例在逐步减小。
也可以尝试使用不同的Scheduler,如下：

from diffusers import LMSDiscreteScheduler

# Replace the scheduler
pipe.scheduler = LMSDiscreteScheduler.from_config(pipe.scheduler.config)

# Print the config
print('Scheduler config:', pipe.scheduler)

# Generate an image with this new scheduler
pipe(prompt="Palette knife painting of an winter cityscape", height=480, width=480,
     generator=torch.Generator(device=device).manual_seed(42)).images[0]
     

Scheduler config: LMSDiscreteScheduler {
  "_class_name": "LMSDiscreteScheduler",
  "_diffusers_version": "0.11.1",
  "beta_end": 0.012,
  "beta_schedule": "scaled_linear",
  "beta_start": 0.00085,
  "clip_sample": false,
  "num_train_timesteps": 1000,
  "prediction_type": "epsilon",
  "set_alpha_to_one": false,
  "skip_prk_steps": true,
  "steps_offset": 1,
  "trained_betas": null
}

5、A DIY Sampling Loop 自定义采样循环
了解了所有组件以后，就可以自己组装来利用pipeline生成图片啦，下面是一个示例：

guidance_scale = 8 #@param
num_inference_steps = 30 #@param
prompt = "Beautiful picture of a wave breaking" #@param
negative_prompt = "zoomed in, blurry, oversaturated, warped" #@param

# Encode the prompt
text_embeddings = pipe._encode_prompt(prompt, device, 1, True, negative_prompt)

# Create our random starting point
latents = torch.randn((1, 4, 64, 64), device=device, generator=generator)
latents *= pipe.scheduler.init_noise_sigma

# Prepare the scheduler
pipe.scheduler.set_timesteps(num_inference_steps, device=device)

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

  # Expand the latents if we are doing classifier free guidance
  latent_model_input = torch.cat([latents] * 2)

  # Apply any scaling required by the scheduler
  latent_model_input = pipe.scheduler.scale_model_input(latent_model_input, t)

  # Predict the noise residual with the UNet
  with torch.no_grad():
    noise_pred = pipe.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample

  # Perform guidance
  noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
  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

# Decode the resulting latents into an image
with torch.no_grad():
  image = pipe.decode_latents(latents.detach())

# View
pipe.numpy_to_pil(image)[0]

可以看到生成效果还是非常不错的。下面将再使用几个常用的组件进行进一步探索。

2.4Additonal Pipelines

本小节将探索一些除根据prompt生成图像之外的pipeline。

2.4.1 Img2Img

到目前为止，在示例中，我们完全从头开始生成图像，从随机潜变量开始并应用完整的扩散采样循环。但我们不必从头开始。Img2Img pipeline首先将现有图像编码为一组潜伏图像，然后向潜伏图像添加一些噪声，并将其用作起点。添加的噪声量和应用的降噪步骤数决定了img2img过程的“强度”。仅添加少量噪声（低强度）将导致非常小的变化，而添加最大噪声量并运行完整的降噪过程将产生与输入几乎不相似的图像，除了整体结构上的一些相似之处。
此管道不需要特殊模型，因此只要模型 ID 与上面的文本到图像示例相同，就不需要下载新文件。

# Loading an Img2Img pipeline
model_id = "stabilityai/stable-diffusion-2-1-base"
img2img_pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id).to(device)

加载一个init_image进行实践：

# Apply Img2Img
result_image = img2img_pipe(
    prompt="An oil painting of a man on a bench",
    image=init_image, # The starting image
    strength=0.6, # 0 for no change, 1.0 for max strength
).images[0]

# View the result
fig, axs = plt.subplots(1, 2, figsize=(12, 5))
axs[0].imshow(init_image);axs[0].set_title('Input Image')
axs[1].imshow(result_image);axs[1].set_title('Result');

可以看到原始图像的一副狗狗被转化为了一张男人坐在长椅上的油画。

2.4.2 In-Painting

如果我们想保持一些输入图像不变，但在其他部分生成一些新的东西，该怎么办？这称为“修复”。虽然可以使用与之前演示相同的模型（通过 StableDiffusionInpaintPipelineLegacy ）来完成，但我们可以通过使用自定义微调版本的 Stable Diffusion 来获得更好的结果，该版本采用mask图片作为额外的条件反射。mask图像的形状应与输入图像相同，要替换的区域为白色，要保持不变的区域为黑色。

以下是我们如何加载这样的管道，并将其应用于“设置”部分中加载的示例图像和掩码：

# Load the inpainting pipeline (requires a suitable inpainting model)
pipe = StableDiffusionInpaintPipeline.from_pretrained("runwayml/stable-diffusion-inpainting")
pipe = pipe.to(device)

# Inpaint with a prompt for what we want the result to look like
prompt = "A small robot, high resolution, sitting on a park bench"
image = pipe(prompt=prompt, image=init_image, mask_image=mask_image).images[0]

# View the result
fig, axs = plt.subplots(1, 3, figsize=(16, 5))
axs[0].imshow(init_image);axs[0].set_title('Input Image')
axs[1].imshow(mask_image);axs[1].set_title('Mask')
axs[2].imshow(image);axs[2].set_title('Result');

以上就是根据mask生成的图像，可以看出mask区域的狗狗变成了机器人，这里面默认才用的模型是CLIPSEG

2.4.3 Depth2Image

Img2Img 很棒，但有时我们想创建一个具有原始但完全不同的颜色或纹理的新图像。可能很难找到一种 Img2Img 强度来保留我们想要的布局，同时又不保留输入颜色。
这时就要进行另一个微调模型Depth2Image了，其在生成时将深度的信息作为额外的条件。该管道使用深度估计模型来创建深度图，然后在生成图像时将其提供给微调的 UNet，以（希望）保留初始图像的深度和结构，同时填充全新的内容。

# Load the Depth2Img pipeline (requires a suitable model)
pipe = StableDiffusionDepth2ImgPipeline.from_pretrained("stabilityai/stable-diffusion-2-depth")
pipe = pipe.to(device)
     

# Inpaint with a prompt for what we want the result to look like
prompt = "An oil painting of a man on a bench"
image = pipe(prompt=prompt, image=init_image).images[0]

# View the result
fig, axs = plt.subplots(1, 2, figsize=(16, 5))
axs[0].imshow(init_image);axs[0].set_title('Input Image')
axs[1].imshow(image);axs[1].set_title('Result');

可以看到整体结构与原图像一致，但有更多的颜色变化，非常奇怪的是右图生成图片的男子也模仿了狗的形状和姿势哈哈哈。