AI 绘画Stable Diffusion 研究（七） 一文读懂 Stable Diffusion 工作原理_noise schedule

大家好，我是风雨无阻。

本文适合人群：

• 想要了解AI绘图基本原理的朋友。

• 对Stable Diffusion AI绘图感兴趣的朋友。

  
  

本期内容：

• Stable Diffusion 能做什么

• 什么是扩散模型

• 扩散模型实现原理

• Stable Diffusion 潜扩散模型

• Stable Diffusion文本如何影响图片生成

• Stable Diffusion Cross-attention 技术

• Stable Diffusion noise schedule 技术

• Stable Diffusion文生图底层运行演示

  
  

一、 Stable Diffusion能做什么

通过前面几篇文章关于 Stable Diffusion 整合包的安装、ControlNet插件的介绍使、sd模型的安装和使用以及文生图功能的介绍后，相信看过的朋友应该都清楚的知道 Stable Diffusion 是做什么的吧？

对于新朋友，想详细了解的话，请前往：

AI 绘画Stable Diffusion 研究（一）sd整合包v4.2 版本安装说明
AI 绘画Stable Diffusion 研究（二）sd模型ControlNet1.1 介绍与安装
AI 绘画Stable Diffusion 研究（三）sd模型种类介绍及安装使用详解
AI 绘画Stable Diffusion 研究（四）sd文生图功能详解（上）
AI 绘画Stable Diffusion 研究（五）sd文生图功能详解（下）
AI 绘画Stable Diffusion 研究（六）sd提示词插件

这里再用最直白的话说一下：SD它是一个text-to-image模型 ，通过给定的 text prompt（文本提示词），可生成一张匹配文本的图片。

二、什么是扩散模型

大家都经常听到有人说，Stable Diffusion是一种潜在扩散模型（Diffusion Models）。

那我们先弄明白什么是扩散模型？

为什么叫扩散模型呢？因为它的数学公式看起来非常像物理上的扩散现象。

1、前向扩散

假如我们训练一个模型如下：

正如上图所示，是一个前向扩散的过程，它是在训练图像上逐渐添加噪声，最后变成完全随机噪声图，并最终无法辨认噪点图对应的初始图片。

这个过程就像是一滴墨水滴在一杯清水里，会慢慢扩散最终均匀分布在清水里一样，且无法判断它最初是从水杯的中心滴入，还是从边缘滴入，扩散这个名字就是这么来的。

2、反向扩散

反向扩散的思想是：输入一张噪点图，反向扩散（Reverse Diffusion），让上述过程获得逆向从随机噪声图生成清晰图像的过程。

从反向扩散的角度来说，我们需要知道有多少“噪点”加入到了某张图片里。

那么要知道答案的方式便是：训练一个神经网络来预测添加的噪点，这在SD里称为噪点预测器（Noise Predicator），其本质是一个U-Net模型。

训练流程为：

（1）、选择一张训练图（例如一张猫的图片）

（2）、生成随机的噪点图

（3）、给这张图继续增加多轮噪点

（4）、训练Noise Predicator，预测加入了多少噪点，通过神经网络训练权重，并展示其正确答案。

反向扩散训练的重点下图中的噪声预测器（Noise Predicator），它可以通过训练得出每次需要减掉的噪声，每次需要减多少噪声是预测出来的，从而实现还原清晰图片的目的。

三、扩散模型实现原理

扩散模型（Diffusion Models）的成功，其实并非横空出世，突然出现在人们的视野中。其实早在2015年就已有人提出了类似的想法，最终在2020年提出了扩散模型的生成技术。

以下是扩散模型推导公式：

更详细的原理：

参考：扩散模型详解原理+代码

通过前面的介绍，我们大概明白了，什么是扩散模型，但这并不是 Stable Diffusion的工作原理。

这是因为：上述扩散过程是在图片空间里完成的，无论是模型训练，还是生成图片的过程，都是需要海量的算力支持和内存需求。

