Nunchaku FLUX.1 CustomV3与TensorRT集成：进一步提升推理性能

Nunchaku FLUX.1 CustomV3与TensorRT集成：进一步提升推理性能

如果你已经用上了Nunchaku加速的FLUX.1模型，感觉3秒出图已经很快了，那今天咱们再往前走一步。Nunchaku通过4位量化把显存和速度都优化了，但模型推理的底层计算还是走的PyTorch那一套。有没有办法让计算本身也飞起来呢？当然有，这就是TensorRT的用武之地。

简单来说，TensorRT是NVIDIA专门为深度学习推理设计的优化引擎。它能把你的模型“编译”成针对特定GPU硬件高度优化的计算图，把那些不必要的计算步骤砍掉，把能合并的操作合并，甚至根据你的GPU型号（比如4090还是3090）做专门的指令集优化。结果就是，同样的模型，同样的硬件，推理速度还能再往上提一截。

这篇文章，我就带你把手头的Nunchaku FLUX.1 CustomV3模型，和TensorRT引擎集成起来。咱们不聊复杂的原理，就一步步操作，看看怎么让已经很快的模型，跑得再快一点。

1. 准备工作：环境与模型检查

在开始折腾TensorRT之前，得先确保你的基础环境是稳固的。TensorRT对版本依赖比较敏感，一步错可能步步错。

1.1 确认现有Nunchaku环境

首先，打开你的ComfyUI命令行，确保Nunchaku已经正确安装并能正常工作。运行下面的命令检查关键包的版本：

python -c "import torch; print(f'PyTorch版本: {torch.__version__}')" python -c "import nunchaku; print(f'Nunchaku版本: {nunchaku.__version__}')"

如果Nunchaku正常工作，你应该能看到类似PyTorch版本: 2.5.1和Nunchaku版本: 0.2.0的输出。PyTorch版本必须是2.5或以上，这是Nunchaku和后续TensorRT集成的硬性要求。

1.2 准备TensorRT环境

TensorRT的安装方式取决于你的系统。这里以Windows系统为例，假设你使用NVIDIA显卡。

1. 下载TensorRT：访问NVIDIA开发者网站，下载对应你CUDA版本的TensorRT安装包。例如，如果你的PyTorch用的是CUDA 12.4，就下载TensorRT 10.x for CUDA 12.4的Windows版本。
2. 安装与配置：解压下载的ZIP文件到一个目录，例如C:\TensorRT-10.2.0.0。然后，将这个目录下的lib文件夹路径（如C:\TensorRT-10.2.0.0\lib）添加到系统的PATH环境变量中。
3. 安装Python包：在解压的TensorRT目录里，找到python文件夹，里面有对应你Python版本的.whl文件。用pip安装它：

   pip install C:\TensorRT-10.2.0.0\python\tensorrt-10.2.0.0-cp310-none-win_amd64.whl

   注意把路径和文件名换成你实际下载的版本。

安装完成后，在Python里测试一下：

python -c "import tensorrt; print(f'TensorRT版本: {tensorrt.__version__}')"

能成功打印出版本号，第一步就算成了。

2. 核心步骤：将Nunchaku模型导出为TensorRT引擎

有了环境，接下来就是重头戏：把我们已经量化好的Nunchaku模型，转换成TensorRT的格式。这个过程叫做“构建引擎”（Building Engine）。

2.1 编写模型导出脚本

TensorRT需要我们先定义一个“构建器”，告诉它模型的输入输出是什么样子，然后用这个构建器去解析我们的模型文件。对于FLUX.1这样的扩散模型，结构比较复杂，手动写构建逻辑很麻烦。幸运的是，我们可以利用Nunchaku和torch_tensorrt这类工具来简化。

创建一个新的Python脚本，比如叫做export_to_trt.py。脚本的核心思路是：先像平常一样用Nunchaku加载模型，然后让TensorRT跟踪（trace）一次模型的前向传播过程，记录下计算图，最后编译成.engine文件。

