解决stable-diffusion-webui-depthmap中TensorRT不支持INT64的问题

在使用stable-diffusion-webui-depthmap项目进行背景移除时，用户可能会遇到一个与TensorRT和ONNX运行时相关的技术问题。本文将详细分析问题原因并提供完整的解决方案。

问题现象

当用户启用"Remove Background"选项时，系统会输出以下警告信息：

TensorRT does not natively support INT64. Attempting to cast down to INT32.

这表明TensorRT运行时遇到了64位整型数据，但无法原生支持，只能尝试将其转换为32位整型。

问题根源分析

1. 模型权重类型不匹配：ONNX模型使用了INT64类型的权重，而TensorRT执行提供程序仅原生支持INT32类型
2. 显式批处理标志冲突：警告信息中还提到了NetworkDefinitionCreationFlag::kEXPLICIT_BATCH标志的使用问题
3. 环境配置不当：可能是由于ONNX运行时和TensorRT的版本不兼容或安装不正确导致

完整解决方案

1. 清理现有环境：

   • 删除stable-diffusion-webui\venv\Lib\site-packages下的onnxruntime相关文件夹
   • 移除stable-diffusion-webui\models\rem_bg目录

2. 重新安装依赖：

   • 进入虚拟环境(venv)
   • 安装特定版本的ONNX运行时GPU支持：

     pip install onnxruntime-gpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/
     

   • 安装带有GPU支持的rembg包：

     pip install rembg[gpu,cli]
     

技术原理

1. TensorRT的限制：TensorRT设计时主要针对32位浮点和整型运算优化，对64位整型支持有限
2. 类型转换机制：ONNX运行时会自动尝试将INT64转换为INT32，但可能影响精度
3. 批处理模式：显式批处理标志(kEXPLICIT_BATCH)改变了TensorRT处理输入的方式

注意事项

1. 确保CUDA版本(12.1)与安装的onnxruntime-gpu版本兼容
2. 清理环境时要彻底，避免残留文件影响新安装
3. 使用正确的pip源安装特定版本的包
4. 安装完成后建议重启stable-diffusion-webui服务

通过以上步骤，用户应该能够成功解决TensorRT不支持INT64的问题，并使背景移除功能恢复正常工作。该解决方案不仅解决了类型不匹配问题，还优化了GPU加速性能，使背景移除操作更加高效。