秒速出图！体验 TensorRT 加速 Stable Diffusion 图像创作

作者：顾静

TensorRT 如何加速 Stable Diffusion？

生成式 AI 图像内容生成技术近年来发展迅速，可以根据人类语言描述生成图片，在时尚、建筑、动漫、广告、游戏等领域有着广泛应用。

Stable Diffusion WebUI 是 Github 上最为热门的利用生成式 AI 进行图像生成的项目。它采用 ClipText 对文字进行编码，然后采用 UNet+Scheduler 在潜在表示空间（latent space）上进行 Diffusion，最后采用 Autoencoder Decoder 将第二步生成的扩散信息再转为图像。

Stable Diffusion PipelineDiffusion 模型最大的痛点是生成图片的速度过慢。Stable Diffusion 采用了多种方式加速图片生成，令实时图像生成成为可能。Stable Diffusion 使用编码器将图片从 3512512 转为 46464，极大地降低了计算量。它在潜在表示空间（latent space）上进行 Diffusion 过程，大大减少计算复杂度，同时也能保证不错的图片生成效果。在 GPU 上大约 4s 可以生成一张描述复杂的图片。

对于很多 ToC 应用来说，一张图片生成需要 4s 耗时仍然过长。TensorRT 是由 NVIDIA 提供的高性能深度学习推理框架，通过优化编译器和 runtime 来提升延迟敏感应用的并发度。TensorRT 可以优化几乎所有的深度神经网络，包括 CNN、RNN 和 Transformer。具体优化措施包括以下 6 点。

减少混合精度，支持 FP32、TF32、FP16 和 INT8
优化 GPU 内存带宽
自动调整核函数，为目标 GPU 选择最佳的算法
动态 Tensor 内存分配，提升内存使用效率
可扩展设计以处理多个计算流
时间融合：随着时间步长优化 RNN

TensorRT 的基本流程如下图所示，可以分为构建期和运行期。

TensorRT Pipeline

基于阿里云容器服务 ACK 的实战体验

云原生 AI 套件

云原生 AI 套件是阿里云容器服务 ACK 提供的云原生 AI 技术和产品方案，帮助企业更快、更高效地落地云原生 AI 系统。

本文将介绍如何基于阿里云容器服务 ACK 云原生 AI 套件，利用 TensorRT 加速 Stable Diffusion 图像生成。

环境配置

参考文档安装云原生 AI 套件 [ 1] 。
登陆容器服务管理控制台 [ 2] ，在左侧导航栏选择集群 > 应用 > 云原生 AI 套件。 等待开发控制台准备就绪后，单击开发控制台。
在开发控制台左侧，选择 Notebook，在 Notebook 页面右上角，单击创建 Notebook 创建新的 Notebook 环境。Notebook 资源需要 CPU：16C，内存：32G，GPU 显存：16GB。

准备 Stable Diffusion + TensorRT 环境

在新建的 Notebook 中输入以下命令安装所需依赖。

! pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple
! pip install --upgrade "torch <2.0.0"
! pip install --upgrade "tensorrt>=8.6"
! pip install --upgrade "accelerate" "diffusers==0.21.4" "transformers"
! pip install --extra-index-url https://pypi.ngc.nvidia.com --upgrade "onnx-graphsurgeon" "onnxruntime" "polygraphy"
! pip install  polygraphy==0.47.1 -i https://pypi.ngc.nvidia.com

下载数据集

import diffusers
import torch
import tensorrt
from diffusers.pipelines.stable_diffusion import StableDiffusionPipeline
from diffusers import DDIMScheduler

# 默认从huggingface下载，如果机器无法访问huggingface，也可以使用本地模型。
# 如使用本地模型，替换runwayml/stable-diffusion-v1-5为本地模型地址即可
model_path = "runwayml/stable-diffusion-v1-5"
scheduler = DDIMScheduler.from_pretrained(model_path, subfolder="scheduler")

使用 TensorRT 生成序列化网络 (计算图 TRT 的内部表示)

# 使用自定义的pipeline
pipe_trt = StableDiffusionPipeline.from_pretrained(
    model_path,
    custom_pipeline="stable_diffusion_tensorrt_txt2img",
    revision='fp16',
    torch_dtype=torch.float16,
    scheduler=scheduler,
)

# 设置缓存地址
# 会在缓存地址下生成engine文件夹，包含clip.plan、unet.plan、vae.plan文件。A10上首次生成plan文件需要约35分钟
pipe_trt.set_cached_folder(model_path, revision='fp16')
pipe_trt = pipe_trt.to("cuda")

使用编译后的模型进行推理

# 生成图片
prompt = "A beautiful ship is floating in the clouds, unreal engine, cozy indoor lighting, artstation, detailed, digital painting, cinematic"
neg_prompt = "ugly"

import time
start_time = time.time()
image = pipe_trt(prompt, negative_prompt=neg_prompt).images[0]
end_time = time.time()
print("time: "+str(round(end_time-start_time, 2))+"s")
display(image)

生成单张图片耗时为 2.31s。