HunyuanDiT-TensorRT-libtorch

用TensorRT和libtorch简单实现了HunyuanDiT模型的pipeline推理。

准备

• 安装TensorRT, TensorRT10的api变化了, 建议用TensorRT8以下的版本
• 从huggingface下载模型
• 安装pytorch, onnx等依赖

导出4个onnx模型用于pipeline

修改export.py中的args 执行:

python export.py

你会得到bert, t5, hunyuan, vae四个onnx模型 你可以用onnxsim将它们简化 执行:

onnxsim bert.onnx bert-sim.onnx
onnxsim t5.onnx t5-sim.onnx
onnxsim hunyuan.onnx hunyuan-sim.onnx
onnxsim vae.onnx vae-sim.onnx

onnx很大的情况下, 简化的耗时也很长

onnx转换到tensorrt

这里我用了trtexec转化, 比较省事

trtexec --onnx=bert-sim.onnx --saveEngine=bert.plan --fp16
trtexec --onnx=t5-sim.onnx --saveEngine=t5.plan --fp16
trtexec --onnx=hunyuan-sim.onnx --saveEngine=hunyuan.plan --fp16
trtexec --onnx=vae-sim.onnx --saveEngine=vae.plan --fp16

tensorrt转换的过程也很慢

编译安装python包

执行:

python setup.py install

包名是: py_hunyuan_dit

推理一个文生图

修改run.py中的4个模型路径, 修改推理步数, 默认100比较慢

执行:

python run.py

你会看到生成的图片