适用于 ONNX Runtime CUDA 的 Stable Diffusion Turbo

简介

本仓库托管了 SD Turbo 的优化 ONNX 模型，用于加速在 Nvidia GPU 上使用 ONNX Runtime CUDA 执行提供程序的推理。它无法在 CPU 和 DirectML 等其他提供程序上运行。

这些模型由 Olive 使用如下命令生成：

python stable_diffusion.py --provider cuda --model_id stabilityai/sd-turbo --optimize --use_fp16_fixed_vae

有关如何使用本仓库中托管的 ONNX 文件运行 SDXL 流水线，请参阅使用说明。

模型描述

• 开发者： Stability AI
• 模型类型： 基于扩散的文本到图像生成模型
• 许可证： STABILITY AI 非商业研究社区许可证
• 模型描述： 这是 SD-Turbo 模型的转换版本，用于在 ONNX Runtime 上使用 CUDA 执行提供程序进行推理。

性能

延迟

以下是在 NVIDIA A100-SXM4-80GB GPU 上生成 512x512 尺寸图像的平均延迟：

Static 表示引擎是为给定的批大小和图像大小组合构建的，并使用 CUDA graph 来加速。

使用示例

按照演示说明操作。示例步骤如下：

0. 按照这些说明安装 nvidia-docker。

1. 克隆 onnxruntime 仓库。

git clone https://github.com/microsoft/onnxruntime
cd onnxruntime

2. 从本仓库下载 ONNX 文件

git lfs install
git clone https://huggingface.co/tlwu/sd-turbo-onnxruntime

3. 启动 docker

docker run --rm -it --gpus all -v $PWD:/workspace nvcr.io/nvidia/pytorch:23.10-py3 /bin/bash

4. 从源码构建 ONNX Runtime

export CUDACXX=/usr/local/cuda-12.2/bin/nvcc
git config --global --add safe.directory '*'
sh build.sh --config Release  --build_shared_lib --parallel --use_cuda --cuda_version 12.2 \
            --cuda_home /usr/local/cuda-12.2 --cudnn_home /usr/lib/x86_64-linux-gnu/ --build_wheel --skip_tests \
            --use_tensorrt --tensorrt_home /usr/src/tensorrt \
            --cmake_extra_defines onnxruntime_BUILD_UNIT_TESTS=OFF \
            --cmake_extra_defines CMAKE_CUDA_ARCHITECTURES=80 \
            --allow_running_as_root
python3 -m pip install build/Linux/Release/dist/onnxruntime_gpu-*-cp310-cp310-linux_x86_64.whl --force-reinstall

如果 GPU 不是 A100，请根据 GPU 计算能力在命令行中更改 CMAKE_CUDA_ARCHITECTURES=80（例如 RTX 4090 为 89，RTX 3090 为 86）。如果您的机器内存小于 64GB，请将 --parallel 替换为 --parallel 4 --nvcc_threads 1 以避免内存不足。

5. 安装库和依赖项

python3 -m pip install --upgrade pip
cd /workspace/onnxruntime/python/tools/transformers/models/stable_diffusion
python3 -m pip install -r requirements-cuda12.txt
python3 -m pip install --upgrade polygraphy onnx-graphsurgeon --extra-index-url https://pypi.ngc.nvidia.com

6. 执行 ONNX Runtime 优化推理

python3 demo_txt2img.py \
  \"starry night over Golden Gate Bridge by van gogh\" \
  --version sd-turbo   \
  --engine-dir /workspace/sd-turbo-onnxruntime