文生图模型StableDiffusion技术原理与训练推理方法

StableDiffusion是一个文生图、图生图的生成式模型。

基础论文地址:https://arxiv.org/pdf/2112.10752

基于潜在扩散模型的⾼分辨率图像合成

摘要

通过将图像生成过程分解为一系列去噪自动编码器的应用,扩散模型(DMs)在图像数据及其他领域实现了最先进的合成结果。此外,它们的公式允许引入指导机制来控制图像生成过程,而无需重新训练。然而,由于这些模型通常直接在像素空间中操作,优化强大的DMs常常需要消耗数百个GPU天,并且由于顺序评估,推理成本高昂。为了在有限的计算资源上训练DMs,同时保留其质量和灵活性,将其应用于强大的预训练自动编码器的潜在空间。与之前的工作相比,在这种表示上训练扩散模型首次能够在复杂性降低和细节保留之间达到近乎最优的平衡点,极大地提升了视觉保真度。通过在模型架构中引入交叉注意力层,将扩散模型转变为强大的通用条件输入生成器,如文本或边界框,并且以卷积方式实现高分辨率合成成为可能。潜在扩散模型(LDMs)在图像修复和类别条件图像合成方面达到了新的最先进分数,并在包括无条件图像生成、文本到图像合成和超分辨率在内的各种任务中表现出色,同时显著降低了与基于像素的DMs相比的计算需求。

1、介绍

图像合成,特别是复杂自然场景的高分辨率合成,目前主要由基于似然的模型主导,这些模型可能包含数十亿参数的自回归(AR)变换器。基于去噪自动编码器层次结构的扩散模型在图像合成及其他领域展示了令人印象深刻的结果,并在类别条件图像合成和超分辨率方面定义了最先进的水平。作为基于似然的模型,它们不会出现模式崩溃和训练不稳定性问题,并且通过大量利用参数共享,它们可以建模自然图像的高度复杂分布。

民主化高分辨率图像合成扩散模型。为了提高这个强大模型类的可访问性,同时减少其显著的资源消耗,需要一种方法来降低训练和采样的计算复杂性。因此,在不损害其性能的情况下降低 DMs 的计算需求是提高其可访问性的关键。

转向潜在空间

方法从分析已经在像素空间中训练的扩散模型开始:图2展示了训练模型的速率-失真权衡。与任何基于似然的模型一样,学习可以大致分为两个阶段:首先是感知压缩阶段,该阶段去除了高频细节,但仍然学习到很少的语义变化。在第二阶段,实际的生成模型学习数据的语义和概念组成(语义压缩)。因此,目标是首先找到一个感知上等效但计算上更合适的空间,在这个空间中将训练扩散模型用于高分辨率图像合成。

这种方法的一个显著优势是只需要训练一次通用的自动编码阶段,因此可以为多个扩散模型训练重用它,或者探索可能完全不同的任务。这使得能够高效地探索大量用于各种图像到图像和文本到图像任务的扩散模型。对于后者,设计了一种架构,将变换器连接到扩散模型的 UNet 主干,并启用了任意类型的基于令牌的条件机制。

总的来说,工作做出了以下贡献:

1、与纯粹基于 Transformer 的方法相比,方法在处理高维数据时更具扩展性,因此可以:

    - 在压缩级别上提供比以前的工作更忠实和详细的重建(见图 1)。

    - 有效地应用于百万像素图像的高分辨率合成。

2、在多个任务(无条件图像合成、修复、随机超分辨率)和数据集上取得了有竞争力的性能,同时显著降低了计算成本。与基于像素的扩散方法相比,还显著降低了推理成本。

3. 表明,与以前同时学习编码器/解码器架构和基于分数的先验的工作不同,方法不需要对重建和生成能力进行微妙的加权。这确保了极其忠实的重建,并且对潜在空间的正则化要求非常低。

4. 发现,对于密集条件任务,如超分辨率、修复和语义合成,模型可以以卷积方式应用,并渲染出约 10242 像素的大而一致的图像。

5. 此外,设计了一种基于交叉注意力的通用条件机制,实现了多模态训练。用它来训练类别条件、文本到图像和布局到图像模型。

6. 最后,在 https://github.com/CompVis/latent-diffusion 上发布了预训练的潜在扩散和自动编码模型,这些模型可能可用于训练 DMs 之外的各种任务。

2、相关工作

