文档反馈官招募中，报名立赚积分兑换代金券！> HOT
文档中心>TencentOS Server>工具指南>AI 使用>TencentOS Server 3 运行主要深度学习推理框架及热门模型示例>TensorRT>Stable Diffusion

Stable Diffusion

最近更新时间：2024-08-13 09:55:21

我的收藏

本页目录：

本指导适用于在 TencentOS Server 3上使用 TensorRT 推理框架运行 Stable Diffusion 的官方 Demo，以 Docker 方式启动。

TensorRT 环境准备

1. 从 GitHub 下载 TensorRT 开源仓库到本地。
#下载TensorRT开源仓库到本地，版本为8.6.1
git clone git@github.com:NVIDIA/TensorRT.git -b release/8.6 --single-branch
cd TensorRT
2. 启动 TensorRT NGC 容器。
docker run -it --gpus all --name=tensorrt -e HF_ENDPOINT="https://hf-mirror.com" -v $PWD:/workspace nvcr.io/nvidia/pytorch:23.06-py3 /bin/bash
此时会从 nvcr 拉取 docker 镜像，请确保网络环境较好，直到镜像里所有层下载成功。成功后会直接进入容器内部。

检查必要的包的版本

1. 确保使用较新的 pip 包和 TensorRT 包，并使用命令更新。
#更新pip包
python3 -m pip install --upgrade pip
2. 检查 TensorRT 版本，最低要求为8.6.0。
#检查TensorRT版本
python3 -c 'import tensorrt;print(tensorrt.__version__)'
本指导下载的 TensorRT 版本为8.6.1，如版本低于8.6.0，请卸载 TensorRT 并重新安装。
#卸载TensorRT包
python3 -m pip uninstall tensorrt

#下载新的TensorRT包，且指定版本，这里指定为8.6.1
python3 -m pip install tensorrt==8.6.1
python3 -c 'import tensorrt;print(tensorrt.__version__)'

下载运行框架和模型必要的包

需要安装的包位于 demo/Diffusion/requirements.txt 里，需要安装里面所有的包。
#安装运行框架和模型所需的包
export TRT_OSSPATH=/workspace
cd $TRT_OSSPATH/demo/Diffusion 
pip3 install -r requirements.txt

#创建输出文件夹(注意一定要先创建文件夹再运行模型，否则后续会报错找不到路径)
mkdir -p onnx engine output
export onnx_dir=onnx
export engine_dir=engine
export output_dir=output
到此环境已经备好，所需的包也已经准备完毕。
#基础包              对应版本
diffusers           0.14.0
onnx                1.13.1
onnx-graphsurgeon   0.3.26
onnxruntime         1.14.1
polygraphy          0.47.1
tensorrt            8.6.1.6
tokenizers          0.13.2
torch               1.13.0
transformers        4.26.1

运行模型

下载模型权重地址换源

由于中国大陆无法下载 Hugging Face 网站模型，首先需要对下载网站换源，使用国内镜像网站的 HF-Mirror 模型。
说明：
如果 docker run 的时候加上了-e HF_ENDPOINT="https://hf-mirror.com"，则此步可以跳过。
#单次有效，退出容器且暂停容器运行后失效，再次重启进入容器需重新输入此条命令
export HF_ENDPOINT="https://hf-mirror.com"

#设为默认，永久有效，即便退出容器且暂停容器运行，再次进入容器后也可直接运行模型（推荐使用此方法）
echo 'export HF_ENDPOINT="https://hf-mirror.com"' >> ~/.bashrc

Stable Diffusion 模型在 Hugging Face 下载需要 User access tokens，登录 Hugging Face 网站获取 User access tokens，可参见 Hugging Face User access tokens。

1. 单击个人账户 >Settings > Access Tokens > New token：

2. Name 输入创建的 token 名，Type 选 Fine-grained 即可，将生成的 token 输入到系统环境中。
#输入token，example_token为在网站上生成的token
export HF_TOKEN=example_token
Demo Stable Diffusion 有三种用法：
根据给定的 prompt 生成图像（demo_txt2img.py）。
根据给定的 prompt 对已有图像进行风格迁移（demo_img2img.py）。
根据给定的 prompt 对已有图像的mask分割区域做修改（demo_inpaint.py）。
不同用法的具体可选参数可以使用以下命令查看：
#查看
python3 demo_txt2img.py --help
python3 demo_img2img.py --help
python3 demo_inpaint.py --help

