File size: 3,849 Bytes
76de374 a542584 76de374 a542584 76de374 a542584 76de374 a542584 76de374 a542584 76de374 a542584 76de374 a542584 76de374 a542584 76de374 a542584 76de374 a542584 0988d89 a542584 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 |
# 混元-DiT TensorRT 加速
语言: [**English**](https://huggingface.co/Tencent-Hunyuan/TensorRT-libs/blob/main/README.md) | **中文**
我们提供了将 [混元-DiT](https://github.com/Tencent/HunyuanDiT) 中的文生图模型转换为 TensorRT 的代码和相关依赖用于推理加速
(比 Flash Attention 更快). 您可以使用以下步骤使用我们 TensorRT 模型, 基于 **TensorRT-9.2.0.5** 和 **cuda (11.7 或 11.8)**.
> ⚠️ **重要提醒 (关于TensorRT加速版本测试的建议)**:
> 我们建议用户在 Compute Capability>=8.0 的 NVIDIA 显卡上测试 TensorRT 版本, (例如, RTX4090, RTX3090, H800, A10/A100/A800 等)
> 您可以从[这里](https://developer.nvidia.com/cuda-gpus#compute)查询到您显卡对应的 Compute Capability。
> 对于 Compute Capability<8.0 的 NVIDIA 显卡，如果您要尝试 TensorRT 版本，有可能遭遇 TensorRT Engine 文件无法生成的错误或推理性能较差
> 的问题，主要原因在于 TensorRT 没有在该架构上支持 fused mha kernel。
## 🛠 构建步骤
### 1. 从 Huggingface 下载 TensorRT 的依赖文件
```shell
cd HunyuanDiT
# Download the dependencies
huggingface-cli download Tencent-Hunyuan/TensorRT-libs --local-dir ./ckpts/t2i/model_trt
```
### 2. 安装 TensorRT 依赖
```shell
sh trt/install.sh
```
### 3. 构建 TensorRT engine
#### 方法1: 使用预构建的 engine
本仓库提供了一些预构建的 TensorRT engine.
| 支持的 GPU | 文件链接 | 远程地址 |
|:----------------:|:---------------------------------------------------------------------------------------------------------------:|:---------------------------------:|
| GeForce RTX 3090 | [HuggingFace](https://huggingface.co/Tencent-Hunyuan/TensorRT-engine/blob/main/engines/RTX3090/model_onnx.plan) | `engines/RTX3090/model_onnx.plan` |
| GeForce RTX 4090 | [HuggingFace](https://huggingface.co/Tencent-Hunyuan/TensorRT-engine/blob/main/engines/RTX4090/model_onnx.plan) | `engines/RTX4090/model_onnx.plan` |
| A100 | [HuggingFace](https://huggingface.co/Tencent-Hunyuan/TensorRT-engine/blob/main/engines/A100/model_onnx.plan) | `engines/A100/model_onnx.plan` |
可以使用以下命令下载并放置在指定的位置
```shell
huggingface-cli download Tencent-Hunyuan/TensorRT-engine <远程地址> --local-dir ./ckpts/t2i/model_trt/engine
```
#### 方法2: 自行构建 engine
如果您使用不同于上面表格中的 GPU, 可以使用以下命令构建适配于当前 GPU 的 engine.
```shell
# 首先设置 TensorRT 构建相关的环境变量，我们提供了一个脚本来一键设置
source trt/activate.sh
# 构建 TensorRT engine. 默认会读取当前目录下的 ckpts 文件夹
sh trt/build_engine.sh
```
### 4. 使用 TensorRT 模型进行推理.
```shell
# 使用 prompt 强化 + 文生图 TensorRT 模型进行推理
python sample_t2i.py --prompt "渔舟唱晚" --infer-mode trt
# 关闭 prompt 强化 (可以在显存不足时使用)
python sample_t2i.py --prompt "渔舟唱晚" --infer-mode trt --no-enhance
```
### 5. 提示
为了性能考虑, 我们限制 TensorRT engine 只支持以下输入的形状. 未来我们会验证并尝试支持任意的形状.
```python
STANDARD_SHAPE = [
[(768, 768), (1024, 1024), (1280, 1280)], # 1:1
[(1024, 768), (1152, 864), (1280, 960)], # 4:3
[(768, 1024), (864, 1152), (960, 1280)], # 3:4
[(1280, 768)], # 16:9
[(768, 1280)], # 9:16
]
```
## ❓ Q&A
参考 [Q&A](./QA.md) 获取更多关于构建 TensorRT Engine 的问题解答.
|