使用 KerasCV 中的 Stable Diffusion 进行高性能图像生成

作者: fchollet, lukewood, divamgupta
创建日期: 2022/09/25
最后修改日期: 2022/09/25
描述: 使用 KerasCV 的 Stable Diffusion 模型生成新图像。

在 Colab 中查看 • GitHub 源代码

概述

在本教程中，我们将展示如何使用 KerasCV 实现的 stability.ai's 文本到图像模型 Stable Diffusion 根据文本提示生成新图像。

Stable Diffusion 是一个强大、开源的文本到图像生成模型。尽管存在多个开源实现，允许你轻松地从文本提示创建图像，但 KerasCV 提供了一些独特的优势。这些包括 XLA 编译 和 混合精度 支持，它们结合在一起实现了最先进的生成速度。

在本指南中，我们将探索 KerasCV 的 Stable Diffusion 实现，展示如何利用这些强大的性能提升，并探讨它们带来的性能优势。

注意: 要在 torch 后端运行本教程，请在每个地方设置 jit_compile=False。Stable Diffusion 的 XLA 编译目前不支持 torch。

要开始，先安装一些依赖并整理一些导入：

!pip install -q --upgrade keras-cv
!pip install -q --upgrade keras  # 升级到 Keras 3.

import time
import keras_cv
import keras
import matplotlib.pyplot as plt

介绍

与大多数教程不同，我们通常首先解释一个主题，然后展示如何实现它，文本到图像生成更容易展现而不是叙述。

看看 keras_cv.models.StableDiffusion() 的强大功能。

首先，我们构建一个模型：

model = keras_cv.models.StableDiffusion(
    img_width=512, img_height=512, jit_compile=False
)

通过使用该模型检查点，您承认其使用受 CreativeML Open RAIL-M 许可条款的约束，详见 https://raw.githubusercontent.com/CompVis/stable-diffusion/main/LICENSE

接下来，我们给它一个提示：

images = model.text_to_image("photograph of an astronaut riding a horse", batch_size=3)


def plot_images(images):
    plt.figure(figsize=(20, 20))
    for i in range(len(images)):
        ax = plt.subplot(1, len(images), i + 1)
        plt.imshow(images[i])
        plt.axis("off")


plot_images(images)

 50/50 ━━━━━━━━━━━━━━━━━━━━ 63s 211ms/step

真是令人难以置信！

但这不是该模型唯一能做的事情。我们试试一个更复杂的提示：

images = model.text_to_image(
    "cute magical flying dog, fantasy art, "
    "golden color, high quality, highly detailed, elegant, sharp focus, "
    "concept art, character concepts, digital painting, mystery, adventure",
    batch_size=3,
)
plot_images(images)

 50/50 ━━━━━━━━━━━━━━━━━━━━ 10s 209ms/step

可能性真是无穷无尽（或者至少延伸到 Stable Diffusion 的潜在流形的边界）。

等等，这到底是怎么工作的？

与您此时可能期待的不同，Stable Diffusion 并不是真的在魔法上运行。它是一种“潜在扩散模型”。让我们深入探讨一下这意味着什么。

您可能熟悉 超分辨率 的概念：可以训练一个深度学习模型来 去噪 输入图像——从而将其转变为更高分辨率的版本。深度学习模型并不是通过神奇地恢复来自嘈杂、低分辨率输入的信息来完成的——实际上，模型利用其训练数据分布来幻觉输入下最可能的视觉细节。要了解更多关于超分辨率的信息，您可以查看以下 Keras.io 教程：

• 使用高效的子像素 CNN 进行图像超分辨率
• 增强的深度残差网络用于单图像超分辨率

当你将这个想法推向极限时，你可能会开始问——如果我们仅仅在纯噪声上运行这样的模型会怎么样？然后模型将“去噪声”，并开始幻觉出一张全新的图像。通过多次重复这个过程，你可以把一小块噪声变成越来越清晰和高分辨率的人工图片。

这是潜在扩散的关键思想，提出于2020年的 高分辨率图像合成与潜在扩散模型。 要深入理解扩散，你可以查看 Keras.io 教程 去噪声扩散隐式模型。

现在，要从潜在扩散转向文本到图像系统，你仍然需要添加一个关键特性：通过提示关键词控制生成的视觉内容的能力。这是通过“条件化”实现的，条件化是一种经典的深度学习技术，它将代表一段文本的向量与噪声片段连接在一起，然后在一个{图像: 标题}对的数据集上训练模型。

这催生了稳定扩散架构。稳定扩散由三个部分组成：