想象一下：一张512 x 512的图片（包含3个颜色通道：红、绿、蓝），它的空间是786432维，也就是说我们要为一张图片指定这么多的值。因此，基本无法在单个GPU上运行。

Stable Diffusion就是降低算力和内存需求的一个解决方案。它让Stable Diffusion在消费级GPU上运行成为了可能。

四、Stable Diffusion 潜扩散模型

Stable Diffusion 它是一个Latent Diffusion Model（潜扩散模型）。其方式是将图片压缩到一个“潜空间”（Latent Space）中，而不是在高维的图片空间里工作。潜空间比图片空间小了48倍，所以它可以节省大量计算，继而运行速度更快。

扩散过程会分成很多步循环，而每一步的过程如下图所示，将文本描述、隐变量、步数等数值传入UNet，生成新的隐变量，而这个过程会涉及一些模型。

在最后一步循环，将隐特征经由 Variational Autoencoder（VAE）解码成图像。

这个过程的核心思想就是：压缩图像，它通过变分自编码器 Variational Autoencoder（VAE）模型，把图像压缩到极致，我们把此类压缩方式称作降维，这种降维级别的压缩不丢失重要信息。

经过压缩后，图像被称作低维潜在（Latent）“图像”，作为U-net的输入，去潜空间（Latent Space）里一步一步降噪后，完成反向扩散的低维“图片”还得通过VAE的解码器，把图像从潜空间转换回像素空间（Pixel Space）。

VAE包含2部分：Encoder与Decoder。

• Encoder将一张图片压缩到“潜空间”里的一个低维空间表示

• Decoder从“潜空间”里的表示恢复为一张图片

下列代码演示了VAE模型的使用方法，其中load_vae为根据配置init_config去初始化模型，然后从预训练模型model.ckpt中读取参数，预训练模型的first_stage_model即指代VAE模型。

from ldm.models.autoencoder import AutoencoderKL
#VAE模型
def load_vae():
    #初始化模型
    init_config = {
        "embed_dim": 4,
        "monitor": "val/rec_loss",
        "ddconfig":{
          "double_z": True,
          "z_channels": 4,
          "resolution": 256,
          "in_channels": 3,
          "out_ch": 3,
          "ch": 128,
          "ch_mult":[1,2,4,4],
          "num_res_blocks": 2,
          "attn_resolutions": [],
          "dropout": 0.0,
        },
        "lossconfig":{
          "target": "torch.nn.Identity"
        }
    }
    vae = AutoencoderKL(**init_config)
    #加载预训练参数
    pl_sd = torch.load("model.ckpt", map_location="cpu")
    sd = pl_sd["state_dict"]
    model_dict = vae.state_dict()
    for k, v in model_dict.items():
        model_dict[k] = sd["first_stage_model."+k]
    vae.load_state_dict(model_dict, strict=False)

    vae.eval()
    return vae

#测试vae模型
def test_vae():
    vae = load_vae()
    img = load_image("girl_and_horse.png")  #(1,3,512,512)   
    latent = vae.encode(img).sample()       #(1,4,64,64)
    samples = vae.decode(latent)            #(1,3,512,512)
    save_image(samples,"vae.png")

test_vae()
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

五、Stable Diffusion 文本如何影响图片生成

在 Stable Diffusion 模型中，prompt 是通过引导向量（guidance vector）来控制 U-Net 的。具体来说，prompt 会被编码成一个文本嵌入向量（text embeddings），然后与其他输入一起传递给 U-Net。

通过这种方式，prompt 能够影响 U-Net 的输出，从而在生成过程中引导模型产生符合预期的结果，即通过 prompt 产生我们想要的图。

在Stable Diffusion模型限制prompt在75个单词。

下图是文本提示词（text prompt）如何处理并输入到Noise Predictor的流程：

根据上图，我们可以看到这个流程：

