博主先前整理并简单介绍了AI绘图工具的部署资源与攻略，觉得其中Stable Diffusion部分不够带劲，故开始试图从论文与代码中一探究竟。前文链接如下：

要点初见：AI绘图工具的部署资源、攻略整理（上篇）_BingLiHanShuang的博客-CSDN博客_ai绘图

要点初见：AI绘图工具的部署资源、攻略整理（下篇）_BingLiHanShuang的博客-CSDN博客

最近Stable Diffusion实在是太火爆了，在B站上看up主分析论文的视频，分P分析到引言部分就戛然而止，后续视频全是各种整合包的分享与实操，单个视频的播放量也直线上升hhh

本文将分为2个部分：

上半部分深入分析Stable Diffusion所对应的论文High-Resolution Image Synthesis with Latent Diffusion Models，即《具有潜在扩散模型的高分辨率图像合成》，论文链接如下：https://arxiv.org/pdf/2112.10752.pdf

下半部分深入Stable Diffusion项目代码，代码主要由Python撰写，分析文本转图像部分的代码（模型核心部分将在下篇进行分析）。

一、Stable Diffusion论文分析

《具有潜在扩散模型的高分辨率图像合成》在概述部分将Stable Diffusion分为2个阶段：

第一个阶段为感知压缩阶段（perceptual compression），训练了一个预训练的自编码器（pretrained autoencoders）用于下采样、上采样，自编码器学习到的是一个潜在的空间（latent space），比像素空间小很多，扩散模型在该潜在空间中训练；

第二阶段是扩散模型，语义压缩阶段，在下采样、上采样之间，引入了一个针对文本、边界框、图像的交叉注意力层（cross-attention layers）。

这样整个模型只需要训练降采样和插值之间的部分即可，大大降低了对算力的要求，降低了训练时间，这个模型也被叫做Latent Diffusion Models。

文章对核心模型的介绍主要分为感知压缩模型、潜在扩散模型、调节机制三部分。

1、感知压缩模型

感知压缩模型由一个通过感知损失和基于补丁的对抗性目标组合训练的自动编码器组成，这确保通过强制执行局部真实性将重建限制在图像流形中，并避免仅依靠像素空间损失（例如L2或L1目标）引入的模糊。

为了避免任意高方差的潜在空间，论文尝试了KL-reg、VQ-reg两种正则化方法：其中KL-reg类似于VAE，对学习潜伏的标准法线施加轻微的KL惩罚；而VQ-reg在解码器中使用矢量量化层，该模型可以解释为量化层被解码器吸收的VQGAN。

2、潜在扩散模型

扩散模型是概率模型，旨在通过逐渐去噪正态分布变量来学习数据分布。

文中的模型可以根据不同图像提供对应的归纳偏差，其中包括主要从2D卷积层构建底层UNet的能力，并使重新加权的边界进一步将目标集中在感知最相关位的目标上。文中模型的核心神经网络被实现为时序UNet，由于前向的过程是固定的，因此可以在训练期间从E中高效地获得zt，并且可以通过D单次将来自数据分布的样本解码到图像空间，如下图。

3、调节机制

调节机制方面，通过使用交叉注意机制增强其底层UNet主干，将DM转变为更灵活的条件图像生成器，从而能对于学习各种输入模式的基于注意的模型有效。

文章最后介绍了模型的局限性，一是顺序采样过程仍然比GAN慢，二是在高精度场景的效果可能不如意（虽然目前4倍超分辨率效果很好）。作者接下来将基于GAN的方法，研究结合对抗性训练和基于可能性的目标的两阶段方法在多大程度上歪曲了数据。（如何在确保真实性的基础上确保多样性）

二、Stable Diffusion代码分析

本文暂先分析文本转图像部分的执行代码scripts\txt2img.py：

1、库文件导入

import argparse, os, sys, globimport cv2import torchimport numpy as npfrom omegaconf import OmegaConffrom PIL import Imagefrom tqdm import tqdm, trangefrom imwatermark import WatermarkEncoderfrom itertools import islicefrom einops import rearrangefrom torchvision.utils import make_gridimport timefrom pytorch_lightning import seed_everythingfrom torch import autocastfrom contextlib import contextmanager, nullcontextfrom ldm.util import instantiate_from_configfrom ldm.models.diffusion.ddim import DDIMSamplerfrom ldm.models.diffusion.plms import PLMSSamplerfrom diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyCheckerfrom transformers import AutoFeatureExtractor

