Stable Diffusion Turbo ONNX Runtime CUDA 版本

简介

本仓库托管了 SD Turbo 的优化 ONNX 模型，用于在 ONNX Runtime CUDA 执行提供商下加速推理，适用于 Nvidia GPU。该模型无法在 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 计算能力（如 RTX 4090 为 89，RTX 3090 为 86）在命令行中修改 CMAKE_CUDA_ARCHITECTURES=80。如果您的机器内存小于 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