首先，Tokenizer（分词器）将每个输入的单词转为一个数字，我们称为token。

然后，每个token转为一个768维的向量，称为词嵌入（embedding）。

最后，由Text Transformer处理词嵌入，并可以被Noise predictor进行消费。

1、分词器 （Tokenizer）

人类可以读懂单词，但计算机只能读懂数字。所以这也是为什么文本提示词首先要转为单词。

文本提示词（text prompt）首先由一个CLIP tokenizer做分词。

CLIP是一个深度学习模型，由Open AI开发，用于为任何图片生成文本描述。

以下是CLIP具体的实例

展示了如何将文本“苹果”通过CLIP技术转化为能输入到神经网络中训练的tokens数据。

这里使用Python和OpenAI库来实现。

（1）、安装依赖库

pip install torch openai
1

（2)、导入相关库

import torch import openai
1

（3）、加载CLIP模型

model, preprocess = openai.clip.load("ViT-B/32")
1

（4）、准备输入文本

text_description = "苹果"
1

(5)、将文本转换为tokens

使用CLIP模型的tokenize方法将文本转换为tokens

text_tokens = openai.clip.tokenize(text_description)
1

这里，text_tokens是一个PyTorch张量，形状为(1, N)，其中N是文本描述中的token数量。

在这个例子中，N=3，因为"苹果"被分成了3个tokens。

(6)、 查看tokens

print(f"Tokens: {text_tokens}")
1

输出结果可能类似于：

Tokens: tensor([[49406, 3782, 49407]])
1

这里，49406表示开始符号（start-of-sentence），3782表示“苹果”，49407表示结束符号（end-of-sentence）。

通过以上步骤，我们将文本“苹果”转换为了tokens。

PS：

• Stable Diffusion v1使用了CLIP模型的tokenizer

• Tokenizer只能将其在训练过程中见到过的单词进行分词

  例如:假设CLIP模型里有“dream”与“beach”单词，但是没有“dreambeach”单词。

  Tokenizer会将“dreambeach”分成2个单词“dream”与“beach”。

• 1个单词并非代表1个token，而是有可能进一步进行拆分

• 空格也是token的一部分

  例如:短语 “dream beach” 产生了两个token “dream” 和 “[space]beach”。

  这些标记与 “dreambeach” 产生的标记不同，后者是 “dream” 和 “beach”（beach 前没有空格）。

  
  

2、词嵌入（Embedding）

(1)、为什么需要词嵌入（Embedding）？

因为有些单词相互之间是非常相似，我们希望利用到这些语义信息。

例如：

man、gentleman、guy的词嵌入是非常相近的，因此它们可以相互替换。

Monet、Manet以及Degas都以印象派的风格绘画，但是方式各不相同。

这些名字看起来是非常相似，但是在词嵌入(Embedding)里是不一样的。

（2）、词嵌入（Embedding） 是如何工作的？

Embedding 将输入的tokens转换为一个连续的向量来表示，这个向量可以捕捉文本中的语义信息。在我们的例子中，"苹果"的tokens经过CLIP模型的encode_text方法后，会得到一个特征向量。

这个特征向量是一个高维空间中的点，通常具有固定的维度（在CLIP模型中，维度为512）。请注意，由于模型权重和随机性的原因，每次运行时生成的特征向量可能略有不同。以下是一个示例输出：

print(f"Text features: {text_features}")
1

输出结果可能类似于：

Text features: tensor([[-0.0123,  0.0345, -0.0678, ...,  0.0219, -0.0456,  0.0789]])
1

这里，text_features是一个形状为(1, 512)的PyTorch张量，其中包含了“苹果”这个词的向量表示。神经网络可以利用这个向量表示进行训练和预测任务。

Stable diffusion v1使用Open AI的ViT-L/14模型，词嵌入为768维的向量。

3、文本转换器(text transformer)

（1）、为什么需要text transformer ？