扩散概率模型 (DM)在密度估计和样本质量方面达到了最先进的水平。这些模型的生成能力源于当其底层神经骨干实现为 UNet 时,自然适合类似图像数据的归纳偏置。通常在使用重加权目标进行训练时,可以达到最佳合成质量。在这种情况下,DM 对应于有损压缩器,并允许在图像质量和压缩能力之间进行权衡。然而,在像素空间中评估和优化这些模型存在推理速度低和训练成本高的缺点。虽然前者可以通过高级采样策略和分层方法部分解决,但在高分辨率图像数据上训练总是需要计算昂贵的梯度。通过提出的 LDMs 解决了这两个缺点,LDMs 在较低维度的压缩潜在空间上工作。这使得训练在计算上更便宜,并在几乎不降低合成质量的情况下加快了推理速度(见图 1)。

提出的 LDMs 由于其卷积骨干,在更高维度的潜在空间上更温和地扩展。因此,可以自由选择压缩级别,在保证高保真重建的同时,最优地调解学习强大的第一阶段,而不会将过多的感知压缩留给生成扩散模型(见图 1)。

2、方法

为了降低训练用于高分辨率图像合成的扩散模型的计算需求,观察到,尽管扩散模型允许通过减少相应的损失项来忽略感知上无关的细节,但它们仍然需要在像素空间中进行昂贵的函数评估,这导致在计算时间和能源资源方面有巨大的需求。

提出通过引入压缩学习阶段与生成学习阶段的明确分离来规避这一缺点(见图 2)。为了实现这一点,我们利用一个自动编码模型,该模型学习一个在感知上等同于图像空间的空间,但提供了显著降低的计算复杂性。

这种方法提供了几个优势:

1. 通过离开高维图像空间,获得了计算上更加高效的扩散模型,因为采样是在低维空间中进行的。

2. 利用了扩散模型从其 UNet 架构继承的归纳偏置,这使得它们在处理具有空间结构的数据时特别有效,从而减轻了以前方法所需的那种激进且降低质量的压缩水平。

3. 最后,获得了通用的压缩模型,其潜在空间可以用于训练多个生成模型,并且还可以用于其他下游应用,例如单图像 CLIP 引导的合成。

3.1 感知图像压缩(Perceptual Image Compression)是一种图像压缩方法,它通过利用自编码器(autoencoder)对图像进行编码和解码,以实现图像的压缩。在这个过程中,编码器将图像编码为潜在表示(latent representation),解码器则从潜在表示中重建图像。

具体来说,给定一个 RGB 空间中的图像 x∈R^H×W×3,编码器 E 将 x 编码为潜在表示 z=E(x),解码器 D 从潜在表示中重建图像,得到 ˜x=D(z)=D(E(x))。其中,z∈R^h×w×c,编码器将图像下采样了一个因子 f=H/h=W/w。

为了避免潜在表示空间的任意高方差,作者实验了两种不同的正则化方法:KL-reg. 和 VQ-reg.。KL-reg. 类似于 VAE,对学习到的潜在表示施加了轻微的 KL 惩罚,使其趋向于标准正态分布;VQ-reg. 则在解码器中使用了向量量化层(vector quantization layer)。

与之前的工作不同,作者的压缩模型利用了学习到的潜在空间 z=E(x)的二维结构,因此可以使用相对温和的压缩率,并获得非常好的重建效果。这是因为之前的工作依赖于对学习到的空间 z 的任意 1D 排序来对其分布进行自回归建模,从而忽略了 z 的固有结构。

总的来说,感知图像压缩通过利用自编码器和正则化方法,实现了对图像的高效压缩,同时保留了图像的细节信息。

3.2 扩散模型(Diffusion Models)是一种概率模型,旨在通过逐步去噪一个正态分布的变量来学习数据分布 p(x),这相当于学习一个固定马尔可夫链的反向过程,长度为 T。

对于图像合成,最成功的模型依赖于 p(x) 的变分下界的重加权变体,这类似于去噪分数匹配。这些模型可以解释为一系列等权重的去噪自编码器 ǫθ(xt, t);t = 1...T,它们被训练来预测其输入 xt 的去噪变体,其中 xt 是输入 x 的带噪版本。相应的目标函数可以简化为

其中 t 从 {1, ..., T} 中均匀采样。

通过训练好的感知压缩模型(由 E 和 D 组成),现在可以访问一个高效的、低维度的潜在空间,在这个空间中,高频的、不可察觉的细节被抽象掉了。与高维的像素空间相比,这个空间更适合基于似然的生成模型,因为它们现在可以(i)专注于数据的重要语义部分,(ii)在更低维度、计算效率更高的空间中进行训练。

