使用 ONNX Runtime 和 Olive 加速 SD Turbo 和 SDXL Turbo 推理
介绍
SD Turbo 和 SDXL Turbo 是两个快速的生成式文本到图像模型，能够在短短一步内生成可用的图像，这比以前的 Stable Diffusion 模型通常需要的 30 多步是一个显著的改进。SD Turbo 是 Stable Diffusion 2.1 的蒸馏版本，SDXL Turbo 是 SDXL 1.0 的蒸馏版本。我们之前展示过如何使用 ONNX Runtime 加速 Stable Diffusion 推理。ONNX Runtime 不仅在使用 SD Turbo 和 SDXL Turbo 时提供性能优势，而且还使模型可以在 Python 之外的语言中使用，例如 C# 和 Java。
性能提升
在这篇文章中，我们将介绍 ONNX Runtime CUDA 和 TensorRT 执行提供程序中的优化，这些优化显著加速了 SD Turbo 和 SDXL Turbo 在 NVIDIA GPU 上的推理。
在所有测试的（批次大小，步数）组合中，ONNX Runtime 都优于 PyTorch，SDXL Turbo 模型的吞吐量提升高达 229%，SD Turbo 模型吞吐量提升高达 120%。ONNX Runtime CUDA 在动态形状方面表现出色，但在静态形状方面也比 PyTorch 有显著改进。
如何运行 SD Turbo 和 SDXL Turbo
为了使用 ONNX Runtime CUDA 执行提供程序加速推理，请访问我们在 Hugging Face 上优化的 SD Turbo 和 SDXL Turbo 版本。
这些模型由 Olive 生成，Olive 是一款易于使用的、硬件感知型的模型优化工具。请注意，为了获得最佳性能，必须通过命令行启用 fp16 VAE，如共享的优化版本所示。有关如何使用 Hugging Face 上托管的 ONNX 文件运行 SD 和 SDXL 管道的说明，请参阅 SD Turbo 使用示例 和 SDXL Turbo 使用示例。
要使用 ONNX Runtime TensorRT 执行提供程序加速推理，请按照此处的说明进行操作。
以下是使用 SDXL Turbo 模型并由文本提示引导生成图像的示例
python3 demo_txt2img_xl.py \
--version xl-turbo \
"little cute gremlin wearing a jacket, cinematic, vivid colors, intricate masterpiece, golden ratio, highly detailed"
图 1. 使用 SDXL Turbo 通过文本提示生成的穿着夹克的小可爱小妖精图像。
请注意，示例图像是在 4 个步骤中生成的，这表明 SD Turbo 和 SDXL Turbo 能够以比以前的 Stable Diffusion 模型更少的步骤生成可用的图像。
要以用户友好的方式试用 Stable Diffusion 模型，请参阅我们的 Automatic1111 SD WebUI 的 ONNX Runtime 扩展。此扩展可以在 NVIDIA GPU 上优化执行 Stable Diffusion UNet 模型，并使用 ONNX Runtime CUDA 执行提供程序对经过 Olive 优化的模型进行推理。目前，此扩展仅针对 Stable Diffusion 1.5 进行了优化。SD Turbo 和 SDXL Turbo 模型也可以使用，但性能优化仍在进行中。
C# 和 Java 中 Stable Diffusion 的应用
利用 ONNX Runtime 的跨平台、性能和可用性优势，社区成员也贡献了自己的 Stable Diffusion 示例和 UI 工具，并使用 ONNX Runtime 进行推理。
这些社区贡献包括 OnnxStack，这是一个 .NET 库，它基于我们之前的 C# 教程构建，为用户在使用 C# 和 ONNX Runtime 进行推理时，提供了各种 Stable Diffusion 模型的多种功能。
此外，Oracle 还发布了一个使用 Java 的 Stable Diffusion 示例，它在 ONNX Runtime 之上运行推理。这个项目也基于我们的 C# 教程。
基准测试结果
我们使用 A100-SXM4-80GB 的 Standard_ND96amsr_A100_v4 VM 和配备 RTX-4090 GPU 的 联想台式机（WSL Ubuntu 20.04）对 SD Turbo 和 SDXL Turbo 模型进行了基准测试，以使用 LCM 调度器和 fp16 模型生成 512x512 分辨率的图像。结果根据以下规格测量：
- onnxruntime-gpu==1.17.0（从源代码构建）
- torch==2.1.0a0+32f93b1
- tensorrt==8.6.1
- transformers==4.36.0
- diffusers==0.24.0
- onnx==1.14.1
- onnx-graphsurgeon==0.3.27
- polygraphy==0.49.0
要复现这些结果，我们建议使用“使用示例”部分中链接的说明。
由于 SDXL Turbo 的原始 VAE 无法以 fp16 精度运行，我们在测试 SDXL Turbo 时使用了 sdxl-vae-fp16-fix。它的输出与原始 VAE 略有差异，但解码后的图像对于大多数用途来说足够接近。
静态形状的 PyTorch 管道已应用通道优先内存格式和带 reduce-overhead 模式的 torch.compile。
以下图表显示了每秒图像吞吐量与不同（批次大小，步数）组合下各种框架的对比。值得注意的是，每个条形上方的标签表示相对于 Torch Compile 的加速百分比——例如，在第一个图表中，对于（批次，步数）组合（4，1），ORT_TRT（静态）比 Torch（编译）快 31%。
我们选择使用 1 步和 4 步，因为 SD Turbo 和 SDXL Turbo 都可以在最短 1 步内生成可用的图像，但通常在 3-5 步内生成质量最佳的图像。
SDXL Turbo
下面的图表展示了 SDXL Turbo 模型在静态和动态形状下的每秒图像吞吐量。结果是在 A100-SXM4-80GB GPU 上针对不同的（批次大小，步数）组合收集的。对于动态形状，TensorRT 引擎支持批次大小 1 到 8 和图像大小 512x512 到 768x768，但它针对批次大小 1 和图像大小 512x512 进行了优化。
SD Turbo
接下来的两张图表展示了 SD Turbo 模型在 A100-SXM4-80GB GPU 上静态和动态形状下的每秒图像吞吐量。
最后一组图表展示了 SD Turbo 模型在 RTX-4090 GPU 上静态和动态形状下的每秒图像吞吐量。在此动态形状测试中，TensorRT 引擎是为批次大小 1 到 8（针对批次大小 1 优化）和固定图像大小 512x512 构建的，因为内存限制。
ONNX Runtime 下 SD Turbo 和 SDXL Turbo 的速度如何？
这些结果表明，在静态和动态形状下，ONNX Runtime 在所有所示的（批次，步数）组合中，其 CUDA 和 TensorRT 执行提供程序均显著优于 PyTorch。这个结论适用于两种模型尺寸（SD Turbo 和 SDXL Turbo）以及两种测试的 GPU。值得注意的是，ONNX Runtime 与 CUDA（动态形状）在（批次，步数）组合（1，4）下比 Torch Eager 快 229%。
此外，由于在大多数（批次，步数）组合下 ORT_TRT 吞吐量高于相应的 ORT_CUDA 吞吐量，使用 TensorRT 执行提供程序的 ONNX Runtime 在静态形状方面表现稍好。当用户在图定义时知道批次和图像大小（例如，用户只计划生成批次大小为 1 且图像大小为 512x512 的图像）时，通常首选静态形状。在这些情况下，静态形状具有更快的性能。但是，如果用户决定切换到不同的批次和/或图像大小，TensorRT 必须创建一个新的引擎（这意味着磁盘上的引擎文件数量会增加一倍）并切换引擎（这意味着加载新引擎会花费额外时间）。
另一方面，在使用 A100-SXM4-80GB GPU 时，ONNX Runtime 与 CUDA 执行提供程序通常是 SD Turbo 和 SDXL Turbo 模型动态形状的更好选择，但在使用 RTX-4090 GPU 时，ONNX Runtime 与 TensorRT 执行提供程序在大多数（批次，步数）组合下动态形状的表现稍好。使用动态形状的好处是，当批次和图像大小直到图执行时才可知时（例如，为一个图像运行批次大小为 1 且图像大小为 512x512，为另一个图像运行批次大小为 4 且图像大小为 512x768），用户可以更快地运行推理。在这些情况下使用动态形状时，用户只需构建并保存一个引擎，而无需在推理期间切换引擎。
GPU 优化
除了我们之前的 Stable Diffusion 博客中介绍的技术外，ONNX Runtime 还应用了以下优化措施，以获得本文所述的 SD Turbo 和 SDXL Turbo 结果：
- 启用静态形状输入的 CUDA 图。
- 添加 Flash Attention V2。
- 删除文本编码器中的额外输出（保留由 clip_skip 参数指定的隐藏状态输出）。
- 添加 SkipGroupNorm 融合，将组归一化与其前面的 Add 节点融合。
此外，我们还增加了对新功能的支持，包括用于潜在一致性模型（LCM）的 LoRA 权重。
后续步骤
未来，我们计划继续改进 Stable Diffusion 工作，更新演示以支持新功能，如 IP Adapter 和 Stable Video Diffusion。ControlNet 支持也将很快推出。
我们还在努力优化 SD Turbo 和 SDXL Turbo 在我们现有 Stable Diffusion web UI 扩展上的性能，并计划帮助将这两种模型支持添加到 ONNX Runtime 社区成员开发的 Windows UI 中。
此外，关于如何使用 C# 和 ONNX Runtime 运行 SD Turbo 和 SDXL Turbo 的教程即将发布。在此期间，请查看我们之前关于 Stable Diffusion 的教程。
资源
查看本文讨论的一些资源
- SD Turbo：Olive 优化的 ONNX Runtime CUDA 模型，托管在 Hugging Face 上。
- SDXL Turbo：Olive 优化的 ONNX Runtime CUDA 模型，托管在 Hugging Face 上。
- Stable Diffusion GPU 优化：ONNX Runtime GitHub 仓库中关于使用 NVIDIA GPU 优化 Stable Diffusion 的说明。
- Automatic1111 SD WebUI 的 ONNX Runtime 扩展：启用在 NVIDIA GPU 上优化执行 Stable Diffusion UNet 模型的扩展。
- OnnxStack：社区贡献的 .NET 库，支持使用 C# 和 ONNX Runtime 进行 Stable Diffusion 推理。
- SD4J (Java 中的 Stable Diffusion)：Oracle 提供的 Java 和 ONNX Runtime Stable Diffusion 示例。
- 使用 C# 和 ONNX Runtime 进行 Stable Diffusion 推理：之前发布的 C# 教程。
