nodes/huchukato-ComfyUI-RIFE-TensorRT-Auto
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge pull request #1 from silveroxides/PR14-reaperhammer

Browse Source 
Merge pr from reaperhammer
This commit is contained in:
Silver
2025-11-29 20:12:43 +01:00
committed by GitHub
parent 41686c5c64 d61e185923
commit 00cbebe151
9 changed files with 361 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+145 -20
View File
@@ -3,23 +3,149 @@ import os
from comfy.model_management import get_torch_device
from .vfi_utilities import preprocess_frames, postprocess_frames, generate_frames_rife, logger
from .trt_utilities import Engine
from .utilities import download_file, ColoredLogger
import folder_paths
import time
from polygraphy import cuda
import comfy.model_management as mm
import tensorrt
import json
ENGINE_DIR = os.path.join(folder_paths.models_dir, "tensorrt", "rife")
# Image dimensions for TensorRT engine building
IMAGE_DIM_MIN = 256
IMAGE_DIM_OPT = 512
IMAGE_DIM_MAX = 3840
# Logger for this module
rife_logger = ColoredLogger("ComfyUI-Rife-Tensorrt")
# Function to load configuration
def load_node_config(config_filename="load_rife_config.json"):
"""Loads node configuration from a JSON file."""
current_dir = os.path.dirname(__file__)
config_path = os.path.join(current_dir, config_filename)
default_config = {
"model": {
"options": ["rife49_ensemble_True_scale_1_sim"],
"default": "rife49_ensemble_True_scale_1_sim",
"tooltip": "Default model (fallback from code)"
},
"precision": {
"options": ["fp16", "fp32"],
"default": "fp16",
"tooltip": "Default precision (fallback from code)"
}
}
try:
with open(config_path, 'r') as f:
config = json.load(f)
rife_logger.info(f"Successfully loaded configuration from {config_filename}")
return config
except FileNotFoundError:
rife_logger.warning(f"Configuration file '{config_path}' not found. Using default fallback configuration.")
return default_config
except json.JSONDecodeError:
rife_logger.error(f"Error decoding JSON from '{config_path}'. Using default fallback configuration.")
return default_config
except Exception as e:
rife_logger.error(f"An unexpected error occurred while loading '{config_path}': {e}. Using default fallback.")
return default_config
# Load the configuration once when the module is imported
LOAD_RIFE_NODE_CONFIG = load_node_config()
class LoadRifeTensorrtModel:
@classmethod
def INPUT_TYPES(cls):
# Use the pre-loaded configuration
model_config = LOAD_RIFE_NODE_CONFIG.get("model", {})
precision_config = LOAD_RIFE_NODE_CONFIG.get("precision", {})
# Provide sensible defaults if keys are missing in the config
model_options = model_config.get("options", ["rife49_ensemble_True_scale_1_sim"])
model_default = model_config.get("default", "rife49_ensemble_True_scale_1_sim")
model_tooltip = model_config.get("tooltip", "Select a RIFE model.")
precision_options = precision_config.get("options", ["fp16", "fp32"])
precision_default = precision_config.get("default", "fp16")
precision_tooltip = precision_config.get("tooltip", "Select precision.")
return {
"required": {
"model": (model_options, {"default": model_default, "tooltip": model_tooltip}),
"precision": (precision_options, {"default": precision_default, "tooltip": precision_tooltip}),
}
}
RETURN_NAMES = ("rife_trt_model",)
RETURN_TYPES = ("RIFE_TRT_MODEL",)
CATEGORY = "tensorrt"
DESCRIPTION = "Load RIFE tensorrt models, they will be built automatically if not found."
FUNCTION = "load_rife_tensorrt_model"
def load_rife_tensorrt_model(self, model, precision):
tensorrt_models_dir = os.path.join(folder_paths.models_dir, "tensorrt", "rife")
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"{model}.onnx")
# Build tensorrt model path with detailed naming
engine_channel = 3
engine_min_batch, engine_opt_batch, engine_max_batch = 1, 1, 1
engine_min_h, engine_opt_h, engine_max_h = IMAGE_DIM_MIN, IMAGE_DIM_OPT, IMAGE_DIM_MAX
engine_min_w, engine_opt_w, engine_max_w = IMAGE_DIM_MIN, IMAGE_DIM_OPT, IMAGE_DIM_MAX
tensorrt_model_path = os.path.join(tensorrt_models_dir, f"{model}_{precision}_{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}_{tensorrt.__version__}.trt")
if not os.path.exists(tensorrt_model_path):
if not os.path.exists(onnx_model_path):
onnx_model_download_url = f"https://huggingface.co/yuvraj108c/rife-onnx/resolve/main/{model}.onnx"
rife_logger.info(f"Downloading {onnx_model_download_url}")
download_file(url=onnx_model_download_url, save_path=onnx_model_path)
else:
rife_logger.info(f"ONNX model found at: {onnx_model_path}")
rife_logger.info(f"Building TensorRT engine for {onnx_model_path}: {tensorrt_model_path}")
mm.soft_empty_cache()
s = time.time()
engine = Engine(tensorrt_model_path)
engine.build(
onnx_path=onnx_model_path,
fp16=True if precision == "fp16" else False,
input_profile=[
{
"img0": [(engine_min_batch, engine_channel, engine_min_h, engine_min_w), (engine_opt_batch, engine_channel, engine_opt_h, engine_opt_w), (engine_max_batch, engine_channel, engine_max_h, engine_max_w)],
"img1": [(engine_min_batch, engine_channel, engine_min_h, engine_min_w), (engine_opt_batch, engine_channel, engine_opt_h, engine_opt_w), (engine_max_batch, engine_channel, engine_max_h, engine_max_w)],
}
],
)
e = time.time()
rife_logger.info(f"Time taken to build: {(e-s)} seconds")
rife_logger.info(f"Loading TensorRT engine: {tensorrt_model_path}")
mm.soft_empty_cache()
engine = Engine(tensorrt_model_path)
engine.load()
return (engine,)
class RifeTensorrt:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"frames": ("IMAGE", ),
"engine": (os.listdir(ENGINE_DIR),),
"clear_cache_after_n_frames": ("INT", {"default": 100, "min": 1, "max": 1000}),
"multiplier": ("INT", {"default": 2, "min": 1}),
"use_cuda_graph": ("BOOLEAN", {"default": True}),
"keep_model_loaded": ("BOOLEAN", {"default": False}),
"frames": ("IMAGE", {"tooltip": "Input frames for video frame interpolation"}),
"rife_trt_model": ("RIFE_TRT_MODEL", {"tooltip": "Tensorrt model built and loaded"}),
"clear_cache_after_n_frames": ("INT", {"default": 100, "min": 1, "max": 1000, "tooltip": "Clear CUDA cache after processing this many frames"}),
"multiplier": ("INT", {"default": 2, "min": 1, "tooltip": "Frame interpolation multiplier"}),
"use_cuda_graph": ("BOOLEAN", {"default": True, "tooltip": "Use CUDA graph for better performance"}),
"keep_model_loaded": ("BOOLEAN", {"default": False, "tooltip": "Keep model loaded in memory after processing"}),
},
}
@@ -31,7 +157,7 @@ class RifeTensorrt:
def vfi(
self,
frames,
engine,
rife_trt_model,
clear_cache_after_n_frames=100,
multiplier=2,
use_cuda_graph=True,
@@ -45,24 +171,21 @@ class RifeTensorrt:
}
cudaStream = cuda.Stream()
engine_path = os.path.join(ENGINE_DIR, engine)
if (not hasattr(self, 'engine') or self.engine_label != engine):
self.engine = Engine(engine_path)
logger(f"Loading TensorRT engine: {engine_path}")
self.engine.load()
self.engine.activate()
self.engine_label = engine
else:
logger(f"Using cached TensorRT engine: {engine_path}")
self.engine.allocate_buffers(shape_dict=shape_dict)
# Use the provided model directly
engine = rife_trt_model
logger(f"Using loaded TensorRT engine")
# Activate and allocate buffers for the engine
engine.activate()
engine.allocate_buffers(shape_dict=shape_dict)
frames = preprocess_frames(frames)
def return_middle_frame(frame_0, frame_1, timestep):
timestep_t = torch.tensor([timestep], dtype=torch.float32).to(get_torch_device())
# s = time.time()
output = self.engine.infer({"img0": frame_0, "img1": frame_1, "timestep": timestep_t}, cudaStream, use_cuda_graph)
output = engine.infer({"img0": frame_0, "img1": frame_1, "timestep": timestep_t}, cudaStream, use_cuda_graph)
# e = time.time()
# print(f"Time taken to infer: {(e-s)*1000} ms")
@@ -71,19 +194,21 @@ class RifeTensorrt:
result = generate_frames_rife(frames, clear_cache_after_n_frames, multiplier, return_middle_frame)
out = postprocess_frames(result)
if not keep_model_loaded:
del self.engine, self.engine_label
engine.reset()
return (out,)
NODE_CLASS_MAPPINGS = {
"RifeTensorrt": RifeTensorrt,
"LoadRifeTensorrtModel": LoadRifeTensorrtModel,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"RifeTensorrt": "⚡ Rife Tensorrt",
"LoadRifeTensorrtModel": "Load Rife Tensorrt Model",
}
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
export_onnx.py
+6 -6
View File
@@ -14,10 +14,10 @@ from onnxsim import simplify
MODEL_TYPE = pathlib.Path(__file__).parent.name
CKPT_NAME_VER_DICT = {
"rife40.pth": "4.0",
"rife41.pth": "4.0",
"rife42.pth": "4.2",
"rife43.pth": "4.3",
"rife44.pth": "4.3",
"rife41.pth": "4.0",
"rife42.pth": "4.2",
"rife43.pth": "4.3",
"rife44.pth": "4.3",
"rife45.pth": "4.5",
"rife46.pth": "4.6",
"rife47.pth": "4.7",
@@ -106,7 +106,7 @@ def export_onnx(ckpt_name, ensemble, scale_factor):
'img1': {2: 'height', 3: 'width'},
'output': {2: 'height', 3: 'width'},
}
with torch.no_grad(): # Disable gradients for efficiency
torch.onnx.export(interpolation_model,
(img0, img1, timestep), # ensemble + scale_factor set in forward fn
@@ -124,7 +124,7 @@ def export_onnx(ckpt_name, ensemble, scale_factor):
onnx_model = onnx.load(onnx_save_path)
onnx.checker.check_model(onnx_model) # Perform a validity check
print("ONNX model validation successful!")
# print(onnx.helper.printable_graph(onnx_model.graph))
sim_model_path = os.path.join("models", f"{ckpt_name.split('.')[0]}_ensemble_{ensemble}_scale_{scale_factor}_sim.onnx")
load_rife_config.json
+16
View File
@@ -0,0 +1,16 @@
{
"model": {
"options": [
"rife47_ensemble_True_scale_1_sim",
"rife48_ensemble_True_scale_1_sim",
"rife49_ensemble_True_scale_1_sim"
],
"default": "rife49_ensemble_True_scale_1_sim",
"tooltip": "RIFE models for video frame interpolation. These models have been tested with tensorrt. Loaded from config."
},
"precision": {
"options": ["fp16", "fp32"],
"default": "fp16",
"tooltip": "Precision to build the tensorrt engines. Loaded from config."
}
}
readme.md
+27 -14
View File
@@ -2,9 +2,9 @@
# ComfyUI Rife TensorRT ⚡
[![python](https://img.shields.io/badge/python-3.10.12-green)](https://www.python.org/downloads/release/python-31012/)
[![cuda](https://img.shields.io/badge/cuda-12.4-green)](https://developer.nvidia.com/cuda-downloads)
[![trt](https://img.shields.io/badge/TRT-10.4.0-green)](https://developer.nvidia.com/tensorrt)
[![python](https://img.shields.io/badge/python-3.12.11-green)](https://www.python.org/downloads/release/python-31211/)
[![cuda](https://img.shields.io/badge/cuda-12.9-green)](https://developer.nvidia.com/cuda-downloads)
[![trt](https://img.shields.io/badge/TRT-10.13.3.9-green)](https://developer.nvidia.com/tensorrt)
[![by-nc-sa/4.0](https://img.shields.io/badge/license-CC--BY--NC--SA--4.0-lightgrey)](https://creativecommons.org/licenses/by-nc-sa/4.0/deed.en)
![node](https://github.com/user-attachments/assets/5fd6d529-300c-42a5-b9cf-46e031f0bcb5)
@@ -40,27 +40,40 @@ cd ./ComfyUI-Rife-Tensorrt
pip install -r requirements.txt
```
## 🛠️ Building Tensorrt Engine
## 🛠️ Supported Models
1. Download one of the following onnx models:
- [rife49_ensemble_True_scale_1_sim.onnx](https://huggingface.co/yuvraj108c/rife-onnx/resolve/main/rife49_ensemble_True_scale_1_sim.onnx)
- [rife48_ensemble_True_scale_1_sim.onnx](https://huggingface.co/yuvraj108c/rife-onnx/resolve/main/rife48_ensemble_True_scale_1_sim.onnx)
- [rife47_ensemble_True_scale_1_sim.onnx](https://huggingface.co/yuvraj108c/rife-onnx/resolve/main/rife47_ensemble_True_scale_1_sim.onnx)
2. Edit onnx/trt paths inside [export_trt.py](./export_trt.py) and build tensorrt engine by running:
- `python export_trt.py`
The following RIFE models are supported and will be automatically downloaded and built:
- **rife49_ensemble_True_scale_1_sim** (default) - Latest and most accurate
- **rife48_ensemble_True_scale_1_sim** - Good balance of speed and quality
- **rife47_ensemble_True_scale_1_sim** - Fastest option
3. Place the exported engine inside ComfyUI `/models/tensorrt/rife` directory
Models are automatically downloaded from [HuggingFace](https://huggingface.co/yuvraj108c/rife-onnx) and TensorRT engines are built on first use.
## ☀️ Usage
- Insert node by `Right Click -> tensorrt -> Rife Tensorrt`
- Image resolutions between `256x256` and `3840x3840` will work with the tensorrt engines
1. **Load Model**: Insert `Right Click -> Add Node -> tensorrt -> Load Rife Tensorrt Model`
- Choose your preferred RIFE model (rife47, rife48, or rife49)
- Select precision (fp16 recommended for speed, fp32 for maximum accuracy)
- The model will be automatically downloaded and TensorRT engine built on first use
2. **Process Frames**: Insert `Right Click -> Add Node -> tensorrt -> Rife Tensorrt`
- Connect the loaded model from step 1
- Input your video frames
- Configure interpolation settings (multiplier, CUDA graph, etc.)
- Image resolutions between `256x256` and `3840x3840` are supported
## 🤖 Environment tested
- Ubuntu 22.04 LTS, Cuda 12.4, Tensorrt 10.4.0, Python 3.10, RTX 3070 GPU
- WSL Ubuntu 24.04.03 LTS, Cuda 12.9, Tensorrt 10.13.3.9, Python 3.12.11, RTX 5080 GPU
- Windows (Not tested, but should work)
## 🚨 Updates
### December 2025
- **Automatic Model Management**: No more manual downloads! Models are automatically downloaded from HuggingFace and TensorRT engines are built on demand
- **Improved Workflow**: New two-node system with `Load Rife Tensorrt Model` + `Rife Tensorrt` for better organization
- **Updated Dependencies**: TensorRT updated to 10.13.3.9 for better performance and compatibility
## 👏 Credits
- https://github.com/styler00dollar/VSGAN-tensorrt-docker
requirements.txt
+4 -2
View File
@@ -1,5 +1,7 @@
einops
colored
polygraphy
tensorrt==10.4.0
cuda-python
tensorrt==10.13.3.9
cuda-python
requests
tqdm
rife_arch.py
+1 -1
View File
@@ -532,6 +532,6 @@ class IFNet(nn.Module):
tmp = self.unet(img0, img1, warped_img0, warped_img1, mask, flow, c0, c1)
res = tmp[:, :3] * 2 - 1
merged[3] = torch.clamp(merged[3] + res, 0, 1)
result = merged[3][:, :, :h, :w]
return result
trt_utilities.py
+57 -5
View File
@@ -136,19 +136,53 @@ class Engine:
self.buffers = OrderedDict()
self.tensors = OrderedDict()
self.cuda_graph_instance = None # cuda graph
self.graph = None
def __del__(self):
# Clean up CUDA graph resources
if hasattr(self, 'cuda_graph_instance') and self.cuda_graph_instance is not None:
try:
cudart.cudaGraphDestroy(self.cuda_graph_instance)
except:
pass
if hasattr(self, 'graph') and self.graph is not None:
try:
cudart.cudaGraphDestroy(self.graph)
except:
pass
del self.engine
del self.context
del self.buffers
del self.tensors
def reset(self, engine_path=None):
del self.engine
del self.context
del self.buffers
del self.tensors
self.engine_path = engine_path
# Clean up CUDA graph resources first
if hasattr(self, 'cuda_graph_instance') and self.cuda_graph_instance is not None:
try:
cudart.cudaGraphDestroy(self.cuda_graph_instance)
except:
pass
self.cuda_graph_instance = None
if hasattr(self, 'graph') and self.graph is not None:
try:
cudart.cudaGraphDestroy(self.graph)
except:
pass
self.graph = None
if hasattr(self, 'engine') and self.engine is not None:
del self.engine
if hasattr(self, 'context') and self.context is not None:
del self.context
if hasattr(self, 'buffers'):
del self.buffers
if hasattr(self, 'tensors'):
del self.tensors
self.engine = None
self.context = None
self.engine_path = engine_path if engine_path else self.engine_path
self.buffers = OrderedDict()
self.tensors = OrderedDict()
@@ -220,6 +254,10 @@ class Engine:
self.engine = engine_from_bytes(bytes_from_path(self.engine_path))
def activate(self, reuse_device_memory=None):
# If engine was reset, reload it
if self.engine is None:
self.load()
if reuse_device_memory:
self.context = self.engine.create_execution_context_without_device_memory()
# self.context.device_memory = reuse_device_memory
@@ -227,6 +265,20 @@ class Engine:
self.context = self.engine.create_execution_context()
def allocate_buffers(self, shape_dict=None, device="cuda"):
# Clean up CUDA graph resources since tensors will be recreated
if hasattr(self, 'cuda_graph_instance') and self.cuda_graph_instance is not None:
try:
cudart.cudaGraphDestroy(self.cuda_graph_instance)
except:
pass
self.cuda_graph_instance = None
if hasattr(self, 'graph') and self.graph is not None:
try:
cudart.cudaGraphDestroy(self.graph)
except:
pass
self.graph = None
nvtx.range_push("allocate_buffers")
for idx in range(self.engine.num_io_tensors):
name = self.engine.get_tensor_name(idx)
utilities.py
+101
View File
@@ -0,0 +1,101 @@
import requests
from tqdm import tqdm
import logging
import sys
class ColoredLogger:
COLORS = {
'RED': '\033[91m',
'GREEN': '\033[92m',
'YELLOW': '\033[93m',
'BLUE': '\033[94m',
'MAGENTA': '\033[95m',
'RESET': '\033[0m'
}
LEVEL_COLORS = {
'DEBUG': COLORS['BLUE'],
'INFO': COLORS['GREEN'],
'WARNING': COLORS['YELLOW'],
'ERROR': COLORS['RED'],
'CRITICAL': COLORS['MAGENTA']
}
def __init__(self, name="MY-APP"):
self.logger = logging.getLogger(name)
self.logger.setLevel(logging.DEBUG)
self.app_name = name
# Prevent message propagation to parent loggers
self.logger.propagate = False
# Clear existing handlers
self.logger.handlers = []
# Create console handler
handler = logging.StreamHandler(sys.stdout)
handler.setLevel(logging.DEBUG)
# Custom formatter class to handle colored components
class ColoredFormatter(logging.Formatter):
def format(self, record):
# Color the level name according to severity
level_color = ColoredLogger.LEVEL_COLORS.get(record.levelname, '')
colored_levelname = f"{level_color}{record.levelname}{ColoredLogger.COLORS['RESET']}"
# Color the logger name in blue
colored_name = f"{ColoredLogger.COLORS['BLUE']}{record.name}{ColoredLogger.COLORS['RESET']}"
# Set the colored components
record.levelname = colored_levelname
record.name = colored_name
return super().format(record)
# Create formatter with the new format
formatter = ColoredFormatter('[%(name)s|%(levelname)s] - %(message)s')
handler.setFormatter(formatter)
self.logger.addHandler(handler)
def debug(self, message):
self.logger.debug(f"{self.COLORS['BLUE']}{message}{self.COLORS['RESET']}")
def info(self, message):
self.logger.info(f"{self.COLORS['GREEN']}{message}{self.COLORS['RESET']}")
def warning(self, message):
self.logger.warning(f"{self.COLORS['YELLOW']}{message}{self.COLORS['RESET']}")
def error(self, message):
self.logger.error(f"{self.COLORS['RED']}{message}{self.COLORS['RESET']}")
def critical(self, message):
self.logger.critical(f"{self.COLORS['MAGENTA']}{message}{self.COLORS['RESET']}")
def download_file(url, save_path):
"""
Download a file from URL with progress bar
Args:
url (str): URL of the file to download
save_path (str): Path to save the file as
"""
GREEN = '\033[92m'
RESET = '\033[0m'
response = requests.get(url, stream=True)
total_size = int(response.headers.get('content-length', 0))
with open(save_path, 'wb') as file, tqdm(
desc=save_path,
total=total_size,
unit='iB',
unit_scale=True,
unit_divisor=1024,
colour='green',
bar_format=f'{GREEN}{{l_bar}}{{bar}}{RESET}{GREEN}{{r_bar}}{RESET}'
) as progress_bar:
for data in response.iter_content(chunk_size=1024):
size = file.write(data)
progress_bar.update(size)
vfi_utilities.py
+4 -4
View File
@@ -7,7 +7,7 @@ import einops
from comfy.model_management import soft_empty_cache, get_torch_device
import numpy as np
from comfy.utils import ProgressBar
from colored import Fore, Back, Style
from colored import Fore, Back, Style
DEVICE = get_torch_device()
@@ -40,7 +40,7 @@ def generate_frames_rife(
multiplier,
return_middle_frame_function
):
output_frames = torch.zeros(multiplier*frames.shape[0], *frames.shape[1:], device="cpu")
out_len = 0
@@ -70,7 +70,7 @@ def generate_frames_rife(
logger("Clearing cache...")
pbar.update(1)
# Append final frame
output_frames[out_len] = frames[-1:]
@@ -80,6 +80,6 @@ def generate_frames_rife(
# clear cache for courtesy
soft_empty_cache()
logger("Final clearing cache done ...")
#
#
res = output_frames[:out_len]
return res