既然通过embedding后可以直接输入到模型中进行训练，为何在stable diffusion中还需要将embedding通过text transformer转换后再作为模型的输入呢？

这是因为Stable Diffusion模型是一个图像生成模型，它需要理解输入文本的语义信息以生成与之相关的图像。直接使用基本的文本embedding可能无法充分捕捉到文本中的复杂语义关系。通过使用text transformer，可以获得一个更丰富、更具表现力的文本表示，这有助于提高生成图像的质量和与输入文本的相关性。

使用text transformer 在捕捉文本语义信息时，能够考虑到更多上下文关系和抽象概念。

这个转换器就像是一个通用的条件（conditioning）适配器。

（2）、text transformer转换示例

下面以"苹果"为例进行说明。

假设我们已经获得了"苹果"的基本embedding（一个形状为(1, 512)的PyTorch张量）：

text_features = tensor([[-0.0123,  0.0345, -0.0678, ...,  0.0219, -0.0456,  0.0789]])
1

接下来，我们将这个张量输入到text transformer中：

transformed_text_features = text_transformer(text_features)
1

经过text transformer处理后，我们可能会得到一个新的张量，如：

print(f"Transformed text features: {transformed_text_features}")
1

输出结果可能类似于：

Transformed text features: tensor([[ 0.0234, -0.0567,  0.0890, ..., -0.0321,  0.0672, -0.0813]])
1

这个新的张量（形状仍为(1, 512)）包含了更丰富的语义信息，例如上下文关系和抽象概念。

这有助于Stable Diffusion模型更好地理解输入文本，并生成与之相关的图像。

请注意：

由于模型权重和随机性的原因，每次运行时生成的特征向量可能略有不同。

此外，具体的变化过程取决于所使用的text transformer结构和参数。

六、Stable Diffusion Cross-attention技术

Cross-attention 是通过提示词产生图片的核心技术。

文本转换器的输出，会被noise predictor在U-Net中使用到多次。

U-Net以一个叫做cross-attention机制的方式来使用它，cross-attention机制允许模型在不同的特征层次上关注相关的区域，从而提高生成结果的质量，这即是prompt适配图片的地方。

下面代码是stable diffusion所使用的transformers块，实现了cross-attention：

class SpatialTransformer(nn.Module):
    """
    Transformer block for image-like data.
    First, project the input (aka embedding)
    and reshape to b, t, d.
    Then apply standard transformer action.
    Finally, reshape to image
    """
    def __init__(self, in_channels, n_heads, d_head,
                 depth=1, dropout=0., context_dim=None):
        super().__init__()
        self.in_channels = in_channels
        inner_dim = n_heads * d_head
        self.norm = Normalize(in_channels)

        self.proj_in = nn.Conv2d(in_channels,
                                 inner_dim,
                                 kernel_size=1,
                                 stride=1,
                                 padding=0)

        self.transformer_blocks = nn.ModuleList(
            [BasicTransformerBlock(inner_dim, n_heads, d_head, dropout=dropout, context_dim=context_dim)
                for d in range(depth)]
        )

        self.proj_out = zero_module(nn.Conv2d(inner_dim,
                                              in_channels,
                                              kernel_size=1,
                                              stride=1,
                                              padding=0))

    def forward(self, x, context=None):
        # note: if no context is given, cross-attention defaults to self-attention
        b, c, h, w = x.shape
        x_in = x
        x = self.norm(x)
        x = self.proj_in(x)
        x = rearrange(x, 'b c h w -> b (h w) c')
        for block in self.transformer_blocks:
            x = block(x, context=context)
        x = rearrange(x, 'b (h w) c -> b c h w', h=h, w=w)
        x = self.proj_out(x)
        return x + x_in
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

七、Stable Diffusion noise schedule 技术

1、什么是 noise schedule ？

噪声通过多次U-Net的处理，最终会输出我们想要的图片。

