Issue Title: 关于 utilities.py 的函数修改和兼容性修复 | Modifications and Compatibility Fixes in utilities.py
背景 | Background
我于2025年1月26日从秋叶整合包分支下载了 stablediffusion,并通过 git clone 下载了 Stable-Diffusion-WebUI-TensorRT。
初次运行时,UI 未显示 TensorRT 标签。
问题与解决过程 | Problem and Resolution Process
- 初步排查
在查找解决方案的过程中发现:
- Bilibili 用户 Poro7 提供了以下建议:
```bash
python -m pip install --pre --extra-index-url pypi.nvidia.com/ tensorrt==9.0.1.post11.dev4 --no-cache-dir
```
但该安装方法仅包含 .pyc 或 .py 文件,不附带 DLL。
- 这导致 TensorRT 缺失关键 DLL(来源:developer.nvidia.com/tensorrt),无法完成初始化。
- 解决方案
通过安装以下开发者工具并配置环境变量,解决了上述问题:
- cuDNN
- TensorRT {这个下载下来是压缩包 记得自己手动添加PATH环境变量}
- 新问题
在安装上述套件后,Stable-Diffusion-WebUI-TensorRT 短暂运行成功,并生成了量化后的模型。然而,在某次运行后,发现 TensorRT 对新版本突然不再兼容。
- 排查后定位到 utilities.py 文件的部分函数报错。
- 兼容性修复 | Compatibility Fixes
为解决新版与旧版兼容性问题,对 utilities.py 进行了以下修改,详细代码如下。
修改内容 | Code Changes
修改函数 allocate_buffers
文件行数: 20-36(示例)
主要修复: 使用 TensorRT 的新 API 修复输入/输出张量分配逻辑。
def allocate_buffers(self, shape_dict=None, device="cuda"):
nvtx.range_push("allocate_buffers")
for idx in range(self.engine.num_io_tensors):
tensor_name = self.engine.get_tensor_name(idx) # 获取张量名称
dtype = trt.nptype(self.engine.get_tensor_dtype(tensor_name)) # 新API获取数据类型
# 获取形状
if shape_dict and tensor_name in shape_dict:
shape = shape_dict[tensor_name].shape
else:
shape = self.context.get_tensor_shape(tensor_name) # 新API获取形状
# 设置输入形状
if self.engine.get_tensor_mode(tensor_name) == trt.TensorIOMode.INPUT:
self.context.set_input_shape(tensor_name, shape) # 新API设置输入形状
# 创建Tensor
tensor = torch.empty(tuple(shape), dtype=numpy_to_torch_dtype_dict[dtype]).to(device=device)
self.tensors[tensor_name] = tensor
nvtx.range_pop()
Patch 注释:
-
[行 23] 修复
get_tensor_dtype的调用以支持新 API。 -
[行 26] 替换形状获取逻辑,调用新 API
get_tensor_shape。 -
[行 30] 调用新 API
set_input_shape修复输入形状设置问题。
修改函数 __str__
文件行数: 38-46(示例)
主要修复: 修复形状打印函数以支持新版 TensorRT 的 API。
def __str__(self):
out = ""
for opt_profile in range(self.engine.num_optimization_profiles):
for binding_idx in range(self.engine.num_io_tensors):
name = self.engine.get_tensor_name(binding_idx)
# 替换为 get_tensor_profile_shape,并调整参数顺序
shape = self.engine.get_tensor_profile_shape(name, opt_profile) # 关键修改
out += f"\t{name} = {shape}\n"
return out
Patch 注释:
-
[行 43] 将
get_binding_shape替换为get_tensor_profile_shape,并调整参数顺序以兼容新版本。
结论 | Conclusion
通过上述修改,成功修复了新版 TensorRT 的兼容性问题。期待官方对此问题的进一步优化或其他解决方案。
如果需要进一步测试或补充信息,请随时联系!
---
English Summary
Background
On January 26, 2025, I cloned Stable-Diffusion-WebUI-TensorRT from the GitHub repository while using the Autumn Integration package of stablediffusion. However, the UI did not display the TensorRT tab upon the first launch.
Problem & Resolution
-
The issue stemmed from missing DLLs in the TensorRT package.
-
After installing CUDA Toolkit, cuDNN, and TensorRT, the tool worked briefly but became incompatible with new versions of TensorRT.
-
The
utilities.pyfile was updated to resolve compatibility issues.
Code Changes
-
Function
allocate_buffers: Fixed input/output tensor allocation logic using the new TensorRT API. -
Function
__str__: Updated tensor shape printing logic to useget_tensor_profile_shape.
Let me know if further details or testing is required. Thank you!