天才程序员周弈帆 | Stable Diffusion 解读（四）：Diffusers实现源码解读

本文来源公众号“天才程序员周弈帆”，仅用于学术分享，侵权删，干货满满。

原文链接：Stable Diffusion 解读（四）：Diffusers实现源码解读

接上一篇文章[天才程序员周弈帆 | Stable Diffusion 解读（三）：原版实现源码解读(篇幅略长，建议收藏！)-优快云博客]，我们来学习Stable Diffusion在Diffusers中的实现。

本文用到的Stable Diffusion版本是v1.5。Diffusers版本是0.25.0。为了提升可读性，本文对源代码做了一定的精简，部分不会运行到的分支会被略过。

1 Diffusers

Diffusers是由Hugging Face维护的一套Diffusion框架。这个库的代码被封装进了一个Python模块里，我们可以在安装了Diffusers的Python环境中用import diffusers随时调用该库。相比之下，Diffusers的代码架构更加清楚，且各类Stable Diffusion的新技术都会及时集成进Diffusers库中。

由于我们已经在上篇文章中学过了Stable Diffusion官方源码，在学习Diffusers代码时，我们只会大致过一过每一段代码是在做什么，而不会赘述Stable Diffusion的原理。

1.1 安装

安装该库时，不需要克隆仓库，只需要直接用pip即可。

pip install --upgrade diffusers[torch]

之后，随便在某个地方创建一个Python脚本文件，输入官方的示例项目代码。

from diffusers import DiffusionPipeline
import torch

pipeline = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
pipeline.to("cuda")
pipeline("An image of a squirrel in Picasso style").images[0].save('output.jpg')

运行代码后，"一幅毕加索风格的松鼠图片"的绘制结果会保存在output.jpg中。我得到的结果如下：

在Diffusers中，from_pretrained函数可以直接从Hugging Face的模型仓库中下载预训练模型。比如，示例代码中from_pretrained("runwayml/stable-diffusion-v1-5", ...)指的就是从模型仓库https://huggingface.co/runwayml/stable-diffusion-v1-5中获取模型。

如果在当前网络下无法从命令行中访问Hugging Face，可以先想办法在网页上访问上面的模型仓库，手动下载v1-5-pruned.ckpt。之后，克隆Diffusers的GitHub仓库，再用Diffusers的工具把Stable Diffusion原版模型文件转换成Diffusers支持的模型格式。

git clone git@github.com:huggingface/diffusers.git
cd diffusers
python scripts/convert_original_stable_diffusion_to_diffusers.py --checkpoint_path <src> --dump_path <dst>

比如，假设你的模型文件存在ckpt/v1-5-pruned.ckpt，你想把输出的Diffusers的模型文件存在ckpt/sd15，则应该输入：

python scripts/convert_original_stable_diffusion_to_diffusers.py --checkpoint_path ckpt/v1-5-pruned.ckpt --dump_path ckpt/sd15 

之后修改示例脚本中的路径，就可以成功运行了。

from diffusers import DiffusionPipeline
import torch

pipeline = DiffusionPipeline.from_pretrained("ckpt/sd15", torch_dtype=torch.float16)
pipeline.to("cuda")
pipeline("An image of a squirrel in Picasso style").images[0].save('output.jpg')

对于其他的原版SD checkpoint（比如在civitai上下载的），也可以用同样的方式把它们转换成Diffusers兼容的版本。

1.2 采样

Diffusers使用Pipeline来管理一类图像生成算法。和图像生成相关的模块（如U-Net，DDIM采样器）都是Pipeline的成员变量。打开Diffusers版Stable Diffusion模型的配置文件model_index.json（在 https://huggingface.co/runwayml/stable-diffusion-v1-5/blob/main/model_index.json 网页上直接访问或者在本地的模型文件夹中找到），我们能看到该模型使用的Pipeline:

{
  "_class_name": "StableDiffusionPipeline",
  ...
}

在diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py中，我们能找到StableDiffusionPipeline类的定义。所有Pipeline类的代码都非常长，一般我们可以忽略其他部分，只看运行方法__call__里的内容。

def __call__(
    self,
    prompt: Union[str,