混元-DiT TensorRT 加速

语言: English | 中文

我们提供了将混元-DiT 中的文生图模型转换为 TensorRT 的代码和相关依赖用于推理加速 (比 Flash Attention 更快). 您可以使用以下步骤使用我们 TensorRT 模型, 基于 TensorRT-10.1.0.27 和 cuda (11.7 或 11.8).

⚠️ 重要提醒 (关于TensorRT加速版本测试的建议):
我们建议用户在 Compute Capability>=8.0 的 NVIDIA 显卡上测试 TensorRT 版本, (例如, RTX4090, RTX3090, H800, A10/A100/A800 等) 您可以从这里查询到您显卡对应的 Compute Capability。对于 Compute Capability<8.0 的 NVIDIA 显卡，如果您要尝试 TensorRT 版本，有可能遭遇 TensorRT Engine 文件无法生成的错误或推理性能较差的问题，主要原因在于 TensorRT 没有在该架构上支持 fused mha kernel。

🛠 构建步骤

1. 从 Huggingface 下载 TensorRT 的依赖文件

cd HunyuanDiT

# Download the dependencies
huggingface-cli download Tencent-Hunyuan/TensorRT-libs --local-dir ./ckpts/t2i/model_trt

2. 安装 TensorRT 依赖

# 使用下面的命令解压并安装 TensorRT 依赖
sh trt/install.sh

# 设置使用或构建 TensorRT 相关的环境变量，我们提供了一个脚本来一键设置
source trt/activate.sh

3. 构建 TensorRT engine

方法1: 使用预构建的 engine

我们提供了一些预构建的 TensorRT Engines, 需要从 Huggingface 下载.

支持的 GPU	远程地址
GeForce RTX 3090	`engines/RTX3090/model_onnx.plan`
GeForce RTX 4090	`engines/RTX4090/model_onnx.plan`
A100	`engines/A100/model_onnx.plan`

可以使用以下命令下载并放置在指定的位置. 注意: 请将 <远程地址> 替换为上表中对应 GPU 的远程地址.

export REMOTE_PATH=<远程地址>
huggingface-cli download Tencent-Hunyuan/TensorRT-engine ${REMOTE_PATH} ./ckpts/t2i/model_trt/engine/
ln -s ${REMOTE_PATH} ./ckpts/t2i/model_trt/engine/model_onnx.plan

方法2: 自行构建 engine

如果您使用不同于上面表格中的 GPU, 可以使用以下命令构建适配于当前 GPU 的 engine.

Hunyuan-DiT v1.2

# Build the TensorRT engine. By default, it will read the `ckpts` folder in the current directory.
sh trt/build_engine.sh

使用旧版本模型, Hunyuan-DiT <= v1.1

# v1.1 
sh trt/build_engine.sh 1.1
# v1.0 
sh trt/build_engine.sh 1.0

如果你看到日志打印出 &&&& PASSED TensorRT.trtexec [TensorRT v10100], 说明引擎文件已经成功生成了。

4. 使用 TensorRT 模型进行推理.

# 重要: 如果您在使用 TensorRT 模型之前没有设置环境变量, 请使用以下命令设置
source trt/activate.sh

# 使用 prompt 强化 + 文生图 TensorRT 模型进行推理
python sample_t2i.py --prompt "渔舟唱晚" --infer-mode trt

# 关闭 prompt 强化 (可以在显存不足时使用)
python sample_t2i.py --prompt "渔舟唱晚" --infer-mode trt --no-enhance

5. 提示

为了性能考虑, 我们限制 TensorRT engine 只支持以下输入的形状. 未来我们会验证并尝试支持任意的形状.

STANDARD_SHAPE = [
    [(1024, 1024), (1280, 1280)],   # 1:1
    [(1280, 960)],    # 4:3
    [(960, 1280)],    # 3:4
    [(1280, 768)],                              # 16:9
    [(768, 1280)],                              # 9:16
]

❓ Q&A

参考 Q&A 获取更多关于构建 TensorRT Engine 的问题解答.