修复Stable-Diffusion 的 TensorRT 不显示和不工作的问题

307 阅读3分钟

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  

  1. 初步排查  

   在查找解决方案的过程中发现:

   - 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),无法完成初始化。

  1. 解决方案  

   通过安装以下开发者工具并配置环境变量,解决了上述问题:

   - CUDA Toolkit  

   - cuDNN  

   - TensorRT {这个下载下来是压缩包 记得自己手动添加PATH环境变量}  

  1. 新问题  

   在安装上述套件后,Stable-Diffusion-WebUI-TensorRT 短暂运行成功,并生成了量化后的模型。然而,在某次运行后,发现 TensorRT 对新版本突然不再兼容。  

   - 排查后定位到 utilities.py 文件的部分函数报错。  

  1. 兼容性修复 | 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.py file was updated to resolve compatibility issues.  

Code Changes  

  1. Function allocate_buffers: Fixed input/output tensor allocation logic using the new TensorRT API.  

  2. Function __str__: Updated tensor shape printing logic to use get_tensor_profile_shape.  

Let me know if further details or testing is required. Thank you!