Stable Diffusion 原理详解（附代码实现）

从结构上来看，CLIP 模型由两个 Encoder 构成，分别是用来编码文本的 TextEncoder 和用来编码图片的 ImageEncoder。CLIP 的训练数据是一堆'图片 - 文本'对形式，其工作模式如下：

1. 训练 TextEncoder 和 ImageEncoder，最大化 ItTt（图片向量与响应的文本向量相似度）
2. 利用分类标签生成句子，'a photo of a {object}'
3. 输入图片获得 It，找到最相似的句子向量 Tk，改句子对应的标签就是图片标签 在完成训练后就可以得到比较出色的文本编码器，而后两步则是为图像分类做准备。

2.3 VAE 模型

VAE 模型在 Diffusion Model 里面并非必要的，VAE 在 Stable Diffusion 中作为一种有效利用资源的方法，减少了图片生成的资源需求。下图是 VAE 的结构，其中 c 是一个可选的条件。

VAE 由 Encoder 和 Decoder 两个部分组成，首先需要输入 x，经过 Encoder 编码后，得到（μ，σ），分别表示均值和方差，这两个变量可以确定一个分布，然后在当前分布中采样出样本 z。z 通常是一个比 x 维度更低的向量。

采样出来的 z 输入 Decoder，我们希望 Decoder 的输出与输入的 x 越接近越好。这样我们就达到了图像压缩的效果。

在训练 Stable Diffusion 时，我们会把图片输入 VAE 的 Encoder，然后再拿来训练 UNet，这样我们就可以在更低的维度空间训练图像生成模型，这样就可以减少资源的消耗。

2.4 UNet 模型

UNet 模型结构与 VAE 非常相似，也是 Encoder-Decoder 架构。在 Stable Diffusion 中，UNet 作为一个噪声预测网络，在绘画过程中需要多次推理。我们先不考虑 VAE 的介入，来看看 UNet 在 Stable Diffusion 中的作用。

实际上 UNet 在 Stable Diffusion 中充当噪声预测的功能。UNet 接收一张带有噪声的图片，输出图片中的噪声，根据带噪声的图片和噪声我们可以得到加噪前的图片。这个降噪的过程通常会重复数十遍。

知道 UNet 的作用后，我们就需要创建数据集了。我们只需要图片即可，拿到图片对该图片进行 n 次加噪，直到原图变成完全噪声。而加噪前的图片可以作为输入，加噪后的数据可以作为输出。如图所示：

在加噪的过程中，噪声逐步增大。因此在降噪的过程中，我们需要有噪声图片，以及当前加噪的 step。下图是噪声预测部分的结构：

最后图像生成的步骤就是不停降噪的步骤：

最后，我们再加入 VAE。我们加噪和预测噪声的步骤不再是作用在原图上，而是作用在 VAE Encoder 的输出上面，这样我们就可以在较小的图像上完成 UNet 的训练，极大减少资源的消耗。

现在只需要在 UNet 的输入中再加入文本变量就是完整的 Stable Diffusion 了。

三、Diffusers 模块

现在我们已经知道 Stable Diffusion 的原理，为了加深理解，下面使用 Diffusers 模块实现 Stable Diffusion 的全过程。下面的代码需要使用到 pytorch、transformers 和 diffusers 模块。

3.1 使用 pipeline

HuggingFace 中的模块提供了许多 pipeline 用于各种任务，而 Stable Diffusion 则是 Text-to-image 类型的任务，我们可以使用下面几句代码完成文生图：

from diffusers import AutoPipelineForText2Image
import torch

pipeline = AutoPipelineForText2Image.from_pretrained(
	"runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16, variant="fp16"
).to("cuda")
image = pipeline(
	"stained glass of darth vader, backlight, centered composition, masterpiece, photorealistic, 8k"
).images[0]
image

生成图像如下：

上面是一种简单的调用方式，下面我们加载各个部件，手动完成图像生成的过程。

3.2 加载各个部件

除了 pipeline 直接加载，我们还可以分部件加载，分别加载 CLIP、UNet 和 VAE，代码如下：

