diff --git a/__init__.py b/__init__.py index c13155b..34bd62c 100644 --- a/__init__.py +++ b/__init__.py @@ -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'] diff --git a/export_onnx.py b/export_onnx.py index 97f2948..96b77fc 100644 --- a/export_onnx.py +++ b/export_onnx.py @@ -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") diff --git a/load_rife_config.json b/load_rife_config.json new file mode 100644 index 0000000..3c01bd0 --- /dev/null +++ b/load_rife_config.json @@ -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." + } +} \ No newline at end of file diff --git a/readme.md b/readme.md index 6374357..d165869 100644 --- a/readme.md +++ b/readme.md @@ -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 diff --git a/requirements.txt b/requirements.txt index b479d99..f24af34 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,7 @@ einops colored polygraphy -tensorrt==10.4.0 -cuda-python \ No newline at end of file +tensorrt==10.13.3.9 +cuda-python +requests +tqdm \ No newline at end of file diff --git a/rife_arch.py b/rife_arch.py index 5b1d722..45a2f10 100644 --- a/rife_arch.py +++ b/rife_arch.py @@ -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 \ No newline at end of file diff --git a/trt_utilities.py b/trt_utilities.py index eaa680b..a41d338 100644 --- a/trt_utilities.py +++ b/trt_utilities.py @@ -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) diff --git a/utilities.py b/utilities.py new file mode 100644 index 0000000..0362ed6 --- /dev/null +++ b/utilities.py @@ -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) \ No newline at end of file diff --git a/vfi_utilities.py b/vfi_utilities.py index f9ec616..ae98f74 100644 --- a/vfi_utilities.py +++ b/vfi_utilities.py @@ -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 \ No newline at end of file