与之前依赖于在高度压缩的离散潜在空间中使用自回归、基于注意力的变换器模型的工作 不同,可以利用模型提供的特定于图像的归纳偏置。这包括主要使用 2D 卷积层构建底层 UNet 的能力,并进一步使用重加权边界聚焦于感知上最相关的部分,现在这个边界可以表示为

模型的神经骨干 ǫθ(◦, t) 被实现为一个时间条件 UNet [69]。由于正向过程是固定的,zt 可以在训练期间从 E 中高效地获得,并且从 p(z) 采样的样本可以通过 D 的一次传递解码到图像空间。

条件机制(Conditioning Mechanisms)

与其他类型的生成模型类似,扩散模型原则上能够建模形式为 p(z|y) 的条件分布。这可以通过条件去噪自编码器 ǫθ(zt, t, y) 来实现,并为通过输入 y(如文本、语义图或其他图像)到图像的翻译任务控制合成过程铺平了道路。

然而,在图像合成的背景下,将扩散模型的生成能力与其他类型的条件(除了类别标签或输入图像的模糊变体)结合在一起,仍然是一个研究不足的领域。

通过增强其底层 UNet 骨干网与交叉注意力机制,将扩散模型转变为更灵活的条件图像生成器,这种机制对于学习基于注意力的各种输入模态模型是有效的。为了从各种模态(如语言提示)预处理 y,我们引入了一个域特定的编码器 τθ,它将 y 投影到一个中间表示 τθ(y) ∈ RM×dτ,然后通过交叉注意力层映射到 UNet 的中间层,实现 Attention(Q, K, V)

 表示一个(扁平化的)中间表示,实现 UNet。

是可学习的投影矩阵。参见图 3 的视觉描述。基于图像条件对,我们通过学习条件 LDM 来实现。

其中,τθ 和 ǫθ 通过方程 3 联合优化。这种条件机制是灵活的,因为 τθ 可以用特定领域的专家参数化,例如,当 y 是文本提示时,可以使用(未掩码的)变换器。

4、实验部分

LDM 提供了灵活且计算上可行的基于扩散的图像合成方法,还包括各种图像模态的高分辨率生成,在以下部分通过实验验证了这一点。然而,首先,分析了模型与基于像素的扩散模型在训练和推理方面的收益。有趣的是,发现在 VQ 正则化的潜在空间中训练的 LDM 能够获得更好的样本质量,尽管 VQ 正则化的第一阶段模型的重建能力略逊于其连续的对应模型。

4.1 关于感知压缩的权衡

分析具有不同下采样因子(缩写为 LDM-f,其中 LDM-1 对应基于像素的 DMs)的 LDM 的行为。为了获得可比较的测试场,在本节的所有实验中,将计算资源固定为单个 NVIDIA A100,并对所有模型进行相同数量的训练步骤和相同数量的参数训练。

图 5 显示了在 ImageNet 数据集上进行 200 万步的类条件模型的样本质量与训练进度的关系。可以看出:

1. 对于 LDM-{1,2},较小的下采样因子导致训练进度缓慢;

2. 过大的f值会导致在相对较少的训练步骤后保真度停滞不前。

回顾上述分析(图 1 和图 2),这归因于:

1. 将大部分感知压缩留给扩散模型;

2. 第一阶段压缩过强导致信息丢失,从而限制了可实现的质量。

LDM-{4-16}在效率和感知忠实结果之间取得了良好的平衡,在 200 万训练步骤后,基于像素的扩散(LDM-1)和 LDM-8 之间的 FID 差距为 38。

在图 6 中,比较了在 CelebA-HQ 和 ImageNet 上训练的模型在不同去噪步骤数量下的采样速度与 DDIM 采样器的关系,并将其与 FID 分数绘制在一起。LDM-{4-8}的表现优于具有不合适的感知和概念压缩比的模型。特别是与基于像素的 LDM-1 相比,它们在同时显著提高样本吞吐量的同时实现了更低的 FID 分数。像 ImageNet 这样的复杂数据集需要降低压缩率以避免降低质量。

综上所述,观察到 LDM-4 和 LDM-8 处于实现高质量合成结果的最佳行为状态。

4.2 用潜在扩散模型生成图像

训练了在 CelebA-HQ、FFHQ、LSUN-Churches 和 -Bedrooms数据集上的 2562 分辨率的无条件模型,并使用 FID和 Precision-and-Recall 评估了 i) 样本质量和 ii) 数据流形的覆盖率。表 1 总结了结果。