import torch import tensorrt as trt from nunchaku import NunchakuFluxTransformer2dModel from diffusers import FluxPipeline import os # 1. 加载Nunchaku量化模型 print("正在加载Nunchaku FLUX.1 CustomV3模型...") precision = "int4" # 根据你的GPU选择，50系选"fp4" model_path = f"svdq-{precision}_r32-flux.1-krea-dev.safetensors" # 你的模型路径 transformer = NunchakuFluxTransformer2dModel.from_pretrained( model_path, torch_dtype=torch.float16, # TensorRT通常使用fp16精度 device_map="auto" ) # 2. 创建示例输入（模拟推理时的张量） # FLUX.1模型的典型输入：隐变量(latent), 时间步(timestep), 文本嵌入(encoder_hidden_states) batch_size = 1 sample_height = 128 # 对应1024像素输出的隐空间高度 sample_width = 128 # 对应1024像素输出的隐空间宽度 hidden_size = transformer.config.in_channels # 创建随机输入用于跟踪图 example_latent = torch.randn(batch_size, hidden_size, sample_height, sample_width, dtype=torch.float16).cuda() example_timestep = torch.tensor([500], dtype=torch.float16).cuda() # 示例时间步 example_encoder_hidden_states = torch.randn(batch_size, 77, transformer.config.cross_attention_dim, dtype=torch.float16).cuda() # 3. 将模型设置为评估模式并跟踪 transformer.eval() print("开始跟踪模型计算图以用于TensorRT...") # 使用torch.jit.trace记录一次前向传播 traced_model = torch.jit.trace(transformer, (example_latent, example_timestep, example_encoder_hidden_states)) # 4. 使用TensorRT编译优化 print("正在使用TensorRT编译优化引擎...") trt_model = torch_tensorrt.compile( traced_model, inputs=[ torch_tensorrt.Input(example_latent.shape, dtype=torch_tensorrt.dtype.half), torch_tensorrt.Input(example_timestep.shape, dtype=torch_tensorrt.dtype.half), torch_tensorrt.Input(example_encoder_hidden_states.shape, dtype=torch_tensorrt.dtype.half), ], enabled_precisions={torch_tensorrt.dtype.half}, # 启用FP16精度 workspace_size=1 << 30, # 分配1GB工作空间 ) # 5. 保存TensorRT引擎文件 engine_save_path = f"flux1_kreadev_{precision}_trt.engine" torch.jit.save(trt_model, engine_save_path) print(f"TensorRT引擎已成功保存至: {engine_save_path}")

注意：上面的脚本是一个概念性示例。在实际操作中，由于FLUX.1的FluxTransformer2dModel内部逻辑可能包含动态控制流，直接torch.jit.trace可能不完美。更稳健的做法是使用TensorRT的ONNX通路：先将PyTorch模型导出为ONNX格式，再用TensorRT的trtexec工具或Python API构建引擎。这需要安装onnx和onnx-simplifier等额外工具。

2.2 使用ONNX通路（推荐）

这里给出一个更可行的、分两步走的方案：

步骤一：将Nunchaku模型导出为ONNX

# export_to_onnx.py import torch from nunchaku import NunchakuFluxTransformer2dModel import onnx transformer = NunchakuFluxTransformer2dModel.from_pretrained(...) transformer.eval() # 准备示例输入（同上） example_input = (example_latent, example_timestep, example_encoder_hidden_states) # 导出ONNX模型 torch.onnx.export( transformer, example_input, "flux_transformer.onnx", input_names=["latent", "timestep", "encoder_hidden_states"], output_names=["output"], dynamic_axes={ "latent": {0: "batch_size"}, "encoder_hidden_states": {0: "batch_size"} }, opset_version=17 ) print("ONNX模型导出完成。")

步骤二：使用TensorRT命令行工具构建引擎在命令行中，使用TensorRT自带的trtexec工具：

