秒速出图！体验 TensorRT 加速 Stable Diffusion 图像创作 - 服务器托管|北京服务器租用|机房托管租用|IDC托管租用|机房机柜带宽租用-价格及费用咨询

作者：顾静

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 =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/AUTOMATI服务器托管网C1111/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

服务器托管，北京服务器托管，服务器租用 http://www.fwqtg.net

相关推荐: Linux 下最主流的文件系统格式——ext

硬盘分成相同大小的单元，我们称为块（Block）。一块的大小是扇区大小的整数倍，默认是 4K。在格式化的时候，这个值是可以设定的。一大块硬盘被分成了一个个小的块，用来存放文件的数据部分。这样一来，如果我们像存放一个文件，就不用给他分配一块连续的空间了。我们可…