在 CelebA-HQ 上,报告了一个新的最先进的 FID 分数 5.11,超过了之前的基于似然性的模型和 GAN。在 LSGM上也表现出色,其中潜在扩散模型与第一阶段模型联合训练。相比之下,我们在固定空间中训练扩散模型,并避免了在潜在空间的重建质量和学习先验之间的权衡难度,参见图 1-2。

在除 LSUN-Bedrooms 数据集外的所有数据集上都优于之前的基于扩散的方法,尽管我们使用的参数数量是其一半,并且所需的训练资源减少了 4 倍。此外,LDM 在 Precision 和 Recall 上始终优于基于 GAN 的方法,从而证实了其基于覆盖模式似然性的训练目标相对于对抗性方法的优势。在图 4 中,我们还展示了每个数据集的定性结果。

4.3 条件潜在扩散模型

4.3.1 LDM中的 Transformer 编码器

将基于交叉注意力的条件机制引入潜在扩散模型(LDMs)。通过这种方式,LDMs 可以用于各种条件模态,如文本到图像建模。作者在 LAION-400M 数据集上训练了一个基于语言提示的 1.45B 参数模型,并使用 BERT-tokenizer 和 Transformer 来推断潜在代码,然后通过交叉注意力将其映射到 UNet。这种结合特定领域专家的方法在学习语言表示和视觉合成方面效果显著,能够很好地推广到复杂的用户定义文本提示。此外,作者还在其他数据集上训练了模型,以评估基于交叉注意力的条件机制的灵活性。最后,作者还评估了他们在 ImageNet 上的最佳性能类条件模型。

4.3.2 以上的卷积采样

LDMs 作为高效通用图像到图像翻译模型的应用。通过将空间对齐的条件信息连接到 ǫθ 的输入,LDMs 可用于语义合成、超分辨率和修复等任务。在语义合成中,作者使用了与语义图配对的景观图像,并将下采样的语义图与 f=4 模型的潜在图像表示连接起来。他们发现模型可以推广到更大的分辨率,并以卷积方式生成高达兆像素级别的图像。作者还利用这一特性将超分辨率模型和修复模型应用于生成 512² 到 1024² 之间的大图像。在这个应用中,信号噪声比(由潜在空间的尺度引起)会显著影响结果。

4.4 具有潜在扩散的超分辨率

LDMs 在超分辨率任务中的应用。通过直接将低分辨率图像作为条件进行连接,LDMs 可以高效地训练超分辨率模型。在第一个实验中,作者遵循 SR3 的方法,将图像退化固定为 4× 下采样的双三次插值,并在 ImageNet 上训练,使用与 SR3 相同的数据处理管道。他们使用了在OpenImages 上预训练的 f=4 自动编码模型,并将低分辨率条件 y 和 UNet 的输入连接起来,即 τθ 是恒等射。

作者的定性和定量结果(见图 9 和表 4)显示了有竞争力的性能,LDM-SR 在 FID 上优于 SR3,而 SR3 在 IS 上表现更好。一个简单的图像回归模型在 PSNR 和 SSIM 分数上表现最好;然而,这些指标与人类感知不太一致,并且更倾向于模糊而不是不完全对齐的高频细节。

此外,作者进行了一项用户研究,将像素基线与 LDM-SR 进行比较。他们遵循 SR3 的方法,向人类受试者展示一个低分辨率图像和两个高分辨率图像,并询问偏好。表 5 的结果证实了 LDM-SR 的良好性能。通过使用后处理引导机制,可以提高 PSNR 和 SSIM,作者通过感知损失实现了基于图像的引导器。

由于双三次退化过程不能很好地推广到不符合此预处理的图像,作者还训练了一个通用模型 LDM-BSR,使用更多样化的退化。

4.5 用潜在扩散模型修复图片

使用潜在扩散模型(LDM)进行图像修复的,与其与其他专门化、最先进的修复方法的性能比较。结果表明,LDM 在图像修复任务中具有竞争力,并且通过引入注意力机制和更大的模型,可以进一步提高性能。

5、结论

潜在扩散模型(Latent Diffusion Models, LDM)的主要贡献和实验结果。

主要贡献

-简单且高效:LDM 提供了一种简单且高效的方法,显著提高了去噪扩散模型的训练和采样效率,同时不降低其质量。

-交叉注意力条件机制:基于 LDM 和交叉注意力条件机制,作者的实验展示了在各种条件图像合成任务中与最先进方法相比的有利结果,而无需特定任务的架构。

实验结果

