ComfyUI-Upscaler-TensorRT-RTX

ComfyUI Upscaler TensorRT-RTX

Overview

This project is a fork of the original ComfyUI-Upscaler-Tensorrt by yuvraj108c. The idea behind this fork is to provide users with control over previously constant or hidden parameters during the engine building process, such as adjusting the minimum, maximum input sizes, or modifying the post-upscale resize settings.

Features

• NVIDIA TensorRT Acceleration: Leverages Tensor Cores for up to 30x faster upscaling compared to standard local methods.
• Dynamic Shape Support: Configure minimum, optimal, and maximum resolutions to handle varying input sizes without rebuilding the engine for every image.
• Automatic Engine Building: Automatically compiles specialized .trt files from your .onnx models if a matching engine isn't found.
• Persistent Timing Cache: Saves hardware-specific optimization data to a local cache to speed up subsequent engine builds.
• Memory Efficient: Optimized for VRAM management with automatic garbage collection and cache clearing after inference.

Examples

Installation

ComfyUI Manager

This project has not been submitted to the ComfyUI Manager registry yet. But you can still install it this way:

1. Open ComfyUI Manager.
2. Click the Custom Nodes Manager to open the custom nodes manager page.
3. On the bottom right corner, click the Install via Git URL button.
4. Enter the URL of this repository: https://github.com/ThreadsOfFate/ComfyUI-Upscaler-TensorRT-RTX.git
5. Click "Confirm".
6. Restart ComfyUI.

Manual

1. On the github page, click on the green <> Code button and then "Download ZIP".
2. Extract the root folder within the downloaded ZIP file to your ComfyUI custom_nodes directory.
3. Using the same python environment that runs ComfyUI, install the required dependencies: python -m pip install -r custom_nodes/ComfyUI-Upscaler-TensorRT-Advanced/requirements.txt.
4. Restart ComfyUI.

Requirements

pip install -r requirements.txt

CUDA Toolkit Required

Please download NVIDIA CUDA Toolkit https://developer.nvidia.com/cuda/toolkit

###Nodes Included 🟦 RTX Engine Builder (TensorRTXEngineLoader) Loads a TensorRT-RTX engine. If the engine file for your specific GPU and resolution settings does not exist, it will automatically build one from the source ONNX model. 🟦 RTX Upscale Image (TensorRTXUpscaleImage) The core execution node. It takes an image and a loaded engine to perform the upscale. It supports an optional resize input for custom final dimensions. 🟦 RTX Engine Dynamic Shape Config (TensorRTXDynamicShapeConfig) Allows you to define the resolution bounds (min/opt/max) for the engine. Min/Max: The range of resolutions the engine can handle. Opt: The resolution the engine is most optimized for.

GPU: NVIDIA RTX Series (TensorRT-RTX requires Tensor Cores). Software: NVIDIA Drivers and a compatible version of TensorRT (handled via requirements.txt).

###Troubleshooting First Build Time: The first time you load a model with new shape settings, it may take several minutes to compile the engine. Check the console for progress via the TQDM bar. VRAM Errors: If the build fails, try reducing the width_max or height_max in the Dynamic Shape Config.

🛠️ Supported Models

• These upscaler models have been tested to work with Tensorrt. Onnx are available here

• The exported tensorrt models support dynamic image resolutions from 256x256 to 1280x1280 px (e.g 960x540, 512x512, 1280x720 etc).

• It can support more if you use Dynamic Shape Config to generate a larger resolution spread.

  • 4x-AnimeSharp
  • 4x-UltraSharp
  • 4x-WTP-UDS-Esrgan
  • 4x_NMKD-Siax_200k
  • 4x_RealisticRescaler_100000_G
  • 4x_foolhardy_Remacri
  • RealESRGAN_x4
  • 4xNomos2_otf_esrgan
  • 4x-ClearRealityV1
  • 4x_UniversalUpscalerV2-Neutral_115000_swaG
  • 4x-UltraSharpV2_Lite

🔧 Custom Models

• In the scripts folder there are conversion scripts that can take a .pth file and convert it to an onnx file
• Also included is a file to build all trt files from onnx files.
• NOTE: Everytime you update your graphics drivers or the tensorrt_rtx package, the trt files will need to be rebuilt