2024不可不会的StableDiffusion之Unet（四）

1. 引言

这是我关于StableDiffusion学习系列的第四篇文章，如果之前的文章你还没有阅读，强烈推荐大家翻看前篇内容。在本文中，我们将学习构成StableDiffusion的第三个基础组件基于Unet的扩散模型，并针该组件的功能进行详细的阐述。

闲话少说，我们直接开始吧！

2. 概览

通常来说一个U-Net包含两个输入：
‒ Noisy latent/Noise : 该Noisy latent主要是由VAE编码器产生并在其基础上添加了噪声;或者如果我们想仅根据文本描述来创建随机的新图像，则可以采用纯噪声作为输入。
‒ Text embeddings: 基于CLIP的将文本输入提示转化为文本语义嵌入（embedding）

U-Net模型的输出是从包含输入噪声的Noisy Latents中预测其所包含的噪声。换句话说，它预测输出的为Noisy Latents减去de-noised latents后的结果。

3. 导入所需的库

让我们接着通过代码来了解 U-Net。我们将首先导入所需的库并加载我们的U-Net模型。

from diffusers import UNet2DConditionModel,LMSDiscreteScheduler

## Initializing a scheduler
scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)
## Initializing the U-Net model
sd_path = r'/media/stable_diffusion/stable-diffusion-v1-4'
unet = UNet2DConditionModel.from_pretrained(sd_path,subfolder="unet",
                                            local_files_only=True,
                                            torch_dtype=torch.float16).to("cuda")

4. 可视化Scheduler

正如大家从上面的代码中观察到的那样，我们不仅导入了 unet，还导入了一个scheduler库。scheduler的目的是确定在扩散过程中的给定的步骤中向latent 添加多少噪声。我们使用以下代码来可视化 scheduler 函数：

## Setting number of sampling steps
scheduler.set_timesteps(51)
plt.plot(scheduler.sigmas)
plt.xlabel("Sampling step")
plt.ylabel("sigma")
plt.title("Schedular routine")
plt.show()

得到结果如下：

可以看到，随着sample step的增大，我们添加噪声的权重在逐渐减小。

5. 可视化扩散过程

扩散过程遵循上述采样Scheduler，我们从高噪声开始，然后逐渐按照scheduler对latent_img添加对应权重的噪声。让我们可视化一下这个过程：

img_path = r'/home/VAEs/Scarlet-Macaw-2.jpg'
img = Image.open(img_path).convert("RGB").resize((512, 512))
latent_img = pil_to_latents(img, vae)
# Random noise
noise = torch.randn_like(latent_img) 
fig, axs = plt.subplots(2, 3, figsize=(16, 12))
for c, sampling_step in enumerate(range(0,51,10)):
    encoded_and_noised = scheduler.add_noise(latent_img, noise, 
                   timesteps=torch.tensor([scheduler.timesteps[sampling_step]]))
    axs[c//3][c%3].imshow(latents_to_pil(encoded_and_noised，vae)[0])
    axs[c//3][c%3].set_title(f"Step - {sampling_step}")
plt.show()

得到结果如下：

6. 生成Unet输入

接着我们需要对latent_img, 随机添加一些噪声，用以作为我们Unet输入的Noisy Latent , 代码如下：

encoded_and_noised = scheduler.add_noise(latent_img, noise, 
                     timesteps=torch.tensor([scheduler.timesteps[40]])) 
latents_to_pil(encoded_and_noised，vae)[0]

得到结果如下：

7. 调用Unet去噪

接着我们就可以使用Unet对其进行去噪了，相应的调用代码如下：

prompt = [""]
text_input = tokenizer(prompt, padding="max_length",
                       max_length=tokenizer.model_max_length,
                       truncation=True,return_tensors="pt")
with torch.no_grad():
    text_embeddings = text_encoder(text_input.input_ids.to("cuda"))[0]

latent_model_input = torch.cat( [encoded_and_noised.to("cuda").float()]).half()
with torch.no_grad():
    noise_pred = unet(latent_model_input,40,encoder_hidden_states=text_embeddings
        )["sample"]
out_images = latent_to_pil((encoded_and_noised - noise_pred),vae)
plt.imshow(out_images[0])
plt.show()

得到去噪后结果如下：

我们也可以使用以下代码来将Unet预测出来的噪声进行可视化，代码如下：

out_noises = latent_to_pil( noise_pred,vae)
plt.imshow(out_noises[0])
plt.show()

得到结果如下：

8. Unet 在SD中的用途

潜在扩散模型使用U-Net分几个步骤通过逐渐减去潜伏空间中的噪声，来达到所需要的输出。在每一步中，添加到latents中的噪声量都会减少，直到我们达到最终的去噪输出。

我们知道，在深度学习领域U-Nets最先被引入用于生物医学图像分割任务。U-Net 有一个编码器和解码器，它们由 ResNet blocks组成。在stable diffusion中的U-Net还具有交叉注意力层，使它们能够根据提供的文本描述来控制调整输出。交叉注意力层通常添加到U-Net的编码器和解码器之间。可以在此处了解有关此 U-Net 架构的更多信息。

9. 总结

本文重点介绍了SD模型中的Unet组件的相关功能和具体工作原理，并详细介绍了其去噪过程；至此，我们完成了稳定扩散模型的三个关键组件，即 CLIP 文本编码器、VAE 和 U-Net。在下一篇文章中，我们将研究使用这些组件的扩散过程。

您学废了嘛！