根据给定的 prompt 生成图像（demo_txt2img.py）

1. 首先，查看运行文件的可选参数：
python3 demo_txt2img.py --help

#可选参数如下：
prompt                Text prompt(s) to guide image generation
options:
-h, --help            show this help message and exit
--version {1.4,1.5,2.0,2.0-base,2.1,2.1-base}    Version of Stable Diffusion
--negative-prompt [NEGATIVE_PROMPT ...]
--repeat-prompt {1,2,4,8,16}    Number of times to repeat the prompt (batch size multiplier)
--height HEIGHT       Height of image to generate (must be multiple of 8)
--width WIDTH         Height of image to generate (must be multiple of 8)
--denoising-steps DENOISING_STEPS
                      Number of denoising steps
--onnx-opset {7,8,9,10,11,12,13,14,15,16,17}
                      Select ONNX opset version to target for exported models
--onnx-dir ONNX_DIR   Output directory for ONNX export
--onnx-refit-dir ONNX_REFIT_DIR
                      ONNX models to load the weights from
--force-onnx-export   Force ONNX export of CLIP, UNET, and VAE models
--force-onnx-optimize
                      Force ONNX optimizations for CLIP, UNET, and VAE models
--engine-dir ENGINE_DIR
                      Output directory for TensorRT engines
--force-engine-build  Force rebuilding the TensorRT engine
--build-static-batch  Build TensorRT engines with fixed batch size.
--build-dynamic-shape
                      Build TensorRT engines with dynamic image shapes.
--build-enable-refit  Enable Refit option in TensorRT engines during build.
--build-all-tactics   Build TensorRT engines using all tactic sources.
--timing-cache TIMING_CACHE
                      Path to the precached timing measurements to accelerate build.
--num-warmup-runs NUM_WARMUP_RUNS
                      Number of warmup runs before benchmarking performance
--nvtx-profile        Enable NVTX markers for performance profiling
--seed SEED           Seed for random generator to get consistent results
--use-cuda-graph      Enable cuda graph
--output-dir OUTPUT_DIR
                      Output directory for logs and image artifacts
--hf-token HF_TOKEN   HuggingFace API access token for downloading model checkpoints
-v, --verbose         Show verbose output
--scheduler {DDIM,EulerA,LMSD,DPM,PNDM}
                      Scheduler for diffusion process
2. 运行模型：
#根据prompt生成图像，使用的Stable Diffusion版本为2.0，生成图像的分辨率为512*512
python3 demo_txt2img.py "a beautiful photograph of Mt. Fuji during cherry blossom" --hf-token=$HF_TOKEN --verbose --version=2.0 --onnx-dir=$onnx_dir --engine-dir=$engine_dir --output-dir=$output_dir --seed=456123 --repeat-prompt=1 --width=512 --height=512 --num-warmup-runs=10
3. 输入命令后会开始下载模型权重，请确保较好的网络连接。
运行结束后可看到如下结果（参见）：
|------------|--------------|
|   Module   |   Latency    |
|------------|--------------|
|    CLIP    |      4.29 ms |
| UNet x 50  |    803.43 ms |
|  VAE-Dec   |     40.22 ms |
|------------|--------------|
|  Pipeline  |    846.50 ms |
|------------|--------------|
Saving image 1 / 1 to: output/txt2img-fp16-a_beautifu-1-3070.png
可以在 demo/Diffusion/output/ 文件夹下看到生成的图像：


根据给定的 prompt 对已有图像进行风格迁移（demo_img2img.py）

直接运行模型

1. 直接运行的话会对网上下载的示例图像做风格迁移：
python3 demo_img2img.py "photorealistic new zealand hills" --hf-token=$HF_TOKEN --verbose --version=2.0 --onnx-dir=$onnx_dir --engine-dir=$engine_dir --output-dir=$output_dir --seed=456123 --repeat-prompt=1 --width=512 --height=512 --num-warmup-runs=10
示例图像：

