Stable Diffusion 原理介绍与源码分析（一）(铡刀演示)

推荐整理分享Stable Diffusion 原理介绍与源码分析（一）(铡刀演示)，希望有所帮助，仅作参考，欢迎阅读内容。

文章相关热门搜索词:铡读什么字,铡刀演示,铡刀的视频,铡头视频,铡刀百科,铡头视频,铡头视频,铡这个字念什么,内容如对您有帮助，希望把文章链接给更多的朋友！

Stable Diffusion 是 Stability AI 公司开源的 AI 文生图扩散模型。之前在文章 扩散模型 (Diffusion Model) 简要介绍与源码分析 中介绍了扩散模型的原理与部分算法代码，满足基本的好奇心后便将其束之高阁，没成想近期 AIGC 的发展速度之快大大出乎我的意料，尤其是亲手跑出下面这张 AI 生成的图像， Stable Diffusion 终又重新回到我的视野：

作为一名算法工程师，需要有一双能看透事物本质的眼睛，这张图片最先吸引我的不是内容，而是其生成质量：图像高清、细节丰富，非之前看到的一些粗陋 Toy 可比，红框中标注出来的不协调之处，也是瑕不掩瑜。因此，进一步分析 Stable Diffusion 整个工程框架的原理，实在是迫在眉睫，期待日后能修复红框中的不协调之处，为 AIGC 的进一步发展做出一个技术人员应有的贡献。

Stable Diffusion 整个框架的源码有上万行，没有必要全部分析。本文以 “文本生成图像（text to image）” 为主线，考察 Stable Diffusion 的运行流程以及各个重要的组成模块，在介绍时采用 “总-分” 的形式，先概括整体框架，再分析各个组件（如 DDPM、DDIM 等），另外针对代码中的部分非主流逻辑，比如 predict_cids、return_ids 这些小细节谈谈我的看法。文章内容较长，准备拆分成多个部分。

源码地址：Stable Diffusion

之前我写过很多代码分析文章，但在我遇到问题重新去翻阅时，发现要快速定位到目标位置并准确理解代码意图，仍然存在很大困难，密密麻麻的整块代码，每一次阅读都仿若初见，不易理解，原因在于摘录时引入过多的实现细节，降低了信息的传播效率。

经过一番思考，我不再图省事，决定采用伪代码的方式记录核心原理。平时我深度分析代码时会采用这种方式，对代码进行额外的抽象，相对会耗些时间，但私以为这是有益处的。举个例子，比如 DDPM 模型前向 Diffusion 的代码，如果我用伪代码的方式去写，将是如下的效果：

可以看到，刨除掉无关的实现细节之后，DDPM 的实现是如此的简洁，倘若再配合一定的注释，可方便快速理解，让人获得一种整体而全面的掌控感。此外还应该在文中多增加框图、模型图等来对代码的实现细节进行更直观的展示。

可以在微信中搜索 “珍妮的算法之路” 或者 “world4458” 关注我的微信公众号, 可以及时获取最新原创技术文章更新.

另外可以看看知乎专栏 PoorMemory-机器学习, 以后文章也会发在知乎专栏中.

首先看下 Stable Diffusion 文本生成图像整体框架（文章绘图吐血…希望有一天 AI 能进行辅助）：

上图框架内的模块较多，从上到下分为 3 块，我在图中使用 Part 1、2、3 进行了标注。框架包含训练 + 采样两个阶段，其中：

训练阶段 （查看图中 Part 1 和 Part 2），主要包含：

采样阶段（查看图中 Part 2 和 Part 3），也就是我们加载模型参数后，输入提示词就能产出图像的阶段。主要包含：