其中起主导作用的是：

cv2（Python版OpenCV库，一款计算机视觉库）；

torch（PyTorch，一款Python机器学习库）；

imwatermark（隐形水印添加库，stable diffusion源码中都被添加了“StableDiffusionV1”的隐形水印，不过我看了webui版没有这个）；

ldm（Python中的扩散模型库，stable diffusion图像生成的核心）；

diffusers（扩散Diffusion模型包）；

transformers（PyTorch的最新自然语言处理库）；

2、NOT SAFE FOR WORK” />

4、内部参数设置

    if opt.laion400m:        print("Falling back to LAION 400M model...")        opt.config = "configs/latent-diffusion/txt2img-1p4B-eval.yaml"        opt.ckpt = "models/ldm/text2img-large/model.ckpt"        opt.outdir = "outputs/txt2img-samples-laion400m"    seed_everything(opt.seed)    config = OmegaConf.load(f"{opt.config}")    model = load_model_from_config(config, f"{opt.ckpt}")    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")    model = model.to(device)    if opt.plms:        sampler = PLMSSampler(model)    else:        sampler = DDIMSampler(model)    os.makedirs(opt.outdir, exist_ok=True)    outpath = opt.outdir

对模型内部参数config、ckpt、outdir、device、model、sampler、outpath等参数的具体设定。

5、隐形水印设置

在这里特别分析一下txt2img.py中的隐形水印这一块：

    print("Creating invisible watermark encoder (see https://github.com/ShieldMnt/invisible-watermark)...")    wm = "StableDiffusionV1"    wm_encoder = WatermarkEncoder()    wm_encoder.set_watermark('bytes', wm.encode('utf-8'))

wm中的内容即是通过imwatermark添加的隐形水印文本，可以改成自己想添加的隐形水印文本，从而在生成的图片中附带着隐形的水印。

txt2img.py中的隐形水印添加代码：

img = put_watermark(img, wm_encoder)

def put_watermark(img, wm_encoder=None):    if wm_encoder is not None:        img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)        img = wm_encoder.encode(img, 'dwtDct')        img = Image.fromarray(img[:, :, ::-1])    return img

需要关注的是相比于imwatermark对普通png图像的隐形水印添加代码，此处将RGB转为了BGR后才叠加了水印，再通过Image.fromarray(img[:, :, ::-1])转回RGB通道。

若想看一张图是否有imwatermark添加的隐形水印，可以使用python执行下述脚本可以获取水印输出（但需要注意是否要把BGR转成RGB）：

#!/usr/bin/env python3import cv2from imwatermark import WatermarkDecoderbgr = cv2.imread('cat_wm.png')decoder = WatermarkDecoder('bytes', 32)watermark = decoder.decode(bgr, 'dwtDct')print(watermark.decode('utf-8'))

6、读取输入的文字描述

    batch_size = opt.n_samples    n_rows = opt.n_rows if opt.n_rows > 0 else batch_size    if not opt.from_file:        prompt = opt.prompt        assert prompt is not None        data = [batch_size * [prompt]]    else:        print(f"reading prompts from {opt.from_file}")        with open(opt.from_file, "r") as f:            data = f.read().splitlines()            data = list(chunk(data, batch_size))    sample_path = os.path.join(outpath, "samples")    os.makedirs(sample_path, exist_ok=True)    base_count = len(os.listdir(sample_path))    grid_count = len(os.listdir(outpath)) - 1    start_code = None    if opt.fixed_code:        start_code = torch.randn([opt.n_samples, opt.C, opt.H // opt.f, opt.W // opt.f], device=device)    precision_scope = autocast if opt.precision=="autocast" else nullcontext

此段代码将输入的所有batch的prompt（即输入的正面负面tag）打包进data中，末尾通过PyTorch的autocast设置半精度格式（即float16，字节数是通常使用的float（float32）的一半，只有2个char的长度），以大幅节省模型数据的IO时间、训练推理时间。autocast的定义如下：

class torch.autocast(device_type, enabled=True, **kwargs)

txt2img.py中通过

with precision_scope("cuda"):

调用，即意味着接下来的代码推理采用GPU资源，并使用Float16格式。半精度Float16对GPU运算的性能优化意义非凡，只要变量的minmax值在float16的定义范围内（具体要根据fp16所属库查询），且能接受相对于float32的精度损失，在GPU中使用fp16的性能相对于fp32能提升30%~50%。

模型核心部分将在下篇进行分析，欢迎交流与讨论！