2. 运行结束后可看到如下结果（参见）：
|------------|--------------|
|   Module   |   Latency    |
|------------|--------------|
|  VAE-Enc   |     12.92 ms |
|    CLIP    |      4.56 ms |
| UNet x 50  |    629.90 ms |
|  VAE-Dec   |     26.67 ms |
|------------|--------------|
|  Pipeline  |    674.55 ms |
|------------|--------------|
Saving image 1 / 1 to: output/img2img-fp16-photoreali-1-3321.png
根据输入的 prompt 做风格迁移后的图像：


选择自己的图像进行风格迁移

除了对示例图像做风格迁移，还可以对自己已有的任何图像做风格迁移，需要添加的参数为 --input-image。
例如在 output 文件夹下已有图像：

路径为 demo/Diffusion/output/txt2img-fp16-a_beautifu-1-2703.png。
#将图像变为抽象风格(注意输出路径的写法)
python3 demo_img2img.py "abstractive the Bund of Shanghai" --hf-token=$HF_TOKEN --verbose --version=2.0 --onnx-dir=$onnx_dir --engine-dir=$engine_dir --output-dir=$output_dir --seed=456123 --repeat-prompt=1 --width=512 --height=512 --num-warmup-runs=10 --input-image="$output_dir/txt2img-fp16-a_beautifu-1-2703.png"
运行结束后可看到如下结果（参见）：
|------------|--------------|
|   Module   |   Latency    |
|------------|--------------|
|  VAE-Enc   |     13.53 ms |
|    CLIP    |      4.48 ms |
| UNet x 50  |    621.11 ms |
|  VAE-Dec   |     26.97 ms |
|------------|--------------|
|  Pipeline  |    666.61 ms |
|------------|--------------|
Saving image 1 / 1 to: output/img2img-fp16-abstractiv-1-8447.png
根据输入的 prompt 做风格迁移后的图像：


根据给定的 prompt 对已有图像的 mask 分割区域做修改（demo_inpaint.py）

直接运行模型：

直接运行的话会对网上下载的示例图像和 mask 分割图像区域做修改：
python3 demo_inpaint.py "a mecha robot sitting on a bench" --hf-token=$HF_TOKEN --verbose --version=2.0 --onnx-dir=$onnx_dir --engine-dir=$engine_dir --output-dir=$output_dir --seed=456123 --repeat-prompt=1 --width=512 --height=512 --num-warmup-runs=10
示例图像：

mask 分割图像：

运行结束后可看到如下结果（参见）：
|------------|--------------|
|   Module   |   Latency    |
|------------|--------------|
|  VAE-Enc   |     12.60 ms |
|    CLIP    |      4.41 ms |
| UNet x 50  |    634.08 ms |
|  VAE-Dec   |     26.25 ms |
|------------|--------------|
|  Pipeline  |    680.40 ms |
|------------|--------------|
Saving image 1 / 1 to: output/inpaint-fp16-a_mecha_ro-1-9059.png
修改后的图像：


选择自己的图像进行修改：

选择自己的图像进行修改时，需要加上 --input-image=<path to image> 和 --mask-image=<path to mask> 两个参数，请确保示例图像和 mask 分割 label 图像具有同样的图像分辨率（高度和宽度相同）。如果两张图像分辨率不同，则会运行网上下载的图像进行模型推理。

注意事项

在切换不同版本的 Stable Diffusion 模型推理时，onnx 文件夹和 engine 文件夹需要新建，来存储不同版本模型的中间状态。
#使用1.5版本的Stable Diffusion时
mkdir -p onnx-1.5 engine-1.5
export onnx_dir=onnx-1.5
export engine_dir=engine-1.5

python3 demo_inpaint.py "a mecha robot sitting on a bench" --hf-token=$HF_TOKEN --verbose --version=1.5 --onnx-dir=$onnx_dir --engine-dir=$engine_dir --output-dir=$output_dir --seed=456123 --repeat-prompt=1 --width=512 --height=512 --num-warmup-runs=10
说明：
由于 OpenCloudOS 是 TencentOS Server 的开源版本，理论上上述文档当中的所有操作同样适用于 OpenCloudOS。

参考文档

TensorRT Diffusion Demo
Hugging Face 镜像网站
Hugging Face User access tokens

中国站
目录