在这多次处理中，每一次的降噪的幅度是不同的，所以我们就要通过schedulers来控制每次降噪的幅度（幅度一般是递减的）。这个技术就叫做 noise schedule。

如图：

那么为什么要使用 noise schedule 技术呢？

在 Stable Diffusion 这种生成模型中，U-Net 是一个核心组件，用于从噪声图像中逐步恢复出原始图像。在多次迭代过程中，降噪幅度逐渐减小的原因是为了更精细地恢复图像的细节和结构。

Stable Diffusion 的过程可以看作是一个逆向扩散过程，它从一个高度噪声的图像开始，然后通过多个步骤逐渐去除噪声以重建原始图像。在这个过程中，U-Net 被用来预测每一步的降噪操作。

在前几轮迭代中，图像中的噪声较大，因此需要较大的降噪幅度来消除这些噪声。随着迭代次数的增加，图像中的噪声逐渐减小，因此降噪幅度也应相应减小。这样做的目的是避免过度平滑或损坏已经恢复的图像细节。

通过逐渐减小降噪幅度，U-Net 可以更好地控制去噪过程，使其在保留图像细节的同时有效地去除噪声。这有助于生成更清晰、更真实的图像。

这里举一个文生图的代码，用于说明noise schedule技术：

def txt2img():
    #unet
    unet = load_unet()
    #调度器
    scheduler = lms_scheduler()
    scheduler.set_timesteps(100)
    #文本编码
    prompts = ["a photograph of an astronaut riding a horse"]
    text_embeddings = prompts_embedding(prompts)
    text_embeddings = text_embeddings.cuda()     #(1, 77, 768)
    uncond_prompts = [""]
    uncond_embeddings = prompts_embedding(uncond_prompts)
    uncond_embeddings = uncond_embeddings.cuda() #(1, 77, 768)
    #初始隐变量
    latents = torch.randn( (1, 4, 64, 64))  #(1, 4, 64, 64)
    latents = latents * scheduler.sigmas[0]    #sigmas[0]=157.40723
    latents = latents.cuda()
    #循环步骤
    for i, t in enumerate(scheduler.timesteps):  #timesteps=[999.  988.90909091 978.81818182 ...100个
        latent_model_input = latents  #(1, 4, 64, 64)  
        sigma = scheduler.sigmas[i]
        latent_model_input = latent_model_input / ((sigma**2 + 1) ** 0.5)
        timestamp = torch.tensor([t]).cuda()

        with torch.no_grad():  
            noise_pred_text = unet(latent_model_input, timestamp, text_embeddings)
            noise_pred_uncond = unet(latent_model_input, timestamp, uncond_embeddings)
            guidance_scale = 7.5 
            noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

            latents = scheduler.step(noise_pred, i, latents)

    vae = load_vae()
    latents = 1 / 0.18215 * latents
    image = vae.decode(latents.cpu())  #(1, 3, 512, 512)
    save_image(image,"txt2img.png")

txt2img()
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

八、Stable Diffusion 文生图底层运行演示

在文本生成图的场景下，我们给SD模型输入一组文本提示词，它可以返回一张图片。

第一步、 Stable Diffusion在潜空间里生成一个随机张量。

我们通过设置随机种子seed来控制这个张量的生成。如果我们设置这个随机种子为一个特定的值，则会得到相同的随机张量。这就是我们在潜空间里的图片。但是当前还全是噪点。

第二步、 Noise predictor U-Net将潜噪点图已经文本提示词作为输入，并预测噪点

此噪点同样也在潜空间内（一个4 x 64 x 64的张量）

第三步、从潜图片中抽取潜噪点，并生成了新的潜图片

第二步 与 第三步重复特定采样次数，例如20次。

第四步、VAE 的decoder将潜图片转回像素空间

这便是我们通过SD模型最终得到的图片。

参考资料：

1. How does Stable Diffusion work?

2. stable-diffusion

3.扩散模型详解原理+代码