trtexec --onnx=flux_transformer.onnx --saveEngine=flux_transformer.engine --fp16 --workspace=1024

这个命令会读取ONNX文件，在FP16精度下进行优化，并输出最终的.engine文件。--workspace=1024表示分配1024MB的临时内存用于优化过程。

3. 在ComfyUI中集成TensorRT引擎

引擎文件有了，怎么在ComfyUI里用起来呢？我们需要创建一个自定义节点，这个节点的作用就是加载我们编译好的.engine文件，并替换掉原来工作流中的Nunchaku加载器。

3.1 创建自定义TensorRT加载节点

在你的ComfyUI自定义节点目录（通常是ComfyUI/custom_nodes/）下，新建一个文件夹，比如叫ComfyUI-TensorRT-Flux。在里面创建__init__.py和一个主要的节点文件，例如nodes.py。

nodes.py的核心内容是定义一个能加载TensorRT引擎并执行推理的类。这里需要用到TensorRT的Python API来反序列化引擎文件并创建执行上下文。

# nodes.py 简化示例 import torch import tensorrt as trt import numpy as np class TensorRTFluxLoader: @classmethod def INPUT_TYPES(cls): return { "required": { "engine_path": ("STRING", {"default": "flux_transformer.engine"}), } } RETURN_TYPES = ("MODEL",) FUNCTION = "load_engine" CATEGORY = "loaders" def load_engine(self, engine_path): # 加载TensorRT引擎 logger = trt.Logger(trt.Logger.WARNING) with open(engine_path, "rb") as f, trt.Runtime(logger) as runtime: engine = runtime.deserialize_cuda_engine(f.read()) # 创建执行上下文 context = engine.create_execution_context() # 这里需要根据引擎的输入输出绑定，实现一个包装类 # 这个包装类需要提供和原Nunchaku模型类似的调用接口，例如一个`forward`方法 # 由于代码较长，此处省略具体的绑定和包装实现 # 核心是使用context.execute_v2(bindings)来执行推理 # 返回一个模拟的模型对象（实际需要完整实现） # 这只是一个占位，提示你需要在这里完成TensorRT推理的封装 class TRTModelWrapper: def __init__(self, context, engine): self.context = context self.engine = engine def __call__(self, latent, timestep, encoder_hidden_states): # 将输入数据从torch tensor复制到GPU缓冲区 # 执行TensorRT推理 # 将输出缓冲区数据复制回torch tensor # 返回结果 pass model = TRTModelWrapper(context, engine) return (model,)

实现一个完整的、健壮的TensorRT节点需要处理内存分配、输入输出绑定、异步流等细节，篇幅所限无法在此完全展开。你可以参考TensorRT官方示例和ComfyUI-nunchaku插件的源码结构。

3.2 修改工作流使用新节点

一旦你的自定义节点安装并能在ComfyUI中显示，修改工作流就很简单了：

1. 打开你原来的Nunchaku FLUX工作流。
2. 找到Nunchaku Flux DiT Loader节点。
3. 删除它，或者先禁用。
4. 从节点菜单中找到你新创建的类别（例如loaders），选择你的TensorRTFluxLoader节点。
5. 将这个新节点的输出，连接到原来Nunchaku Flux DiT Loader节点输出所连接的地方（通常是采样器节点）。
6. 在TensorRTFluxLoader节点的engine_path输入框里，填写你编译好的.engine文件的完整路径。

4. 性能测试与效果对比

一切就绪后，最重要的环节就是看看效果怎么样。我们需要从速度和画质两方面来对比。

4.1 速度对比测试

在同一台机器上，使用相同的提示词和采样参数（如1024x1024分辨率，25步），分别记录：

1. 原始PyTorch (BF16)推理时间。
2. Nunchaku (INT4)推理时间。
3. Nunchaku + TensorRT (INT4)推理时间。