• 一个文本编码器，将你的提示转化为潜在向量。
• 一个扩散模型，反复“去噪”一个64x64的潜在图像片段。
• 一个解码器，将最终64x64的潜在片段转换为更高分辨率的512x512图像。

首先，你的文本提示通过文本编码器投影到潜在向量空间，文本编码器只是一个经过预训练的冻结语言模型。然后那个提示向量被连接到一个随机生成的噪声片段上，扩散模型在一系列“步骤”上反复“去噪”这个片段（你运行的步骤越多，你的图像就越清晰和美观——默认值是50步）。

最后，64x64的潜在图像通过解码器进行处理，以适当地以高分辨率渲染它。

总而言之，这是一个相当简单的系统——Keras实现包含四个文件，总行数不到500行代码：

• text_encoder.py: 87 行代码
• diffusion_model.py: 181 行代码
• decoder.py: 86 行代码
• stable_diffusion.py: 106 行代码

但这个相对简单的系统在处理数十亿张图片及其标题后开始显得像魔法一样。正如费曼所说的关于宇宙：“这并不复杂，只是量太大了！”

KerasCV的优势

在许多公开可用的稳定扩散实现中，为什么你应该使用 keras_cv.models.StableDiffusion？

除了易于使用的API外，KerasCV的稳定扩散模型还具有一些强大的优势，包括：

• 图形模式执行
• 通过jit_compile=True进行XLA编译
• 支持混合精度计算

当这些结合在一起时，KerasCV稳定扩散模型的运行速度比简单实现快几个数量级。本节展示了如何启用所有这些功能，以及使用它们所带来的性能提升。

为了进行比较，我们进行了基准测试，将 HuggingFace扩散器实现的 稳定扩散与KerasCV实现进行比较。 这两个实现都被要求生成3幅图像，每幅图像的步骤数量为50。在本次基准测试中，我们使用了Tesla T4 GPU。

我们所有的基准测试都是开源的，可以在GitHub上运行，并可以在Colab上重新运行以复制结果。 基准测试结果显示在下表中：

在Tesla T4上执行时间提高了30%！虽然在V100上的提升较小，但我们通常期望基准测试的结果在所有NVIDIA GPU上始终支持KerasCV。

为完整起见，报告了冷启动和热启动的生成时间。冷启动执行时间包括模型创建和编译的一次性成本， 因此，在生产环境中这可以忽略不计（在那里你会多次重用相同的模型实例）。尽管如此，这里是冷启动的数字：

尽管执行本指南的运行时间结果可能有所不同，在我们的测试中，KerasCV实现的稳定扩散显著快于其PyTorch对应实现。这主要归因于XLA编译。

注意：每种优化的性能收益在硬件配置之间可能会显著不同。

要开始，我们先基准测试我们的未优化模型：

benchmark_result = []
start = time.time()
images = model.text_to_image(
    "一个可爱的水獺在彩虹漩涡中握着贝壳，水彩画",
    batch_size=3,
)
end = time.time()
benchmark_result.append(["标准", end - start])
plot_images(images)

print(f"标准模型: {(end - start):.2f} 秒")
keras.backend.clear_session()  # 清除会话以节省内存。

 50/50 ━━━━━━━━━━━━━━━━━━━━ 10s 209ms/step
标准模型: 10.57 秒

混合精度

“混合精度”是指使用float16精度进行计算，同时以float32格式存储权重。这是为了利用现代NVIDIA GPU上float16操作的内核比其float32对应物快得多的事实。

在Keras中启用混合精度计算（因此对于keras_cv.models.StableDiffusion）只需调用：

keras.mixed_precision.set_global_policy("mixed_float16")

就这么简单。开箱即用——它可以正常工作。

model = keras_cv.models.StableDiffusion(jit_compile=False)

print("计算数据类型:", model.diffusion_model.compute_dtype)
print(
    "变量数据类型:",
    model.diffusion_model.variable_dtype,
)

通过使用这个模型检查点，您确认其使用受制于CreativeML Open RAIL-M许可证的条款，网址为：https://raw.githubusercontent.com/CompVis/stable-diffusion/main/LICENSE
计算数据类型: float16
变量数据类型: float32

正如你所看到的，上面构造的模型现在使用混合精度计算；利用float16操作的速度进行计算，同时以float32精度存储变量。

# 预热模型以在基准测试之前运行图形追踪。
model.text_to_image("预热模型", batch_size=3)

