# 混元-DiT TensorRT 加速

[English](https://huggingface.co/Tencent-Hunyuan/TensorRT-libs/blob/main/README.md) | 中文

我们提供了将 [混元-DiT](https://github.com/Tencent/HunyuanDiT) 中的文生图模型转换为 TensorRT 的代码和相关依赖用于推理加速
（比 Flash Attention 更快）。 您可以使用以下步骤使用我们 TensorRT 模型。

## 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

# 方式1: 构建 TensorRT engine. 默认会读取当前目录下的 ckpts 文件夹
sh trt/build_engine.sh

# 方式2: 如果您的模型目录不是 ckpts, 需要指定模型目录
sh trt/build_engine.sh </path/to/ckpts>
```

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

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

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