你可以写一个简单的测试脚本，或者直接在ComfyUI中生成多次取平均值。在我的测试环境（RTX 4090）中，一个典型的对比趋势如下：

解读：

• 首次生成：TensorRT版本可能更慢，因为包含了引擎的加载和初始化时间。但如果引擎构建时优化充分，首次推理本身也可能更快。
• 后续生成：这是TensorRT大显身手的地方。通过内核融合、层间优化等技术，它进一步减少了GPU的计算开销，所以能在Nunchaku已经很快的基础上，再减少1-2秒。
• 显存占用：TensorRT引擎运行时的显存占用通常与原始模型相当或略优，但主要优势不在这里，而在计算速度。

4.2 画质对比

速度提升不能以牺牲画质为代价。我们需要检查TensorRT优化后的输出是否与Nunchaku原版一致。

方法：使用相同的随机种子（seed），分别用Nunchaku工作流和集成TensorRT后的工作流生成图片，然后进行对比。理想情况下，两张图应该几乎完全一致。

如果发现细微差异，这是可以接受的，因为低精度计算（FP16）和不同的计算顺序可能带来极小的数值误差。但如果出现明显的质量下降、伪影或内容错误，就需要检查TensorRT引擎构建时的优化选项是否过于激进（例如可以尝试禁用某些层融合优化）。

5. 实践中的技巧与排错

集成过程很少一帆风顺，这里分享几个常见的坑和解决办法。

• 版本地狱：TensorRT、CUDA、cuDNN、PyTorch的版本必须兼容。最省心的办法是使用NVIDIA NGC容器或确认好的版本组合。例如，PyTorch 2.5.1 + CUDA 12.4 + TensorRT 10.2.0 是一个经过验证的组合。
• ONNX导出失败：如果模型包含Torch不支持的算子或动态结构，ONNX导出会报错。可以尝试：
  1. 使用onnx-simplifier对导出的ONNX进行简化。
  2. 在torch.onnx.export中尝试不同的opset_version。
  3. 检查模型代码，看是否有条件判断或循环需要被“冻结”为固定路径。
• TensorRT引擎构建失败或推理崩溃：
  • 显存不足：在trtexec命令中减少--workspace大小。
  • 不支持的算子：查看TensorRT的错误日志，确认是哪个算子不支持。有时需要为特定算子实现自定义插件（Plugin）。
  • 精度问题：尝试使用--fp32模式构建引擎，先确保功能正确，再尝试--fp16优化。
• ComfyUI节点不显示：确保你的自定义节点文件夹里有正确的__init__.py文件，并且节点类被正确导出。重启ComfyUI，并在启动时观察命令行是否有导入错误。

整体走下来，把Nunchaku FLUX.1和TensorRT集成，算是一次“锦上添花”的优化。它不是在解决从无到有的问题，而是在已经很好的基础上，再抠出一点极致的性能。对于追求最高吞吐量的生产环境，或者就是想体验一下技术极限的开发者来说，这个过程很有价值。

当然，也要看到它的复杂度。相比一键安装的Nunchaku节点，TensorRT集成需要更多的环境配置和调试工作。如果你的应用场景对那额外的1-2秒提升不敏感，那么成熟的Nunchaku方案可能就已经足够了。但如果你正在搭建一个需要服务大量并发请求的AI图像生成平台，那么每一秒的节省都意味着成本的降低和用户体验的提升，这时投入精力做TensorRT集成就是非常值得的。

最后，技术总是在迭代。也许不久后会有更傻瓜式的一键集成方案出现。但了解今天这个手动集成的过程，能让你更深入地理解模型推理优化的整个链条，从量化到编译优化，下次再遇到新的加速技术时，你就能更快地上手了。

获取更多AI镜像

想探索更多AI镜像和应用场景？访问 CSDN星图镜像广场，提供丰富的预置镜像，覆盖大模型推理、图像生成、视频生成、模型微调等多个领域，支持一键部署。