start = time.time()
images = model.text_to_image(
    "一只可爱的魔法飞狗，幻想艺术，"
    "金色，高质量，高度细致，优雅，清晰聚焦，"
    "概念艺术，角色概念，数字画作，神秘，冒险",
    batch_size=3,
)
end = time.time()
benchmark_result.append(["混合精度", end - start])
plot_images(images)

print(f"混合精度模型: {(end - start):.2f} 秒")
keras.backend.clear_session()

 50/50 ━━━━━━━━━━━━━━━━━━━━ 42s 132ms/step
 50/50 ━━━━━━━━━━━━━━━━━━━━ 6s 129ms/step
混合精度模型: 6.65 秒

XLA 编译

TensorFlow和JAX内置了XLA: 加速线性代数编译器。keras_cv.models.StableDiffusion开箱即用支持jit_compile参数。将该参数设置为True可以启用XLA编译，从而显著加快速度。

让我们在下面使用它：

# 为了基准测试的目的，重设为默认值。
keras.mixed_precision.set_global_policy("float32")

model = keras_cv.models.StableDiffusion(jit_compile=True)
# 在我们基准测试模型之前，我们运行一次推理以确保TensorFlow
# 图形已被追踪。
images = model.text_to_image("一个牛油果沙发", batch_size=3)
plot_images(images)

通过使用这个模型检查点，您确认其使用受制于CreativeML Open RAIL-M许可证的条款，网址为：https://raw.githubusercontent.com/CompVis/stable-diffusion/main/LICENSE
 50/50 ━━━━━━━━━━━━━━━━━━━━ 48s 209ms/step

让我们基准测试我们的XLA模型：

start = time.time()
images = model.text_to_image(
    "一只可爱的水獺在彩虹漩涡中抱着贝壳，水彩画",
    batch_size=3,
)
end = time.time()
benchmark_result.append(["XLA", end - start])
plot_images(images)

print(f"使用 XLA: {(end - start):.2f} 秒")
keras.backend.clear_session()

50/50 ━━━━━━━━━━━━━━━━━━━━ 11s 210ms/step
使用 XLA: 10.63 秒

在 A100 GPU 上，我们获得了大约 2 倍的加速。 太棒了！

整合全部内容

那么，如何组装世界上性能最高的稳定扩散推理管道（截至 2022 年 9 月）。

只需这两行代码：

keras.mixed_precision.set_global_policy("mixed_float16")
model = keras_cv.models.StableDiffusion(jit_compile=True)

使用此模型检查点，即表示您承认其使用受 CreativeML Open RAIL-M 许可证条款的约束，网址为 https://raw.githubusercontent.com/CompVis/stable-diffusion/main/LICENSE

并且要使用它...

# 确保先预热模型
images = model.text_to_image(
    "泰迪熊进行机器学习研究",
    batch_size=3,
)
plot_images(images)

50/50 ━━━━━━━━━━━━━━━━━━━━ 48s 131ms/step

到底有多快？ 让我们来看看！

start = time.time()
images = model.text_to_image(
    "一个神秘的黑暗陌生人造访埃及的大金字塔，"
    "高质量，高度详细，优雅，清晰聚焦，"
    "概念艺术，角色概念，数字绘画",
    batch_size=3,
)
end = time.time()
benchmark_result.append(["XLA + Mixed Precision", end - start])
plot_images(images)

print(f"XLA + mixed precision: {(end - start):.2f} seconds")

50/50 ━━━━━━━━━━━━━━━━━━━━ 6s 130ms/step
XLA + 混合精度: 6.66 秒

让我们检查结果：

print("{:<22} {:<22}".format("模型", "运行时间"))
for result in benchmark_result:
    name, runtime = result
    print("{:<22} {:<22}".format(name, runtime))

模型                  运行时间               
标准                  10.572920799255371    
混合精度              6.651048421859741     
XLA                   10.632121562957764    
XLA + 混合精度        6.659237861633301     

我们经过优化的模型在 A100 GPU 上只需四秒钟就能从文本提示生成三幅新图像。

结论

KerasCV 提供了最先进的稳定扩散实现——并且通过使用 XLA 和混合精度，提供了截至 2022 年 9 月可用的最快稳定扩散管道。

通常，在 keras.io 教程的最后，我们会给您一些未来的方向，以便继续学习。 这次，我们给您留下一个想法：

去通过模型运行您自己的提示！这绝对是个乐趣！

如果您有自己的 NVIDIA GPU，或 M1 MacBookPro，您也可以在本地机器上运行该模型。 （请注意，在 M1 MacBookPro 上运行时，您不应启用混合精度，因为它在苹果的 Metal 运行时中尚未得到良好支持。）