作者：顾静
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。
性能测试
性能测试基于 Github 上的 lambda-diffusers [ 3] 项目，prompt 数量为 1，batch_size 为 50，重复 100 次。GPU 型号为 A10，对应的 ECS 规格为 ecs.gn7i-c8g1.2xlarge。
![图片]()
![图片]()
从实验结果看，开启 xformers 和 TensorRT 优化后，Stable Diffusion 图片生成时间平均减少 44.7%，显存减少 37.6%。
欢迎扫码添加阿里云 ACK 云原生 AI 套件微信小助手邀您进入云原生 AI 套件微信交流群，或扫描右侧钉钉群码进入钉钉群。
![图片]()
参考文献：
TensorRT
https://github.com/NVIDIA/TensorRT
Stable Diffusion WebUI
https://github.com/AUTOMATIC1111/stable-diffusion-webui
利用 GPU 加速生成式 AI 图像内容生成
https://www.nvidia.cn/webinars/sessions/?session_id=230919-29256
相关链接：
[1] 安装云原生 AI 套件
https://help.aliyun.com/zh/ack/cloud-native-ai-suite/user-guide/deploy-the-cloud-native-ai-suite?spm=a2c4g.11186623.0.0.7e223d92U1aVNf
[2] 容器服务管理控制台
https://account.aliyun.com/login/login.htm?oauth_callback=https%3A%2F%2Fcs.console.aliyun.com%2F
[3] lambda-diffusers
https://github.com/LambdaLabsML/lambda-diffusers
