Nodes/ComfyUI-TensorRT-Reforge/TensorRT Exporter Reforge
ComfyUI Node

TensorRT Exporter Reforge

Exporting a TensorRT Engine So Everything Else Is Fast

By zaochuan5854·Created 5 months ago·Updated 2 months ago· 23
TensorRT Exporter Reforge
      model_name
      opt_width512
      opt_height512
      opt_batch_size1
      enable_lorafalse
      opt_context_mult1
      num_video_frames1
      filename_prefixtensorrt/
      min_width0
      max_width0
      min_height0
      max_height0
      min_batch_size0
      max_batch_size0
      min_context_mult0
      max_context_mult0

      This is the node you run once and then pray you never need again. TensorRT Exporter Reforge takes a model you already have - a checkpoint or a diffusion model sitting in ComfyUI - and compiles it into a hardware-tuned TensorRT engine. The first build takes minutes and eats every bit of VRAM it can find. The reward is a .engine (or .bundle) file that makes every generation after it dramatically cheaper.

      Why bother? Because this pack's real target is models that are slow for what they are. Anima is the flagship case: it's only a 2B model but it wants 30–50 steps and famously runs "about 10s on a 5090 and over a minute on a 4060." That's exactly the profile TensorRT was built to fix. And because this is a fork of comfyanonymous's original ComfyUI_TensorRT reworked to survive modern PyTorch, you also get Flux, SDXL, SD 1.5, SD3/SD3.5, AuraFlow, and SVD out of the box.

      How it works

      Export is a two-stage pipeline. First the node loads your model, traces it, and exports an ONNX graph using PyTorch's dynamo=True path (which handles the mathematically fiddly modern architectures the old tracing-only versions choked on). Opset is picked per model - 18 for most, 25 for Anima, whose architecture needs the newer features. Then it feeds that ONNX into the TensorRT builder with a strongly-typed network and an optimization profile, and writes out the serialized engine.

      That profile is the whole game. TensorRT lets you specify min / opt / max shapes, and the engine performs best at the opt shape. The node also keeps a timing cache, so if you re-export the same model, the algorithm search is reused and the rebuild is dramatically faster than the first one.

      The clever part is Refit for LoRA. Traditional TensorRT meant a full 3–10 minute rebuild every time you swapped a LoRA. Here, if you enable the LoRA option, the builder sets the REFIT flag and the exporter hashes every ONNX initializer against your model's state dict to build a weight map - then it packages engine, map, and (for Anima) the LLM adapter's ONNX into a single .bundle file. At load time, LoRA weight deltas get injected straight into the compiled engine. You get LoRA flexibility in seconds instead of a rebuild. That's the feature the README is most proud of, and rightly so.

      The inputs that matter

      The node has a lot of dials. Most of them, you'll leave alone. The ones that matter:

      • model_name - a combo listing everything in models/checkpoints and models/diffusion_models. Pick your base model.
      • opt_width / opt_height / opt_batch_size - the exact dimensions and batch you'll generate at. Set these to what you actually use; TensorRT is fastest at the opt shape. There are min_* / max_* fields if you want a dynamic-shape engine, but every dynamic range you add costs build time and some speed, so don't add range you don't need.
      • enable_lora - the experimental LoRA toggle. If you think you'll ever want LoRAs on this engine, turn it on now; it's what produces the .bundle and the refit magic. Constant folding is disabled during export to protect the weight mapping, which is part of why it's marked experimental.
      • opt_context_mult - the multiplier on context length (how many tokens of conditioning the engine accommodates). Flux and Anima default to a 512-token context; bump this if you feed in very long prompts.
      • num_video_frames - only for SVD; the temporal frame count becomes part of the engine.
      • filename_prefix - defaults to tensorrt/, so output lands in ComfyUI/output/tensorrt/. The filename encodes the whole spec (..._BS1_W1024_H1024_CM1_LoRA.SDXL.engine), which the loader's companion node reads back.

      There are no outputs - this node exists purely for its side effect of writing the engine file.

      Installing it

      Same pack as the loader - this is the exporter half of one install. You need a CUDA 12.x environment (12.8 recommended; CUDA 11 won't work, 13.x isn't supported yet):

      cd ComfyUI/custom_nodes
      git clone https://github.com/zaochuan5854/ComfyUI-TensorRT-Reforge.git
      pip install -r requirements.txt
      

      That pulls NVIDIA's tensorrt-cu12 wheels (10.15.1+) plus onnx and onnxruntime-gpu. Restart ComfyUI afterward.

      Common issues

      Export fails on Windows with Dynamic VRAM. This is the documented big one: ONNX export conflicts with ComfyUI's dynamic VRAM management during FX graph decomposition. Restart ComfyUI with --disable-dynamic-vram. If your ComfyUI version doesn't recognize that flag and refuses to start, the README says you can safely omit it.

      The engine doesn't match your latent. Export is shape-locked by design. If you later try to sample at a size outside the engine's profile, you get a runtime error. Either re-export with the sizes you want (dynamic ranges if you bounce around) or stay disciplined about resolution. This catches everyone once.

      First build is slow and VRAM-hungry. That's normal - it's compiling a whole model. If it OOMs on a smaller card, export at a lower resolution or batch, then rely on the timing cache for subsequent builds. And if you're on CUDA 11, stop before you start: this pack's requirements are cu12-only.

      CategoryTensorRT

      Inputs (16)

      NameTypeDefaultDescription
      model_nameCOMBO0 options:
      opt_widthINT512
      opt_heightINT512
      opt_batch_sizeINT1
      enable_loraBOOLEANfalse
      opt_context_multINT1
      num_video_framesINT1
      filename_prefixSTRINGtensorrt/
      min_widthINT0
      max_widthINT0
      min_heightINT0
      max_heightINT0
      min_batch_sizeINT0
      max_batch_sizeINT0
      min_context_multINT0
      max_context_multINT0

      Outputs (0)

      No outputs