311 lines
12 KiB
Python
311 lines
12 KiB
Python
"""
|
|
Main node definitions for the TensorRT_RTX upscaling extension.
|
|
|
|
This module provides ComfyUI nodes for building, loading, and executing
|
|
TensorRT-optimized upscaling engines. It includes support for dynamic
|
|
shape configurations and automatic engine compilation from ONNX models.
|
|
"""
|
|
import os
|
|
import time
|
|
import gc
|
|
import logging
|
|
|
|
import numpy as np
|
|
import tensorrt_rtx as trt
|
|
import torch
|
|
|
|
import comfy.model_management as mm
|
|
import folder_paths
|
|
|
|
from .utilities import Engine
|
|
from comfy.utils import ProgressBar
|
|
from spandrel import ModelLoader, ImageModelDescriptor
|
|
|
|
|
|
IMAGE_DIM_MIN = 256
|
|
IMAGE_DIM_OPT = 512
|
|
IMAGE_DIM_MAX = 1280
|
|
|
|
IMAGE_BATCH_MIN = 1
|
|
IMAGE_BATCH_MAX = 1
|
|
|
|
def update_folder_names_and_paths(key, targets=[]):
|
|
"""
|
|
Ensures specific folder paths are registered within ComfyUI's path manager.
|
|
|
|
Args:
|
|
key (str): The key to register in folder_names_and_paths.
|
|
targets (list): Fallback keys to inherit paths from if the primary key is missing.
|
|
"""
|
|
# check for existing key
|
|
base = folder_paths.folder_names_and_paths.get(key, ([], {}))
|
|
base = base[0] if isinstance(base[0], (list, set, tuple)) else []
|
|
# find base key & add w/ fallback, sanity check + warning
|
|
target = next((x for x in targets if x in folder_paths.folder_names_and_paths), targets[0])
|
|
orig, _ = folder_paths.folder_names_and_paths.get(target, ([], {}))
|
|
folder_paths.folder_names_and_paths[key] = (orig or base, {".onnx"})
|
|
if base and base != orig:
|
|
logging.warning(f"Unknown file list already present on key {key}: {base}")
|
|
|
|
# Add a custom keys for files ending in .onnx
|
|
update_folder_names_and_paths("onnx_upscalers", ["onnx"])
|
|
|
|
class TensorRTXUpscaleImage:
|
|
"""
|
|
A ComfyUI node that performs image upscaling using a loaded TensorRT engine.
|
|
"""
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE", {"tooltip": "Images to be upscaled. Resolution must be between 256 and 1280 px if no build options has been specified."}),
|
|
"engine": ("TRT_RTX_ENGINE", {"tooltip": "TensorRT_RTX engine built and loaded"}),
|
|
"scale_by": (["2x", "3x", "4x", "custom"], {"default": "2x", "tooltip": "Target upscaling factor"})
|
|
},
|
|
"optional": {
|
|
"width": ("INT", {"default": 1024, "min": 256, "max": 4096}),
|
|
"height": ("INT", {"default": 1024, "min": 256, "max": 4096}),
|
|
},
|
|
}
|
|
|
|
RETURN_NAMES = ("upscaled_images",)
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "upscale"
|
|
CATEGORY = "TensorRT_RTX/Upscale"
|
|
DESCRIPTION = "Upscale images with TensorRT_RTX"
|
|
|
|
def upscale(self, **kwargs):
|
|
"""
|
|
Executes the upscaling inference on the provided images.
|
|
|
|
Args:
|
|
image (torch.Tensor): Input image tensor in NHWC format.
|
|
engine (TRT_RTX_ENGINE): The Tensorrt_RTX object.
|
|
resize (dict, optional): Desired output dimensions. Defaults to 4x scale.
|
|
|
|
Returns:
|
|
tuple[torch.Tensor]: A tuple containing the upscaled image tensor in NHWC format.
|
|
"""
|
|
images = kwargs.get("images")
|
|
engine = kwargs.get("engine")
|
|
scale_by = kwargs.get("scale_by")
|
|
images_bchw = images.permute(0, 3, 1, 2)
|
|
B, C, H, W = images_bchw.shape
|
|
|
|
# Bound checking
|
|
for dim in (H,W):
|
|
if not IMAGE_DIM_MIN <= dim <= IMAGE_DIM_MAX:
|
|
raise ValueError(f"The input image dimension ({dim}) is outside of Tensorrt_RTX engine bounds: {IMAGE_DIM_MIN} to {IMAGE_DIM_MAX}.\nImage Dimensions:({W}x{H})")
|
|
|
|
final_width, final_height = None, None
|
|
|
|
if scale_by == "custom":
|
|
final_width = kwargs.get("width")
|
|
final_height = kwargs.get("height")
|
|
else:
|
|
resize_factor = float(scale_by.split('x')[0])
|
|
final_width = int(W * resize_factor)
|
|
final_height = int(H * resize_factor)
|
|
|
|
logging.info(f"Upscaling {B} images from H:{H}, W:{W} to H:{H*4}, W:{W*4} | Final resolution: H:{final_height}, W:{final_width}")
|
|
|
|
shape_dict = {
|
|
"input": {"shape": (1, 3, H, W)},
|
|
"output": {"shape": (1, 3, H*4, W*4)},
|
|
}
|
|
|
|
# setup engine
|
|
engine.activate()
|
|
engine.allocate_buffers(shape_dict=shape_dict)
|
|
|
|
cudaStream = torch.cuda.current_stream().cuda_stream
|
|
pbar = ProgressBar(B)
|
|
images_list = list(torch.split(images_bchw, split_size_or_sections=1))
|
|
upscaled_frames = torch.empty((B, C, final_height, final_width), dtype=torch.float32, device=mm.intermediate_device()) # offloaded to cpu
|
|
|
|
for i, img in enumerate(images_list):
|
|
result = engine.infer({"input": img}, cudaStream)
|
|
result = result["output"]
|
|
|
|
if W*4 != final_width or H*4 != final_height:
|
|
if W*4 > final_width and H*4 > final_height:
|
|
mode='area'
|
|
antialias=False
|
|
else:
|
|
mode='bicubic'
|
|
antialias=True
|
|
|
|
result = torch.nn.functional.interpolate(result, size=(final_height, final_width), mode=mode, antialias=antialias)
|
|
|
|
upscaled_frames[i] = result.to(mm.intermediate_device())
|
|
pbar.update(1)
|
|
|
|
output = upscaled_frames.permute(0, 2, 3, 1)
|
|
|
|
#Free VRAM
|
|
engine.reset()
|
|
gc.collect()
|
|
mm.soft_empty_cache()
|
|
|
|
logging.info(f"Output shape: {output.shape}")
|
|
return (output,)
|
|
|
|
class TensorRTXEngineLoader:
|
|
"""
|
|
A ComfyUI node that loads a TensorRT engine, building it from ONNX if necessary.
|
|
"""
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
onnx_names = [x for x in folder_paths.get_filename_list("onnx_upscalers") if "rife" not in x.lower()]
|
|
return {
|
|
"required": { "onnx_name": (onnx_names,)},
|
|
"optional": {
|
|
"config": ("TRT_RTX_ENGINE_CONFIG", {"tooltip": "Options for building the TensorRT engine"}),
|
|
}
|
|
}
|
|
|
|
RETURN_NAMES = ("ENGINE",)
|
|
RETURN_TYPES = ("TRT_RTX_ENGINE",)
|
|
CATEGORY = "TensorRT_RTX/Upscaler"
|
|
DESCRIPTION = "Load TensorRT model (the model will be built automatically if not found)"
|
|
FUNCTION = "load_upscaler_tensorrt_model"
|
|
|
|
def load_upscaler_tensorrt_model(self, onnx_name, config=None):
|
|
"""
|
|
Loads the .trt engine or compiles it from .onnx using specified dynamic shapes.
|
|
|
|
Args:
|
|
onnx_name (str): The filename of the source ONNX model.
|
|
config (TensorRTXDynamicShapeOptions, optional): Configuration for dynamic shapes.
|
|
|
|
Returns:
|
|
engine [TRT_ENGINE]: Contains the 'engine' instance.
|
|
"""
|
|
if config is None:
|
|
config = TensorRTXDynamicShapeOptions()
|
|
|
|
tensorrt_models_dir = os.path.join(folder_paths.models_dir, "tensorrt", "upscaler")
|
|
onnx_models_dir = os.path.join(folder_paths.models_dir, "onnx")
|
|
|
|
os.makedirs(tensorrt_models_dir, exist_ok=True)
|
|
os.makedirs(onnx_models_dir, exist_ok=True)
|
|
|
|
onnx_model_path = os.path.join(onnx_models_dir, f"{onnx_name}")
|
|
|
|
# Dynamic shape configuration
|
|
engine_channel = 3
|
|
engine_min_batch, engine_opt_batch, engine_max_batch = config.batch_min, config.batch_opt, config.batch_max
|
|
engine_min_h, engine_opt_h, engine_max_h = config.height_min, config.height_opt, config.height_max
|
|
engine_min_w, engine_opt_w, engine_max_w = config.width_min, config.width_opt, config.width_max
|
|
tensorrt_model_path = os.path.join(tensorrt_models_dir, f"{onnx_name}_{engine_min_batch}x{engine_channel}x{engine_min_h}x{engine_min_w}_{engine_opt_batch}x{engine_channel}x{engine_opt_h}x{engine_opt_w}_{engine_max_batch}x{engine_channel}x{engine_max_h}x{engine_max_w}_{trt.__version__}.trt")
|
|
|
|
# Check if the TRT file exists and build it if it doesn't
|
|
engine = Engine(tensorrt_model_path)
|
|
|
|
if not os.path.exists(tensorrt_model_path):
|
|
logging.info(f"Building TensorRT_RTX engine for {onnx_model_path}: {tensorrt_model_path}")
|
|
|
|
mm.soft_empty_cache()
|
|
s = time.time()
|
|
result = engine.build(
|
|
onnx_path=onnx_model_path,
|
|
input_profile=[
|
|
{"input": [
|
|
(engine_min_batch,engine_channel,engine_min_h,engine_min_w),
|
|
(engine_opt_batch,engine_channel,engine_opt_h,engine_min_w),
|
|
(engine_max_batch,engine_channel,engine_max_h,engine_max_w)
|
|
]},
|
|
]
|
|
)
|
|
if result != 0:
|
|
raise ValueError(f"Failed to build the engine")
|
|
e = time.time()
|
|
logging.info(f"Time taken to build: {(e-s)} seconds")
|
|
|
|
# Load tensorrt model
|
|
logging.info(f"Loading TensorRT_RTX engine: {tensorrt_model_path}")
|
|
mm.soft_empty_cache()
|
|
engine.load()
|
|
|
|
return (engine,)
|
|
|
|
class TensorRTXDynamicShapeOptions:
|
|
"""
|
|
Data container for TensorRT dynamic shape boundaries and optimal values.
|
|
"""
|
|
width_min: int = 256
|
|
width_opt: int = 512
|
|
width_max: int = 1280
|
|
height_min: int = 256
|
|
height_opt: int = 512
|
|
height_max: int = 1280
|
|
batch_min: int = 1
|
|
batch_opt: int = 1
|
|
batch_max: int = 1
|
|
|
|
|
|
class TensorRTXDynamicShapeConfig:
|
|
"""
|
|
A ComfyUI node that allows users to configure TensorRT dynamic shape profiles.
|
|
"""
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"width_min" : ("INT", {"default": TensorRTXDynamicShapeOptions.width_min, "min": 1, "max": 4096}),
|
|
"width_opt" : ("INT", {"default": TensorRTXDynamicShapeOptions.width_opt, "min": 1, "max": 4096}),
|
|
"width_max" : ("INT", {"default": TensorRTXDynamicShapeOptions.width_max, "min": 1, "max": 4096}),
|
|
"height_min" : ("INT", {"default": TensorRTXDynamicShapeOptions.height_min, "min": 1, "max": 4096}),
|
|
"height_opt" : ("INT", {"default": TensorRTXDynamicShapeOptions.height_opt, "min": 1, "max": 4096}),
|
|
"height_max" : ("INT", {"default": TensorRTXDynamicShapeOptions.height_max, "min": 1, "max": 4096}),
|
|
"batch_min" : ("INT", {"default": TensorRTXDynamicShapeOptions.batch_min, "min": 1}),
|
|
"batch_opt" : ("INT", {"default": TensorRTXDynamicShapeOptions.batch_opt, "min": 1}),
|
|
"batch_max" : ("INT", {"default": TensorRTXDynamicShapeOptions.batch_max, "min": 1}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("TRT_RTX_ENGINE_CONFIG",)
|
|
RETURN_NAMES = ("CONFIG",)
|
|
CATEGORY = "TensorRT_RTX/Upscaler"
|
|
FUNCTION = "package"
|
|
|
|
@classmethod
|
|
def VALIDATE_INPUTS(cls, **kwargs):
|
|
"""
|
|
Validates that min <= opt <= max for all shape dimensions.
|
|
"""
|
|
if not all(v is not None for v in kwargs.values()):
|
|
raise ValueError(f"Dynamic shape properties must be numbers")
|
|
|
|
# Check dimensions logic (min <= opt <= max)
|
|
for dim in ['width', 'height', 'batch']:
|
|
mi, opt, ma = kwargs[f'{dim}_min'], kwargs[f'{dim}_opt'], kwargs[f'{dim}_max']
|
|
if not 1 <= mi <= opt <= ma:
|
|
raise ValueError(f"The {dim} bounds are invalid. Must follow: 1 <= min <= opt <= max")
|
|
|
|
return True
|
|
|
|
def package(self, **kwargs):
|
|
"""
|
|
Instantiates and populates a TensorRTXDynamicShapeOptions object.
|
|
"""
|
|
|
|
options = TensorRTXDynamicShapeOptions()
|
|
|
|
for key, value in kwargs.items():
|
|
setattr(options, key, value)
|
|
return (options,)
|
|
|
|
# Node Mappings
|
|
NODE_CLASS_MAPPINGS = {
|
|
"TensorRTXUpscaleImage": TensorRTXUpscaleImage,
|
|
"TensorRTXEngineLoader": TensorRTXEngineLoader,
|
|
"TensorRTXDynamicShapeConfig": TensorRTXDynamicShapeConfig,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"TensorRTXUpscaleImage": "RTX Upscale Image",
|
|
"TensorRTXEngineLoader": "RTX Engine Loader",
|
|
"TensorRTXDynamicShapeConfig": "RTX Engine Dynamic Shape Config",
|
|
}
|