经过上面的介绍，对 Stable Diffusion 整体会有个较清晰的认识，下面就可以按图索骥，将各个重点模块尽力去弄明白。限于个人精力与有限的空闲时间，目前除了 FrozenCLIPEmbedder 和 DPM 算法 (图中没写），Stable Diffusion 的其他模块都大致看了看，包括：

后面会简单介绍一下，记录学习过程。

在阅读代码的过程中，发现有些重量级的论文必须得阅读一下。扩散模型的理论推导还是有些复杂的，有时候公式推导和代码实现相互结合看，可以加深对知识的理解。这里列一下对我阅读代码有很大帮助的论文：

下面对 Stable Diffusion 中的重要组成模块进行简要分析。主要包含：

首先介绍一下 UNetModel 结构，方便后续的文章直接进行引用。

画了一下 Stable Diffusion 中使用的 UNetModel，就不分析代码了，看图很容易将代码写出来。Stable Diffusion 采用 UNetModel 这种 Encoder-Decoder 结构来实现扩散的过程，对噪声进行预估, 网络结构如下：

模型的输入包含三个部分：

模型使用 DownSample 和 UpSample 来对样本进行下采样和上采样，此外出现最多的模块是 ResBlock 以及 SpatialTransformer，其中图中每一个 ResBlock 接收来自上一个模块的输入以及 timesteps 对应的 embedding timestep_emb （大小为 [B, 4*M]，M 是可配置的参数）；而图中每一个 SpatialTransformer接收来自上一个模块的输入以及 context (Prompt 文本的 embedding 表示），使用 Cross Attention，以 context 为 condition，学习 Prompt 和图像的匹配关系。但图上只在虚线框中显示了两个模块有多个输入，其他模块没有画出来）

可以看到，最后模型的输出大小为 [B, C, H, W], 和输入大小相同，也就是说 UNetModel 不改变输入输出的大小。

下面再分别看看 ResBlock、timestep_embedding、context 以及 SpatialTransformer 的实现。

ResBlock 网络结构图如下，它接受两个输入，图像 x 以及 timestep 对应的 embedding：

timestep_embedding 的生成方式如下，用的是 Tranformer（Attention is All you Need）这篇 paper 中的方法：

即 context 的实现。Prompt 使用 CLIP 模型进行编码，我没有对 CLIP 模型详细学习，暂时也没有深入看的打算，后续有机会再补充；代码中使用预训练好的 CLIP 生成 context:

最后再看下 SpatialTransformer 的实现，其模块比较多，在接收图像作为输入时，还使用 context 文本作为 condition 信息，二者使用 Cross Attention 进行建模。进一步展开 SpatialTransformer， 发现包含 BasicTransformerBlock ，它实际调用 Cross Attention 模块，而在 Cross Attention 模块中，图像信息作为 Query，文本信息作为 Key & Value，模型会关注图像和文本各部分内容的相关性：

我觉得可以用一种朴素的想法来理解这里 Cross Attention 的作用，比如训练时给定一张马吃草的图，以及文本提示词：“一匹白色的马在沙漠吃草”，在做 Attention 时，文本中的 “马” 这个关键词和图像中的动物（也是 “马”）的关联性更强，因为权重也更大，而 “一匹”、 “白色”、“沙漠”、 “草” 等权重更低；此时，当模型被训练的很好后，模型不仅将可以学习到图像和文本之间的匹配关系，通过 Attention 还可以学习到文本中的各个关键词想突出图像中哪些主体。

而当我们输入提示词用模型来生成图像时，比如输入 “一匹马在吃草”，由于模型此时已经能捕捉图像和文本的相关性以及文本中的重点信息，当它看到文本 “马”，在黑盒魔法的运作下，会重点突出图像 “马” 的生成；当它看到 “草” 时，便重点突出图像 “草” 的生成，从而尽可能生成和文本进行匹配的图像。

至此，UNetModel 各个重要组件基本介绍完毕。

由于 UNetModel 模型结构并不复杂，看图基本就能写出代码，一图胜千言啊。另外我标注了每个模块输出结果的大小，很方便在大脑中运行模型，哈哈哈。

本文大致介绍了一下 Stable Diffusion 文生图代码的整体框架，列出了扩散模型部分核心论文，简要分析了 UNetModel。后续再分析其他核心组件。

发现 AIGC 的发展实在太快了，学不过来啊… 愈发觉得庄子所言甚是：以有涯随无涯，殆矣！