-广泛的任务:LDM 在多种条件图像合成任务中表现出色,包括但不限于图像修复、超分辨率、语义合成等。

-性能优越:实验结果表明,LDM 在这些任务中与最先进的方法相比具有竞争力,甚至在某些指标上表现更好。

结论

LDM 通过在潜在空间中进行操作,克服了传统基于像素的扩散模型在训练和推理方面的效率问题。通过引入交叉注意力机制,LDM 增强了其对不同输入模态的适应性,并在多种图像合成任务中展示了优越的性能。

未来工作要进一步优化:尽管 LDM 已经展示了显著的性能提升,但仍有进一步优化的空间,特别是在处理更大规模和更复杂的图像数据时。

扩展应用:LDM 的灵活性和高效性使其可以应用于更多领域,如视频生成、3D 模型生成等。

Stable Diffusion是在Latent Diffusion基础上实现的,相比 Latent diffusion 主要有以下几点改进:

训练数据:Latent Diffusion 是采用 Laion-400M 数据训练的,而 Stable Diffusion 是在 Laion-2B-en数据集上训练出来的。明显后者用了更多的训练数据,而且后者还采用了数据筛选来提升样本数据质量,即采用了 Laion-Aesthetic,一个 120M 的训练子集,以专门选择美学评分较高的图像用于文生图任务。

Text Encoder:Latent Diffusion 采用一个随机初始化的 Transformer 来编码 text。而 Stable Diffusion 采用一个预训练好的 Clip text encoder 来编码 text,预训练的 text model 往往要优于从零开始训练的模型。

训练尺寸: Latent Diffusion 是在 256x256 分辨率数据集上训练,而 Stable Diffusion 是先在256x256分辨率上预训练,然后再在 512x512 分辨率上微调优化的,而现在 XL 版本则采用了 Laion-High-Resolution 训练集,一个规模为 170M,图像分辨率大于 1024 的高分辨率训练子集,用于超分辨率任务。

节点的活跃度:Stable diffusion模型假设节点的活跃度是固定不变的,而latent diffusion模型则假设节点的影响力是动态可变的。此外,stable diffusion模型假设网络是稳定的,而latent diffusion模型则假设网络是动态的。这些区别意味着,在使用stable diffusion模型时,我们假设信息的传播取决于节点的活跃度,并且不考虑节点之间的影响力,而在使用latent diffusion模型时,我们假设信息的传播取决于节点的影响力,并且考虑节点之间的影响力。

下面介绍基于StableDiffusion源码进行训练的方法步骤:

1、下载源码与安装环境

Git clone https://github.com/CompVis/stable-diffusion.git

cd stable-diffusion-main

pip install -r environment.yaml

2、用预训练的模型推理

下载模型文件,这里下载的是stable-diffusion-v1-4(在一个文件夹下包含很多文件,都要下载,有21.5G,上载后放到同一个文件夹下),下载页面地址:

https://huggingface.co/CompVis/stable-diffusion-v1-4/tree/main

其他版本的模型文件,可以在https://huggingface.co/models中搜索stable diffusion,查询到后自行选择

文生图,推理代码如下:

from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained(

    "F:/diffusion_model/stable-diffusion-v1-4",

    use_auth_token=True

).to("cuda")

def generate(prompt):

    result=pipe(prompt)

    image = pipe(prompt)["images"][0]

    image.save("txt2im.png")

    return "txt2im.png"

图生图,推理代码如下:

from diffusers import StableDiffusionImg2ImgPipeline

import torch

from PIL import Image

# 初始化模型

pipe = StableDiffusionImg2ImgPipeline.from_pretrained(

    "F:/diffusion_model/stable-diffusion-v1-4",

    use_auth_token=True

).to("cuda")

def generate_img2img(image_path, text):

    # 定义文本提示

    prompt = text

    # 加载要修改的图片

    init_image = Image.open(image_path).convert("RGB")

    # 使用模型进行图像到图像的转换

    result = pipe(prompt=prompt, image=init_image)

    # 获取生成的图像

    generated_image = result["images"][0]

    # 保存生成的图像

    generated_image.save("img2img.png")

return "img2img.png"

如果你用python3.9或更早版本,运行时候可能会报如下错误:

def is_xpu_device(device: str | torch.device = None):报错:Python version 3.9 does not allow writing union types as X | Y:

需要对这个函数的参数表达方式作如下修改:

def is_xpu_device(device: Union[str, torch.device] = None):

3、微调训练

参考另外一个CSDN作者的文章:手把手教你微调Stable Diffusion_stable diffusion微调-CSDN博客

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值