from tqdm.auto import tqdm
from PIL import Image  
import torch  
from transformers import CLIPTextModel, CLIPTokenizer  
from diffusers import AutoencoderKL, UNet2DConditionModel, DDPMScheduler  

# 加载模型  
model_path = "runwayml/stable-diffusion-v1-5"  
vae = AutoencoderKL.from_pretrained(model_path, subfolder="vae")  
tokenizer = CLIPTokenizer.from_pretrained(model_path, subfolder="tokenizer")  
text_encoder = CLIPTextModel.from_pretrained(  
	model_path, subfolder="text_encoder"  
)  
unet = UNet2DConditionModel.from_pretrained(  
	model_path, subfolder="unet"  
)  
scheduler = DDPMScheduler.from_pretrained(model_path, subfolder="scheduler")
# 使用 gpu 加速  
torch_device = "cuda"  
vae.to(torch_device)  
text_encoder.to(torch_device)  
unet.to(torch_device)

在这里我们还加载了 Scheduler，后续会使用 Scheduler 管理降噪的步骤。

3.3 对文本进行编码

下面我们使用 CLIP 模型对文本进行编码，这里要使用到 tokenizer 和 text_encoder：

# 对文本进行编码  
prompt = ["a photograph of an astronaut riding a horse"]  
height = 512 # default height of Stable Diffusion  
width = 512 # default width of Stable Diffusion  
num_inference_steps = 25 # Number of denoising steps  
guidance_scale = 7.5 # Scale for classifier-free guidance  
batch_size = len(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(torch_device))[0]

其中 text_embeddings 就是文本编码结果。

3.4 获取潜变量

在训练过程中潜变量 Latent 是由 VAE 的 Encoder 得到的，而在生成过程中，Latent 则是符合一定分别的随机噪声。代码如下：

# 获取 latent  
latents = torch.randn(  
	(batch_size, unet.config.in_channels, height // 8, width // 8),  
	device=torch_device,  
)  
latents = latents * scheduler.init_noise_sigma

torch.randn 可以得到方差为 1 的噪声，而 latents * scheduler.init_noise_sigma 则把方差修改为 scheduler.init_noise_sigma。

3.5 降噪

接下来就是重复多次 UNet 推理，得到降噪后的 Latent：

# 降噪  
scheduler.set_timesteps(num_inference_steps)  
for t in tqdm(scheduler.timesteps):  
	latent_model_input = latents  
	latent_model_input = scheduler.scale_model_input(latent_model_input, timestep=t)  
	with torch.no_grad():  
		# 预测噪声
		noise_pred = unet(
			latent_model_input, 
			t, 
			encoder_hidden_states=text_embeddings
		).sample 
	# 降噪 
	latents = scheduler.step(noise_pred, t, latents).prev_sample

最后得到的 latents 变量就是降噪后的结果，在训练过程中对应 VAE Encoder 的输出，因此我们还需要使用 VAE Decoder 还原出图片。

3.6 VAE 解码

下面就是使用 VAE Decoder 解码出原图：

# 使用 vae 解码  
latents = 1 / 0.18215 * latents  
with torch.no_grad():  
	image = vae.decode(latents).sample  
	image = (image / 2 + 0.5).clamp(0, 1).squeeze()  
	image = (image.permute(1, 2, 0) * 255).to(torch.uint8).cpu().numpy()  
	images = (image * 255).round().astype("uint8")  
	image = Image.fromarray(image)  
	image.show()

最后生成如下图片：

四、总结

今天我们以 GAN 开始，介绍了 AI 绘画领域的一些模型，并把 Stable Diffusion 作为今天的主角，详解介绍了 Stable Diffusion 的实现原理。

我们还使用 Diffusers 模块实现了 Stable Diffusion 生成图像的代码。在 Stable Diffusion 中，还有诸如 LoRA、Controlnet 等相关技术，在本文没有详细提到。而这些东西在 AI 绘画中却非常重要，也让 AI 绘画可以应用在更多领域。

我们可以使用 Stable Diffusion Webui 等工具使用 LoRA 和 Controlnet 等工具，我们还可以在 Diffusers 中使用这些根据。后续我们将介绍 Diffusers 模块如何加载 LoRA 等附加网络。