使用 Triton 和 TensorRT 部署 Stable Diffusion 模型#

此示例演示了如何通过利用 TensorRT demo 管道和实用程序在 Triton 中部署 Stable Diffusion 模型。

此示例使用 TensorRT demo 作为基础，包含一个可重用的基于 python 的后端，/backend/diffusion/model.py，适用于部署 Diffusion 模型的多个版本和配置。

有关 Stable Diffusion 的更多信息，请访问 stable-diffusion-v1-5，stable-diffusion-xl。有关 TensorRT 实现的更多信息，请参阅 TensorRT demo。

[!Note] 此示例作为示例代码提供，在生产环境中使用前应进行审核。

| 要求 | 构建服务器镜像 | Stable Diffusion v1.5 | Stable Diffusion XL | 发送推理请求 | 模型配置 | 示例客户端 | 已知问题和限制 |

要求#

以下说明需要安装了 Docker 的 Linux 系统。对于 CUDA 支持，请确保您的 CUDA 驱动程序满足 深度学习框架支持矩阵的“NVIDIA 驱动程序”部分中的要求。

构建 Triton 推理服务器镜像#

该示例基于 nvcr.io/nvidia/tritonserver:24.08-py3 docker 镜像和 TensorRT OSS v10.4 设计。

提供了一组方便的脚本，用于基于 nvcr.io/nvidia/tritonserver:24.01-py3 镜像创建 docker 镜像，并安装了 TensorRT Stable Diffusion demo 的依赖项。

Triton 推理服务器 + TensorRT OSS#

克隆仓库#

git clone https://github.com/triton-inference-server/tutorials.git --single-branch
cd tutorials/Popular_Models_Guide/StableDiffusion

构建 Tritonserver Diffusion Docker 镜像#

./build.sh

包含的模型#

default 构建包括位于 /diffusion-models 文件夹中的模型配置文件。为 stable_diffusion_1_5 和 stable_diffusion_xl 提供了示例配置。

模型工件和引擎文件未包含在镜像中，但在单独的步骤中构建到卷挂载目录中。

构建和运行 Stable Diffusion v 1.5#

启动 Tritonserver Diffusion 容器#

以下命令启动一个容器并将当前目录作为 workspace 卷挂载。

./run.sh

构建 Stable Diffusion v 1.5 引擎#

[!Note]

模型 stable-diffusion-v1-5 需要登录 huggingface 并接受使用条款和条件。请相应地设置环境变量 HF_TOKEN。

./scripts/build_models.sh --model stable_diffusion_1_5

预期输出#

 diffusion-models
|-- stable_diffusion_1_5
|   |-- 1
|   |   |-- 1.5-engine-batch-size-1
|   |   |-- 1.5-onnx
|   |   |-- 1.5-pytorch_model
|   `-- config.pbtxt

启动服务器实例#

[!Note] 我们使用 EXPLICIT 模型控制模式进行演示，以控制加载哪个 stable diffusion 版本。对于生产部署，请参阅 安全部署注意事项，了解有关 EXPLICIT 模式相关风险的更多信息。

tritonserver --model-repository diffusion-models --model-control-mode explicit --load-model stable_diffusion_1_5

预期输出#

<SNIP>
I0229 20:15:52.125050 749 server.cc:676]
+----------------------+---------+--------+
| Model                | Version | Status |
+----------------------+---------+--------+
| stable_diffusion_1_5 | 1       | READY  |
+----------------------+---------+--------+

<SNIP>

构建和运行 Stable Diffusion XL#

启动 Tritonserver Diffusion 容器#

以下命令启动一个容器并将当前目录作为 workspace 卷挂载。

./run.sh

构建 Stable Diffusion XL 引擎#

./scripts/build_models.sh --model stable_diffusion_xl

预期输出#

 diffusion-models
 |-- stable_diffusion_xl
    |-- 1
    |   |-- xl-1.0-engine-batch-size-1
    |   |-- xl-1.0-onnx
    |   `-- xl-1.0-pytorch_model
    `-- config.pbtxt

启动服务器实例#

[!Note] 我们使用 EXPLICIT 模型控制模式进行演示，以控制加载哪个 stable diffusion 版本。对于生产部署，请参阅 安全部署注意事项，了解有关 EXPLICIT 模式相关风险的更多信息。

tritonserver --model-repository diffusion-models --model-control-mode explicit --load-model stable_diffusion_xl

预期输出#

<SNIP>
I0229 20:22:22.912465 1440 server.cc:676]
+---------------------+---------+--------+
| Model               | Version | Status |
+---------------------+---------+--------+
| stable_diffusion_xl | 1       | READY  |
+---------------------+---------+--------+

<SNIP>

发送推理请求#

我们提供了一个示例 client 应用程序，以简化发送和接收请求。

启动 Tritonserver Diffusion 容器#

在与服务器不同的终端中，启动一个新容器。

以下命令启动一个容器并将当前目录作为 workspace 卷挂载。

./run.sh

向 Stable Diffusion 1.5 发送提示#

python3 client.py --model stable_diffusion_1_5 --prompt "butterfly in new york, 4k, realistic" --save-image

示例输出#

Client: 0 Throughput: 0.7201335361144658 Avg. Latency: 1.3677194118499756
Throughput: 0.7163933558221957 Total Time: 1.395881175994873

如果给定 --save-image，则输出图像将保存为 jpeg。

client_0_generated_image_0.jpg

sample_generated_image

向 Stable Diffusion XL 发送提示#

python3 client.py --model stable_diffusion_xl --prompt "butterfly in new york, 4k, realistic" --save-image

示例输出#

Client: 0 Throughput: 0.1825067711674996 Avg. Latency: 5.465569257736206
Throughput: 0.18224859609447058 Total Time: 5.487010717391968

如果给定 --save-image，则输出图像将保存为 jpeg。

client_0_generated_image_0.jpg

sample_generated_image

示例客户端#

示例 client 应用程序使用户能够快速测试不同并发场景下的 diffusion 模型。有关客户端应用程序选项的完整列表和描述，请使用

python3 client.py --help

发送并发请求#

要增加负载和并发性，用户可以使用 clients 和 requests 选项来控制客户端进程的数量和每个客户端发送的请求数量。

示例：十个客户端各自发送十个请求#

以下命令启用十个客户端，每个客户端发送十个请求。每个客户端都是一个独立的进程，与其他九个客户端并行地一个接一个地发送其请求。

python3 client.py --model stable_diffusion_xl --requests 10 --clients 10

已知问题和限制#

  1. 与它所基于的 demo 不同，diffusion 后端尚不支持使用可选的 refiner 模型。另请参阅 demo_txt2img_xl.py

在本页中