From dfbf6377f4cc8d9d0757fa81cb067cd9c6afc3aa Mon Sep 17 00:00:00 2001 From: Threads_Of_Fate Date: Thu, 23 Apr 2026 10:28:20 -0400 Subject: [PATCH] add files --- .gitignore | 3 + LICENSE | 3 +- README.md | 78 +++++- __init__.py | 20 ++ nodes.py | 309 ++++++++++++++++++++++++ rife_arch.py | 537 +++++++++++++++++++++++++++++++++++++++++ scripts/export_onnx.py | 142 +++++++++++ scripts/export_trt.py | 44 ++++ trt_utilities.py | 455 ++++++++++++++++++++++++++++++++++ vfi_utilities.py | 112 +++++++++ 10 files changed, 1700 insertions(+), 3 deletions(-) create mode 100644 .gitignore create mode 100644 __init__.py create mode 100644 nodes.py create mode 100644 rife_arch.py create mode 100644 scripts/export_onnx.py create mode 100644 scripts/export_trt.py create mode 100644 trt_utilities.py create mode 100644 vfi_utilities.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..313b49c --- /dev/null +++ b/.gitignore @@ -0,0 +1,3 @@ +__pycache__ +RIFE_RTX_ENGINE.cache +*.bak \ No newline at end of file diff --git a/LICENSE b/LICENSE index 261eeb9..7a4a3ea 100644 --- a/LICENSE +++ b/LICENSE @@ -1,3 +1,4 @@ + Apache License Version 2.0, January 2004 http://www.apache.org/licenses/ @@ -198,4 +199,4 @@ distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and - limitations under the License. + limitations under the License. \ No newline at end of file diff --git a/README.md b/README.md index 11213e4..2fbaf6e 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,76 @@ -# ComfyUI-RIFE-TensorRT-RTX -ComfyUI custom node that performs RIFE Interpolation using Tensorrt-RTX +# ComfyUI RIFE TensorRT-RTX + +[![python](https://img.shields.io/badge/python-3.13.9-green)](https://www.python.org/downloads/release/python-31313/) +[![cuda](https://img.shields.io/badge/cuda-13.0-green)](https://developer.nvidia.com/cuda-downloads) +[![trt](https://img.shields.io/badge/TRT_RTX-1.4.0.76-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) + +## Overview + +This project is a fork of the original [ComfyUI_RIFE_TensorRT_Auto](https://github.com/silveroxides/ComfyUI_RIFE_TensorRT_Auto) by [silveroxides](https://github.com/silveroxides). +This project uses TensorRT-RTX to do RIFE interpolation + +## Features +- **NVIDIA TensorRT Acceleration**: Leverages Tensor Cores for up to 30x faster frame interpolation 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 + +![Simple Example](assets/example.png) + +## 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-RIFE-TensorRT-RTX` +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](https://developer.nvidia.com/cuda/toolkit) + +###Nodes Included +🟦 RIFE RTX Engine Builder (RIFERTXEngineLoader) + 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. +🟦 RIFE RTX Frame Interpolation (RIFERTXFrameInterpolation) + 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. +🟦 RIFE RTX Engine Config (RIFERTXEngineConfig) + 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 + +rife47_ensemble_True_scale_1_sim +rife48_ensemble_True_scale_1_sim +rife49_ensemble_True_scale_1_sim diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..a6d53a9 --- /dev/null +++ b/__init__.py @@ -0,0 +1,20 @@ +"""ComfyUI-Rife-TensorRT Extension Initialization. + +This module acts as the entry point for ComfyUI, exporting the node classes +and display names required to register the TensorRT-RTX RIFE implementation +into the ComfyUI node graph. + +Attributes: + NODE_CLASS_MAPPINGS (dict): Maps internal lookup strings to node classes. + NODE_DISPLAY_NAME_MAPPINGS (dict): Maps internal lookup strings to + human-readable UI names. + WEB_DIRECTORY (str): Path to the directory containing frontend + javascript extensions. +""" + +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +WEB_DIRECTORY = "./js" + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] + diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..eff3437 --- /dev/null +++ b/nodes.py @@ -0,0 +1,309 @@ +"""TensorRT-RTX RIFE Implementation for ComfyUI. + +This module provides custom nodes for video frame interpolation (VFI) using +NVIDIA's TensorRT-RTX. It includes nodes for configuring engine resolutions, +building/loading optimized .trt engines, and performing high-speed inference. +""" + +import json +import os +import sys + +import folder_paths +import logging +import tensorrt_rtx as trt +import time +import torch + +import comfy.model_management as mm + +from pathlib import Path +from polygraphy import cuda +from .vfi_utilities import preprocess_frames, postprocess_frames, generate_frames_rife +from .trt_utilities import Engine + +def update_folder_names_and_paths(key, targets=[]): + """Registers specific folder paths within ComfyUI's path manager. + + Ensures that custom subdirectories (like onnx_rife) are recognized by + ComfyUI's internal file system and UI loaders. + + Args: + key (str): The primary 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}") + + +# Auto-detect CUDA toolkit and add DLL path before importing polygraphy +def _setup_cuda_dll_path(): + """Auto-detects the CUDA toolkit and adds cudart64 DLL path on Windows. + + This is necessary for TensorRT-RTX to correctly locate runtime libraries + when running in a portable environment or on systems where CUDA is not + globally in the system PATH. + """ + if not sys.platform.startswith("win"): + return + + cuda_root = None + + # Check for CUDA_PATH or CUDA_HOME environment variables + cuda_root = os.environ.get("CUDA_PATH") or os.environ.get("CUDA_HOME") + + if not cuda_root: + # Try default Windows install location + program_files = os.environ.get("PROGRAMFILES") + if program_files: + cuda_base = Path(program_files) / "NVIDIA GPU Computing Toolkit" / "CUDA" + if cuda_base.exists(): + # Find highest version directory + versions = sorted([d for d in cuda_base.iterdir() if d.is_dir()], reverse=True) + if versions: + cuda_root = str(versions[0]) + + if cuda_root: + cuda_path = Path(cuda_root) + # CUDA 13.0+ puts cudart64 in bin/x64 subdirectory + cuda_bin_x64 = cuda_path / "bin" / "x64" + if cuda_bin_x64.exists() and any(cuda_bin_x64.glob("cudart64*.dll")): + os.add_dll_directory(str(cuda_bin_x64)) + return + # Fallback to regular bin directory for older CUDA versions + cuda_bin = cuda_path / "bin" + if cuda_bin.exists() and any(cuda_bin.glob("cudart64*.dll")): + os.add_dll_directory(str(cuda_bin)) + return + + # CUDA toolkit not found - print warning with download link + print("[ComfyUI-Rife-TensorRT] WARNING: CUDA toolkit not found.") + print(" Set CUDA_PATH environment variable or install CUDA toolkit.") + print(" Download: https://developer.nvidia.com/cuda-13-0-2-download-archive") + +# Initialize Environment +_setup_cuda_dll_path() +os.environ["POLYGRAPHY_USE_TENSORRT_RTX"] = "1" +update_folder_names_and_paths("onnx_rife", ["onnx"]) + +class RIFERTXEngineOptions: + """Data container for RIFE shape boundaries. + + Attributes: + min_dim (int): Minimum supported resolution. + opt_dim (int): Targeted optimal resolution. + max_dim (int): Maximum supported resolution. + """ + min_dim: int = 384 + opt_dim: int = 720 + max_dim: int = 1312 + +class RIFERTXEngineConfig: + """ComfyUI Node for configuring resolution profiles. + + Provides a UI interface to set the dynamic shape profiles required + during the TensorRT engine build process. + """ + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "min_dim": ("INT", {"default": 384, "min": 64, "max": 4096, "step": 8, "tooltip": "Minimum resolution dimension"}), + "opt_dim": ("INT", {"default": 720, "min": 64, "max": 4096, "step": 8, "tooltip": "Optimal resolution dimension (most common)"}), + "max_dim": ("INT", {"default": 1312, "min": 64, "max": 4096, "step": 8, "tooltip": "Maximum resolution dimension"}), + } + } + + RETURN_TYPES = ("RIFE_RTX_ENGINE_CONFIG",) + RETURN_NAMES = ("RIFE_ENGINE_CONFIG",) + FUNCTION = "configure" + CATEGORY = "TensorRT_RTX/RIFE" + DESCRIPTION = "Configure custom resolution dimensions for RIFE TensorRT engine." + + def configure(self, min_dim, opt_dim, max_dim): + """Wraps dimensions into a configuration dictionary. + + Args: + min_dim (int): Minimum resolution. + opt_dim (int): Optimal resolution. + max_dim (int): Maximum resolution. + + Returns: + tuple: Contains the configuration dictionary. + """ + config = {"min_d": min_dim, "opt_d": opt_dim, "max_d": max_dim,} + return (config,) + + +class RIFERTXEngineLoader: + """ComfyUI Node for loading or building RIFE TensorRT engines. + + Checks if a .trt engine matching the requested configuration exists. + If not, it compiles the ONNX model into a TensorRT engine. + """ + @classmethod + def INPUT_TYPES(cls): + onnx_names = [x for x in folder_paths.get_filename_list("onnx_rife") if "rife" in x.lower()] + return { + "required": { "onnx_name": (onnx_names,)}, + "optional": { "config": ("RIFE_RTX_ENGINE_CONFIG", {"tooltip": "Options for building the TensorRT engine"}),} + } + + RETURN_NAMES = ("RIFE_ENGINE",) + RETURN_TYPES = ("RIFE_RTX_ENGINE",) + CATEGORY = "TensorRT_RTX/RIFE" + DESCRIPTION = "Load RIFE tensorrt_rtx models, they will be built automatically if not found." + FUNCTION = "load_rife_tensorrt_rtx_model" + + def load_rife_tensorrt_rtx_model(self, onnx_name, config=None): + """Loads or builds the RIFE TensorRT engine. + + Args: + onnx_name (str): Name of the source ONNX file. + config (dict, optional): Resolution profile. Defaults to RIFERTXEngineOptions. + + Returns: + tuple: The loaded Engine object. + """ + if config is None: + config = RIFERTXEngineOptions() + + 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"{onnx_name}") + + logging.info(f"Using resolution profile '{profile_name}': min={config.min_dim}, opt={config.opt_dim}, max={config.max_dim}") + + # Build tensorrt model path with detailed naming (includes profile) + engine_channel = 3 + e_min, e_opt, e_max = config.min_d, config.opt_d, config.max_d + tensorrt_model_path = os.path.join(tensorrt_models_dir, f"{onnx_name}_1x{engine_channel}x{e_min}x{min_}_1x{engine_channel}x{e_opt}x{e_opt}_1x{engine_channel}x{e_max}x{e_max}_{trt.__version__}.trt") + + 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=[ + { + "img0": [(1, engine_channel, e_min, e_min), (1, engine_channel, e_opt, e_opt), (1, engine_channel, e_max, e_max)], + "img1": [(1, engine_channel, e_min, e_min), (1, engine_channel, e_opt, e_opt), (1, engine_channel, e_max, e_max)], + } + ], + ) + + if result != 0: + raise ValueError(f"Failed to build the engine") + e = time.time() + logging.info(f"Time taken to build: {(e-s)} seconds") + + logging.info(f"Loading TensorRT_RTX engine: {tensorrt_model_path}") + mm.soft_empty_cache() + engine.load() + + return (engine,) + +class RIFERTXFrameInterpolation: + """ComfyUI Node for performing VFI inference. + + Uses the loaded TensorRT engine to interpolate frames in a video sequence. + """ + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "frames": ("IMAGE", {"tooltip": "Input frames for video frame interpolation"}), + "engine": ("RIFE_RTX_ENGINE", {"tooltip": "Tensorrt_RTX 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"}), + }, + } + + RETURN_TYPES = ("IMAGE", ) + FUNCTION = "vfi" + CATEGORY = "TensorRT_RTX/RIFE" + OUTPUT_NODE=True + + def vfi( + self, + frames, + engine, + clear_cache_after_n_frames=100, + multiplier=2, + use_cuda_graph=True, + keep_model_loaded=False, + ): + """Interpolates frames using the TensorRT engine. + + Args: + frames (torch.Tensor): Input image batch. + engine (Engine): The active TensorRT engine. + clear_cache_after_n_frames (int): Frequency of CUDA cache clearing. + multiplier (int): Target frame rate multiplier. + use_cuda_graph (bool): Whether to use CUDA graph optimization. + keep_model_loaded (bool): If True, keeps engine in VRAM after use. + + Returns: + tuple: Interpolated image batch. + """ + B, H, W, C = frames.shape + shape_dict = { + "img0": {"shape": (1, 3, H, W)}, + "img1": {"shape": (1, 3, H, W)}, + "output": {"shape": (1, 3, H, W)}, + } + + cudaStream = cuda.Stream() + + # 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(mm.get_torch_device()) + output = engine.infer({"img0": frame_0, "img1": frame_1, "timestep": timestep_t}, cudaStream, use_cuda_graph) + result = output['output'] + + return result + + result = generate_frames_rife(frames, clear_cache_after_n_frames, multiplier, return_middle_frame) + out = postprocess_frames(result) + + if not keep_model_loaded: + engine.reset() + + return (out,) + + +NODE_CLASS_MAPPINGS = { + "RIFERTXFrameInterpolation": RIFERTXFrameInterpolation, + "RIFERTXEngineLoader": RIFERTXEngineLoader, + "RIFERTXEngineConfig": RIFERTXEngineConfig, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "RIFERTXFrameInterpolation": "RIFE RTX Frame Interpolation", + "RIFERTXEngineLoader": "RIFE RTX Engine Loader", + "RIFERTXEngineConfig": "RIFE RTX Engine Config", +} \ No newline at end of file diff --git a/rife_arch.py b/rife_arch.py new file mode 100644 index 0000000..45a2f10 --- /dev/null +++ b/rife_arch.py @@ -0,0 +1,537 @@ +# https://github.com/Fannovel16/ComfyUI-Frame-Interpolation/tree/main/vfi_models/rife + +""" +26-Dez-21 +https://github.com/hzwer/Practical-RIFE +https://github.com/hzwer/Practical-RIFE/blob/main/model/warplayer.py +https://github.com/HolyWu/vs-rife/blob/master/vsrife/__init__.py +""" +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.optim import AdamW +import torch +import torch.nn.functional as F +import torch.nn as nn +import torch.optim as optim +import warnings + +device = "cuda:0" +backwarp_tenGrid = {} + + +class ResConv(nn.Module): + def __init__(self, c, dilation=1): + super(ResConv, self).__init__() + self.conv = nn.Conv2d(c, c, 3, 1, dilation, dilation=dilation, groups=1) + self.beta = nn.Parameter(torch.ones((1, c, 1, 1)), requires_grad=True) + self.relu = nn.LeakyReLU(0.2, True) + + def forward(self, x): + return self.relu(self.conv(x) * self.beta + x) + + +def warp(tenInput, tenFlow): + k = (str(tenFlow.device), str(tenFlow.size())) + if k not in backwarp_tenGrid: + tenHorizontal = ( + torch.linspace(-1.0, 1.0, tenFlow.shape[3], device=device) + .view(1, 1, 1, tenFlow.shape[3]) + .expand(tenFlow.shape[0], -1, tenFlow.shape[2], -1) + ) + tenVertical = ( + torch.linspace(-1.0, 1.0, tenFlow.shape[2], device=device) + .view(1, 1, tenFlow.shape[2], 1) + .expand(tenFlow.shape[0], -1, -1, tenFlow.shape[3]) + ) + backwarp_tenGrid[k] = torch.cat([tenHorizontal, tenVertical], 1).to(device) + + tenFlow = torch.cat( + [ + tenFlow[:, 0:1, :, :] / ((tenInput.shape[3] - 1.0) / 2.0), + tenFlow[:, 1:2, :, :] / ((tenInput.shape[2] - 1.0) / 2.0), + ], + 1, + ) + + g = (backwarp_tenGrid[k] + tenFlow).permute(0, 2, 3, 1) + + if tenInput.type() == "torch.cuda.HalfTensor": + g = g.half() + + return torch.nn.functional.grid_sample( + input=tenInput, + grid=g, + mode="bilinear", + padding_mode="border", + align_corners=True, + ) + + +def conv( + in_planes, + out_planes, + kernel_size=3, + stride=1, + padding=1, + dilation=1, + arch_ver="4.0", +): + if arch_ver == "4.0": + return nn.Sequential( + nn.Conv2d( + in_planes, + out_planes, + kernel_size=kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + bias=True, + ), + nn.PReLU(out_planes), + ) + if arch_ver in ["4.2", "4.3", "4.5", "4.6", "4.7", "4.10"]: + return nn.Sequential( + nn.Conv2d( + in_planes, + out_planes, + kernel_size=kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + bias=True, + ), + nn.LeakyReLU(0.2, True), + ) + + +def conv_woact(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1): + return nn.Sequential( + nn.Conv2d( + in_planes, + out_planes, + kernel_size=kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + bias=True, + ), + ) + + +def conv_woact(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1): + return nn.Sequential( + nn.Conv2d( + in_planes, + out_planes, + kernel_size=kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + bias=True, + ) + ) + + +def deconv(in_planes, out_planes, kernel_size=4, stride=2, padding=1, arch_ver="4.0"): + if arch_ver == "4.0": + return nn.Sequential( + torch.nn.ConvTranspose2d( + in_channels=in_planes, + out_channels=out_planes, + kernel_size=4, + stride=2, + padding=1, + bias=True, + ), + nn.PReLU(out_planes), + ) + if arch_ver in ["4.2", "4.3", "4.5", "4.6", "4.7", "4.10"]: + return nn.Sequential( + torch.nn.ConvTranspose2d( + in_channels=in_planes, + out_channels=out_planes, + kernel_size=4, + stride=2, + padding=1, + bias=True, + ), + nn.LeakyReLU(0.2, True), + ) + + +class Conv2(nn.Module): + def __init__(self, in_planes, out_planes, stride=2, arch_ver="4.0"): + super(Conv2, self).__init__() + self.conv1 = conv(in_planes, out_planes, 3, stride, 1, arch_ver=arch_ver) + self.conv2 = conv(out_planes, out_planes, 3, 1, 1, arch_ver=arch_ver) + + def forward(self, x): + x = self.conv1(x) + x = self.conv2(x) + return x + + +class IFBlock(nn.Module): + def __init__(self, in_planes, c=64, arch_ver="4.0"): + super(IFBlock, self).__init__() + self.arch_ver = arch_ver + self.conv0 = nn.Sequential( + conv(in_planes, c // 2, 3, 2, 1, arch_ver=arch_ver), + conv(c // 2, c, 3, 2, 1, arch_ver=arch_ver), + ) + self.arch_ver = arch_ver + + if arch_ver in ["4.0", "4.2", "4.3"]: + self.convblock = nn.Sequential( + conv(c, c, arch_ver=arch_ver), + conv(c, c, arch_ver=arch_ver), + conv(c, c, arch_ver=arch_ver), + conv(c, c, arch_ver=arch_ver), + conv(c, c, arch_ver=arch_ver), + conv(c, c, arch_ver=arch_ver), + conv(c, c, arch_ver=arch_ver), + conv(c, c, arch_ver=arch_ver), + ) + self.lastconv = nn.ConvTranspose2d(c, 5, 4, 2, 1) + + if arch_ver in ["4.5", "4.6", "4.7", "4.10"]: + self.convblock = nn.Sequential( + ResConv(c), + ResConv(c), + ResConv(c), + ResConv(c), + ResConv(c), + ResConv(c), + ResConv(c), + ResConv(c), + ) + if arch_ver == "4.5": + self.lastconv = nn.Sequential( + nn.ConvTranspose2d(c, 4 * 5, 4, 2, 1), nn.PixelShuffle(2) + ) + if arch_ver in ["4.6", "4.7", "4.10"]: + self.lastconv = nn.Sequential( + nn.ConvTranspose2d(c, 4 * 6, 4, 2, 1), nn.PixelShuffle(2) + ) + + def forward(self, x, flow=None, scale=1): + x = F.interpolate( + x, scale_factor=1.0 / scale, mode="bilinear", align_corners=False + ) + if flow is not None: + flow = ( + F.interpolate( + flow, scale_factor=1.0 / scale, mode="bilinear", align_corners=False + ) + * 1.0 + / scale + ) + x = torch.cat((x, flow), 1) + feat = self.conv0(x) + if self.arch_ver == "4.0": + feat = self.convblock(feat) + feat + if self.arch_ver in ["4.2", "4.3", "4.5", "4.6", "4.7", "4.10"]: + feat = self.convblock(feat) + + tmp = self.lastconv(feat) + if self.arch_ver in ["4.0", "4.2", "4.3"]: + tmp = F.interpolate( + tmp, scale_factor=scale * 2, mode="bilinear", align_corners=False + ) + flow = tmp[:, :4] * scale * 2 + if self.arch_ver in ["4.5", "4.6", "4.7", "4.10"]: + tmp = F.interpolate( + tmp, scale_factor=scale, mode="bilinear", align_corners=False + ) + flow = tmp[:, :4] * scale + mask = tmp[:, 4:5] + return flow, mask + + +class Contextnet(nn.Module): + def __init__(self, arch_ver="4.0"): + super(Contextnet, self).__init__() + c = 16 + self.conv1 = Conv2(3, c, arch_ver=arch_ver) + self.conv2 = Conv2(c, 2 * c, arch_ver=arch_ver) + self.conv3 = Conv2(2 * c, 4 * c, arch_ver=arch_ver) + self.conv4 = Conv2(4 * c, 8 * c, arch_ver=arch_ver) + + def forward(self, x, flow): + x = self.conv1(x) + flow = ( + F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False) + * 0.5 + ) + f1 = warp(x, flow) + x = self.conv2(x) + flow = ( + F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False) + * 0.5 + ) + f2 = warp(x, flow) + x = self.conv3(x) + flow = ( + F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False) + * 0.5 + ) + f3 = warp(x, flow) + x = self.conv4(x) + flow = ( + F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False) + * 0.5 + ) + f4 = warp(x, flow) + return [f1, f2, f3, f4] + + +class Unet(nn.Module): + def __init__(self, arch_ver="4.0"): + super(Unet, self).__init__() + c = 16 + self.down0 = Conv2(17, 2 * c, arch_ver=arch_ver) + self.down1 = Conv2(4 * c, 4 * c, arch_ver=arch_ver) + self.down2 = Conv2(8 * c, 8 * c, arch_ver=arch_ver) + self.down3 = Conv2(16 * c, 16 * c, arch_ver=arch_ver) + self.up0 = deconv(32 * c, 8 * c, arch_ver=arch_ver) + self.up1 = deconv(16 * c, 4 * c, arch_ver=arch_ver) + self.up2 = deconv(8 * c, 2 * c, arch_ver=arch_ver) + self.up3 = deconv(4 * c, c, arch_ver=arch_ver) + self.conv = nn.Conv2d(c, 3, 3, 1, 1) + + def forward(self, img0, img1, warped_img0, warped_img1, mask, flow, c0, c1): + s0 = self.down0( + torch.cat((img0, img1, warped_img0, warped_img1, mask, flow), 1) + ) + s1 = self.down1(torch.cat((s0, c0[0], c1[0]), 1)) + s2 = self.down2(torch.cat((s1, c0[1], c1[1]), 1)) + s3 = self.down3(torch.cat((s2, c0[2], c1[2]), 1)) + x = self.up0(torch.cat((s3, c0[3], c1[3]), 1)) + x = self.up1(torch.cat((x, s2), 1)) + x = self.up2(torch.cat((x, s1), 1)) + x = self.up3(torch.cat((x, s0), 1)) + x = self.conv(x) + return torch.sigmoid(x) + + +""" +currently supports 4.0-4.12 + +4.0: 4.0, 4.1 +4.2: 4.2 +4.3: 4.3, 4.4 +4.5: 4.5 +4.6: 4.6 +4.7: 4.7, 4.8, 4.9 +4.10: 4.10 4.11 4.12 +""" + + +class IFNet(nn.Module): + def __init__(self, arch_ver="4.0"): + super(IFNet, self).__init__() + self.arch_ver = arch_ver + if arch_ver in ["4.0", "4.2", "4.3", "4.5", "4.6"]: + self.block0 = IFBlock(7, c=192, arch_ver=arch_ver) + self.block1 = IFBlock(8 + 4, c=128, arch_ver=arch_ver) + self.block2 = IFBlock(8 + 4, c=96, arch_ver=arch_ver) + self.block3 = IFBlock(8 + 4, c=64, arch_ver=arch_ver) + if arch_ver in ["4.7"]: + self.block0 = IFBlock(7 + 8, c=192, arch_ver=arch_ver) + self.block1 = IFBlock(8 + 4 + 8, c=128, arch_ver=arch_ver) + self.block2 = IFBlock(8 + 4 + 8, c=96, arch_ver=arch_ver) + self.block3 = IFBlock(8 + 4 + 8, c=64, arch_ver=arch_ver) + self.encode = nn.Sequential( + nn.Conv2d(3, 16, 3, 2, 1), nn.ConvTranspose2d(16, 4, 4, 2, 1) + ) + if arch_ver in ["4.10"]: + self.block0 = IFBlock(7 + 16, c=192) + self.block1 = IFBlock(8 + 4 + 16, c=128) + self.block2 = IFBlock(8 + 4 + 16, c=96) + self.block3 = IFBlock(8 + 4 + 16, c=64) + self.encode = nn.Sequential( + nn.Conv2d(3, 32, 3, 2, 1), + nn.LeakyReLU(0.2, True), + nn.Conv2d(32, 32, 3, 1, 1), + nn.LeakyReLU(0.2, True), + nn.Conv2d(32, 32, 3, 1, 1), + nn.LeakyReLU(0.2, True), + nn.ConvTranspose2d(32, 8, 4, 2, 1), + ) + + if arch_ver in ["4.0", "4.2", "4.3"]: + self.contextnet = Contextnet(arch_ver=arch_ver) + self.unet = Unet(arch_ver=arch_ver) + self.arch_ver = arch_ver + + def forward( + self, + img0, + img1, + timestep, + scale_list=[8,4,2,1], + training=True, + fastmode=True, + ensemble=False, + return_flow=False, + ): + img0 = torch.clamp(img0, 0, 1) + img1 = torch.clamp(img1, 0, 1) + + scale_list = [8,4,2,1] + + training=True + fastmode=True + ensemble=True + return_flow=False + + n, c, h, w = img0.shape + ph = ((h - 1) // 64 + 1) * 64 + pw = ((w - 1) // 64 + 1) * 64 + padding = (0, pw - w, 0, ph - h) + img0 = F.pad(img0, padding) + img1 = F.pad(img1, padding) + x = torch.cat((img0, img1), 1) + + if training == False: + channel = x.shape[1] // 2 + img0 = x[:, :channel] + img1 = x[:, channel:] + if not torch.is_tensor(timestep): + timestep = (x[:, :1].clone() * 0 + 1) * timestep + else: + timestep = timestep.repeat(1, 1, img0.shape[2], img0.shape[3]) + + timestep = timestep.to(device) + flow_list = [] + merged = [] + mask_list = [] + + if self.arch_ver in ["4.7", "4.10"]: + f0 = self.encode(img0[:, :3]) + f1 = self.encode(img1[:, :3]) + + + warped_img0 = img0 + warped_img1 = img1 + flow = None + mask = None + block = [self.block0, self.block1, self.block2, self.block3] + + for i in range(4): + if flow is None: + # 4.7+ + if self.arch_ver in ["4.7", "4.10"]: + flow, mask = block[i]( + torch.cat((img0[:, :3], img1[:, :3], f0, f1, timestep), 1), + None, + scale=scale_list[i], + ) + + if ensemble: + f_, m_ = block[i]( + torch.cat( + (img1[:, :3], img0[:, :3], f1, f0, 1 - timestep), 1 + ), + None, + scale=scale_list[i], + ) + flow = (flow + torch.cat((f_[:, 2:4], f_[:, :2]), 1)) / 2 + mask = (mask + (-m_)) / 2 + + else: + # 4.7+ + if self.arch_ver in ["4.7", "4.10"]: + fd, m0 = block[i]( + torch.cat( + ( + warped_img0[:, :3], + warped_img1[:, :3], + warp(f0, flow[:, :2]), + warp(f1, flow[:, 2:4]), + timestep, + mask, + ), + 1, + ), + flow, + scale=scale_list[i], + ) + flow = flow + fd + + # 4.0-4.6 ensemble + if ensemble and self.arch_ver in [ + "4.0", + "4.2", + "4.3", + "4.5", + "4.6", + ]: + f1, m1 = block[i]( + torch.cat( + ( + warped_img1[:, :3], + warped_img0[:, :3], + 1 - timestep, + -mask, + ), + 1, + ), + torch.cat((flow[:, 2:4], flow[:, :2]), 1), + scale=scale_list[i], + ) + f0 = (f0 + torch.cat((f1[:, 2:4], f1[:, :2]), 1)) / 2 + m0 = (m0 + (-m1)) / 2 + + # 4.7+ ensemble + if ensemble and self.arch_ver in ["4.7", "4.10"]: + wf0 = warp(f0, flow[:, :2]) + wf1 = warp(f1, flow[:, 2:4]) + + f_, m_ = block[i]( + torch.cat( + ( + warped_img1[:, :3], + warped_img0[:, :3], + wf1, + wf0, + 1 - timestep, + -mask, + ), + 1, + ), + torch.cat((flow[:, 2:4], flow[:, :2]), 1), + scale=scale_list[i], + ) + fd = (fd + torch.cat((f_[:, 2:4], f_[:, :2]), 1)) / 2 + mask = (m0 + (-m_)) / 2 + + if self.arch_ver in ["4.0", "4.2", "4.3", "4.5", "4.6"]: + flow = flow + f0 + mask = mask + m0 + + if not ensemble and self.arch_ver in ["4.7", "4.10"]: + mask = m0 + + mask_list.append(mask) + flow_list.append(flow) + warped_img0 = warp(img0, flow[:, :2]) + warped_img1 = warp(img1, flow[:, 2:4]) + merged.append((warped_img0, warped_img1)) + + if self.arch_ver in ["4.0", "4.1", "4.2", "4.3", "4.4", "4.5", "4.6"]: + mask_list[3] = torch.sigmoid(mask_list[3]) + merged[3] = merged[3][0] * mask_list[3] + merged[3][1] * (1 - mask_list[3]) + + if self.arch_ver in ["4.7", "4.10"]: + mask = torch.sigmoid(mask) + merged[3] = warped_img0 * mask + warped_img1 * (1 - mask) + + if not fastmode and self.arch_ver in ["4.0", "4.2", "4.3"]: + c0 = self.contextnet(img0, flow[:, :2]) + c1 = self.contextnet(img1, flow[:, 2:4]) + 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/scripts/export_onnx.py b/scripts/export_onnx.py new file mode 100644 index 0000000..96b77fc --- /dev/null +++ b/scripts/export_onnx.py @@ -0,0 +1,142 @@ +# https://github.com/Fannovel16/ComfyUI-Frame-Interpolation/tree/main/vfi_utils.py +# https://github.com/yester31/TensorRT_Examples/blob/main/timm_to_trt_python1/onnx_export.py + +import torch +import pathlib +import traceback +import os +from urllib.parse import urlparse +from torch.hub import download_url_to_file, get_dir +from rife_arch import IFNet +import onnx +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", + "rife45.pth": "4.5", + "rife46.pth": "4.6", + "rife47.pth": "4.7", + "rife48.pth": "4.7", + "rife49.pth": "4.7", + "sudo_rife4_269.662_testV1_scale1.pth": "4.0" +} +BASE_MODEL_DOWNLOAD_URLS = [ + "https://github.com/styler00dollar/VSGAN-tensorrt-docker/releases/download/models/", + "https://github.com/Fannovel16/ComfyUI-Frame-Interpolation/releases/download/models/", + "https://github.com/dajes/frame-interpolation-pytorch/releases/download/v1.0.0/" +] +TORCH_DEVICE = "cuda:0" + +def load_file_from_url(url, model_dir=None, progress=True, file_name=None): + """Load file form http url, will download models if necessary. + + Ref:https://github.com/1adrianb/face-alignment/blob/master/face_alignment/utils.py + + Args: + url (str): URL to be downloaded. + model_dir (str): The path to save the downloaded model. Should be a full path. If None, use pytorch hub_dir. + Default: None. + progress (bool): Whether to show the download progress. Default: True. + file_name (str): The downloaded file name. If None, use the file name in the url. Default: None. + + Returns: + str: The path to the downloaded file. + """ + if model_dir is None: # use the pytorch hub_dir + hub_dir = get_dir() + model_dir = os.path.join(hub_dir, 'checkpoints') + + os.makedirs(model_dir, exist_ok=True) + + parts = urlparse(url) + file_name = os.path.basename(parts.path) + if file_name is not None: + file_name = file_name + cached_file = os.path.abspath(os.path.join(model_dir, file_name)) + if not os.path.exists(cached_file): + print(f'Downloading: "{url}" to {cached_file}\n') + download_url_to_file(url, cached_file, hash_prefix=None, progress=progress) + return cached_file + +def load_file_from_github_release(model_type, ckpt_name): + os.makedirs("models", exist_ok=True) + error_strs = [] + for i, base_model_download_url in enumerate(BASE_MODEL_DOWNLOAD_URLS): + try: + return load_file_from_url(base_model_download_url + ckpt_name, "models") + except Exception: + traceback_str = traceback.format_exc() + if i < len(BASE_MODEL_DOWNLOAD_URLS) - 1: + print("Failed! Trying another endpoint.") + error_strs.append(f"Error when downloading from: {base_model_download_url + ckpt_name}\n\n{traceback_str}") + + error_str = '\n\n'.join(error_strs) + raise Exception(f"Tried all GitHub base urls to download {ckpt_name} but no suceess. Below is the error log:\n\n{error_str}") + +def export_onnx(ckpt_name, ensemble, scale_factor): + print(f"PyTorch version: {torch.__version__}") + print(f"ONNX version: {onnx.__version__}") + device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + print(f"Using device: {device}") + + model_path = load_file_from_github_release(MODEL_TYPE, ckpt_name) + arch_ver = CKPT_NAME_VER_DICT[ckpt_name] + interpolation_model = IFNet(arch_ver=arch_ver) + interpolation_model.load_state_dict(torch.load(model_path)) + interpolation_model.eval().to(TORCH_DEVICE) + + # # dummy data + img0 = torch.randn(1, 3, 512, 512).to(TORCH_DEVICE) + img1 = torch.randn(1, 3, 512, 512).to(TORCH_DEVICE) + timestep = torch.tensor([0.5], dtype=torch.float32).to(TORCH_DEVICE) + + # result = (interpolation_model(img0, img1, timestep)) + # print(result) + + onnx_save_name = f"{ckpt_name.split('.')[0]}_ensemble_{ensemble}_scale_{scale_factor}.onnx" + onnx_save_path = os.path.join("models",onnx_save_name) + + dynamic_axes = { + 'img0': {2: 'height', 3: 'width'}, + '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 + onnx_save_path, + export_params=True, + opset_version=19, + do_constant_folding=True, + input_names=['img0', 'img1', 'timestep'], + output_names=['output'], + dynamic_axes=dynamic_axes + ) + print(f"ONNX model exported to: {onnx_save_path}") + + # Verify the exported ONNX model + 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") + print("=> ONNX simplify start!") + sim_onnx_model, check = simplify(onnx_model) # convert(simplify) + onnx.save(sim_onnx_model, sim_model_path) + print("=> ONNX simplify done!") + + sim_model = onnx.load(sim_model_path) + onnx.checker.check_model(sim_model) + print("=> ONNX Model exported at ", sim_model_path) + print("=> sim ONNX Model check done!") + + +export_onnx(ckpt_name="rife47.pth", ensemble=True, scale_factor=1) \ No newline at end of file diff --git a/scripts/export_trt.py b/scripts/export_trt.py new file mode 100644 index 0000000..74267d4 --- /dev/null +++ b/scripts/export_trt.py @@ -0,0 +1,44 @@ +import sys +import os +sys.path.append(os.path.dirname(os.path.abspath(__file__))) +import torch +import time +from trt_utilities import Engine +import tensorrt + +print("Tensorrt version: ", tensorrt.__version__) + +def export_trt(trt_path=None, onnx_path=None, use_fp16=True): + if trt_path is None: + trt_path = input( + "Enter the path to save the TensorRT engine (e.g ./realesrgan.engine): ") + if onnx_path is None: + onnx_path = input( + "Enter the path to the ONNX model (e.g ./realesrgan.onnx): ") + + engine = Engine(trt_path) + + torch.cuda.empty_cache() + + s = time.time() + ret = engine.build( + onnx_path, + use_fp16, + enable_preview=True, + input_profile=[ + # any sizes from 256x256 to 3840x3840, batch size 1 + { + "img0": [(1, 3, 256, 256), (1, 3, 512, 512), (1, 3, 3840, 3840)], + "img1": [(1, 3, 256, 256), (1, 3, 512, 512), (1, 3, 3840, 3840)], + }, + ], + ) + e = time.time() + print(f"Time taken to build: {(e-s)} seconds") + print(f"Tensorrt engine saved at: {trt_path}") + + return ret + + +export_trt(trt_path="./models/rife49_ensemble_True_scale_1_sim.engine", + onnx_path="./models/rife49_ensemble_True_scale_1_sim.onnx", use_fp16=True) diff --git a/trt_utilities.py b/trt_utilities.py new file mode 100644 index 0000000..143d8e7 --- /dev/null +++ b/trt_utilities.py @@ -0,0 +1,455 @@ +"""TensorRT Acceleration Utilities for RIFE. + +This module provides the low-level interface for NVIDIA TensorRT, handling +engine compilation, timing caches, memory allocation, and optimized inference +using CUDA Graphs. +""" + +import copy +import io +import os + +import cuda.bindings.runtime as cudart +import onnx +import numpy as np +import tensorrt_rtx as trt +import torch + +from collections import OrderedDict +from logging import error, warning +from onnxconverter_common import float16 +from onnxsim import simplify + +# Use polygraphy TensorRTX backend +os.environ["POLYGRAPHY_USE_TENSORRT_RTX"] = "1" + +from polygraphy.backend.common import bytes_from_path +from polygraphy import util +from polygraphy.backend.trt import ModifyNetworkOutputs, Profile +from polygraphy.backend.trt import ( + engine_from_bytes, + engine_from_network, + network_from_onnx_path, + network_from_onnx_bytes, + save_engine, +) +from polygraphy.logger import G_LOGGER + +from torch.cuda import nvtx +from tqdm import tqdm + +G_LOGGER.module_severity = G_LOGGER.ERROR + +# Map of numpy dtype -> torch dtype +numpy_to_torch_dtype_dict = { + np.uint8: torch.uint8, + np.int8: torch.int8, + np.int16: torch.int16, + np.int32: torch.int32, + np.int64: torch.int64, + np.float16: torch.float16, + np.float32: torch.float32, + np.float64: torch.float64, + np.complex64: torch.complex64, + np.complex128: torch.complex128, +} +if np.version.full_version >= "1.24.0": + numpy_to_torch_dtype_dict[np.bool_] = torch.bool +else: + numpy_to_torch_dtype_dict[np.bool] = torch.bool + +# Map of torch dtype -> numpy dtype +torch_to_numpy_dtype_dict = { + value: key for (key, value) in numpy_to_torch_dtype_dict.items() +} + +# Path for your persistent hardware profile +node_dir = os.path.dirname(os.path.realpath(__file__)) +MASTER_CACHE_FILE = f"{node_dir}\\RIFE_RTX_ENGINE.cache" + +def get_timing_cache(config): + """Retrieves or creates a TensorRT timing cache. + + Args: + config (trt.IBuilderConfig): The builder configuration object. + + Returns: + trt.ITimingCache: The loaded or new timing cache. + """ + if os.path.exists(MASTER_CACHE_FILE): + with open(MASTER_CACHE_FILE, "rb") as f: + cache_data = f.read() + return config.create_timing_cache(cache_data) + return config.create_timing_cache(b"") + +def save_timing_cache(timing_cache): + """Saves the current timing cache to disk for future builds. + + Args: + timing_cache (trt.ITimingCache): The cache to serialize. + """ + with open(MASTER_CACHE_FILE, "wb") as f: + f.write(timing_cache.serialize()) + +def CUASSERT(cuda_ret): + """Asserts that a CUDA call was successful. + https://github.com/Jeff-LiangF/streamv2v/blob/18c1a3bd56ff348d54a3300605936980bb13b03c/src/streamv2v/acceleration/tensorrt/utilities.py + + Args: + cuda_ret (tuple): The return value from a cuda-python call. + + Returns: + Any: The result of the CUDA call if successful. + + Raises: + RuntimeError: If the CUDA call returned an error. + """ + err = cuda_ret[0] + if err != cudart.cudaError_t.cudaSuccess: + raise RuntimeError( + f"CUDA ERROR: {err}, error code reference: https://nvidia.github.io/cuda-python/module/cudart.html#cuda.cudart.cudaError_t" + ) + if len(cuda_ret) > 1: + return cuda_ret[1] + return None + +class TQDMProgressMonitor(trt.IProgressMonitor): + """Progress monitor for TensorRT engine builds using TQDM bars.""" + + def __init__(self): + trt.IProgressMonitor.__init__(self) + self._active_phases = {} + self._step_result = True + self.max_indent = 5 + + def phase_start(self, phase_name, parent_phase, num_steps): + """Called when a build phase starts.""" + leave = False + try: + if parent_phase is not None: + nbIndents = ( + self._active_phases.get(parent_phase, {}).get( + "nbIndents", self.max_indent + ) + + 1 + ) + if nbIndents >= self.max_indent: + return + else: + nbIndents = 0 + leave = True + self._active_phases[phase_name] = { + "tq": tqdm( + total=num_steps, desc=phase_name, leave=leave, position=nbIndents + ), + "nbIndents": nbIndents, + "parent_phase": parent_phase, + } + except KeyboardInterrupt: + # The phase_start callback cannot directly cancel the build, so request the cancellation from within step_complete. + _step_result = False + + def phase_finish(self, phase_name): + """Called when a build phase finishes.""" + try: + if phase_name in self._active_phases.keys(): + self._active_phases[phase_name]["tq"].update( + self._active_phases[phase_name]["tq"].total + - self._active_phases[phase_name]["tq"].n + ) + + parent_phase = self._active_phases[phase_name].get("parent_phase", None) + while parent_phase is not None: + self._active_phases[parent_phase]["tq"].refresh() + parent_phase = self._active_phases[parent_phase].get( + "parent_phase", None + ) + if ( + self._active_phases[phase_name]["parent_phase"] + in self._active_phases.keys() + ): + self._active_phases[ + self._active_phases[phase_name]["parent_phase"] + ]["tq"].refresh() + del self._active_phases[phase_name] + pass + except KeyboardInterrupt: + _step_result = False + + def step_complete(self, phase_name, step): + """Called when a step within a phase completes.""" + try: + if phase_name in self._active_phases.keys(): + self._active_phases[phase_name]["tq"].update( + step - self._active_phases[phase_name]["tq"].n + ) + return self._step_result + except KeyboardInterrupt: + # There is no need to propagate this exception to TensorRT. We can simply cancel the build. + return False + +class Engine: + """Wrapper for TensorRT engine management. + + Handles loading, building, and executing inference with optional + CUDA Graph support. + + Attributes: + engine_path (str): File path to the .trt engine. + """ + + def __init__( + self, + engine_path, + ): + self.engine_path = engine_path + self.engine = None + self.context = None + self.buffers = OrderedDict() + self.tensors = OrderedDict() + self.cuda_graph_instance = None # cuda graph + self.graph = None + + def __del__(self): + """Ensures CUDA resources are freed on deletion.""" + 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): + """Cleans up engine and CUDA resources. + + Args: + engine_path (str, optional): New path to set for the engine. + """ + 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() + self.inputs = {} + self.outputs = {} + + def build( + self, + onnx_path, + input_profile=None, + enable_refit=False, + enable_preview=False, + enable_all_tactics=False, + timing_cache=None, + update_output_names=None, + ): + """Compiles an ONNX model into a TensorRT engine. + + Args: + onnx_path (str): Path to source ONNX file. + input_profile (list): List of dicts defining min/opt/max shapes. + **kwargs: Configuration flags (enable_refit, etc). + + Returns: + int: 0 if success, 1 if failure. + """ + p = [Profile()] + if input_profile: + p = [Profile() for i in range(len(input_profile))] + for _p, i_profile in zip(p, input_profile): + for name, dims in i_profile.items(): + assert len(dims) == 3 + _p.add(name, min=dims[0], opt=dims[1], max=dims[2]) + + config_kwargs = {} + if not enable_all_tactics: + config_kwargs["tactic_sources"] = [] + + # Load ONNX file then convert the model to fp16 + # Probably might cause issues if onnx file given is quantized or whatever + onnx_bytes = onnx.load(onnx_path) + onnx_precision = onnx_bytes.graph.input[0].type.tensor_type.elem_type + onnx_simp, check = simplify(float16.convert_float_to_float16(onnx_bytes) if onnx_precision == onnx.TensorProto.FLOAT else onnx_bytes) + + if not check: + error(f"Simplified ONNX model: {onnx_path} could not be validated") + return 1 + + # Keep simplified model in memory to not overwrite the onnx file in models directory + # We dont want to be destructive + onnx_buffer = io.BytesIO() + onnx_buffer.write(onnx_simp.SerializeToString()) + onnx_buffer.seek(0) + onnx_buffer.name = "onnx_simp_mem_rife.onnx" + + #network = network_from_onnx_path(onnx_path, flags=[trt.OnnxParserFlag.NATIVE_INSTANCENORM]) + network = network_from_onnx_bytes(onnx_buffer.getvalue(), flags=[trt.OnnxParserFlag.NATIVE_INSTANCENORM]) + + if update_output_names: + print(f"Updating network outputs to {update_output_names}") + network = ModifyNetworkOutputs(network, update_output_names) + + builder = network[0] + config = builder.create_builder_config() + config.progress_monitor = TQDMProgressMonitor() + + #config.set_flag(trt.BuilderFlag.FP16) if fp16 else None + config.set_flag(trt.BuilderFlag.REFIT) if enable_refit else None + + # Limit workspace (e.g., to 2GB or 4GB depending on your GPU) + # 1 << 30 is 1GB + #config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 4 * (1 << 30)) + config.builder_optimization_level = 5 + + # Use Master Cache + timing_cache = get_timing_cache(config) + config.set_timing_cache(timing_cache, False) + + profiles = copy.deepcopy(p) + for profile in profiles: + # Last profile is used for set_calibration_profile. + calib_profile = profile.fill_defaults(network[1]).to_trt( + builder, network[1] + ) + config.add_optimization_profile(calib_profile) + + result = 0 + + try: + engine = engine_from_network(network, config) + save_engine(engine, path=self.engine_path) + except Exception as e: + error(f"Failed to build engine: {e}") + result = 1 + + # Update Master Cache with new findings + save_timing_cache(timing_cache) + print(f"Updated Master Cache.") + + return result + + def load(self): + """Loads the serialized engine from disk.""" + self.engine = engine_from_bytes(bytes_from_path(self.engine_path)) + + def activate(self, reuse_device_memory=None): + """Initializes the execution context.""" + 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 + else: + self.context = self.engine.create_execution_context() + + def allocate_buffers(self, shape_dict=None, device="cuda"): + """Allocates GPU memory for input and output tensors. + + Args: + shape_dict (dict, optional): Specific shapes for tensors. + device (str): Device to allocate on (default "cuda"). + """ + 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) + binding = self.engine[idx] + if shape_dict and binding in shape_dict: + shape = shape_dict[binding]["shape"] + else: + shape = self.context.get_tensor_shape(name) + + dtype = trt.nptype(self.engine.get_tensor_dtype(name)) + if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT: + self.context.set_input_shape(name, shape) + tensor = torch.empty( + tuple(shape), dtype=numpy_to_torch_dtype_dict[dtype] + ).to(device=device) + self.tensors[binding] = tensor + nvtx.range_pop() + + def infer(self, feed_dict, stream, use_cuda_graph=False): + """Performs inference on the provided data. + + Args: + feed_dict (dict): Input tensors. + stream (cuda.Stream): CUDA stream for execution. + use_cuda_graph (bool): If True, uses/captures CUDA graphs. + + Returns: + OrderedDict: The output tensors. + """ + for name, buf in feed_dict.items(): + self.tensors[name].copy_(buf) + + for name, tensor in self.tensors.items(): + self.context.set_tensor_address(name, tensor.data_ptr()) + + if use_cuda_graph: + if self.cuda_graph_instance is not None: + CUASSERT(cudart.cudaGraphLaunch(self.cuda_graph_instance, stream.ptr)) + CUASSERT(cudart.cudaStreamSynchronize(stream.ptr)) + else: + # do inference before CUDA graph capture + noerror = self.context.execute_async_v3(stream.ptr) + if not noerror: + raise ValueError("ERROR: inference failed.") + # capture cuda graph + CUASSERT( + cudart.cudaStreamBeginCapture(stream.ptr, cudart.cudaStreamCaptureMode.cudaStreamCaptureModeGlobal) + ) + self.context.execute_async_v3(stream.ptr) + self.graph = CUASSERT(cudart.cudaStreamEndCapture(stream.ptr)) + self.cuda_graph_instance = CUASSERT(cudart.cudaGraphInstantiate(self.graph, 0)) + else: + noerror = self.context.execute_async_v3(stream.ptr) + if not noerror: + raise ValueError("ERROR: inference failed.") + + return self.tensors + diff --git a/vfi_utilities.py b/vfi_utilities.py new file mode 100644 index 0000000..ebc094e --- /dev/null +++ b/vfi_utilities.py @@ -0,0 +1,112 @@ +"""Utility functions for RIFE Video Frame Interpolation (VFI). + +This module contains preprocessing, postprocessing, and the core interpolation +loop logic used to generate intermediate frames between a sequence of images. + +https://github.com/Fannovel16/ComfyUI-Frame-Interpolation/blob/main/vfi_utils.py +""" + +import einops +import os +import torch +import typing + +import numpy as np + +from comfy.model_management import soft_empty_cache, get_torch_device +from comfy.utils import ProgressBar + +DEVICE = get_torch_device() + +def preprocess_frames(frames: torch.Tensor) -> torch.Tensor: + """Prepares ComfyUI image batches for model inference. + + Converts images from (Batch, Height, Width, Channels) to + (Batch, Channels, Height, Width) and ensures only RGB channels are present. + + Args: + frames (torch.Tensor): Input tensor in BHWC format. + + Returns: + torch.Tensor: Preprocessed tensor in BCHW format. + """ + return einops.rearrange(frames[..., :3], "n h w c -> n c h w") + +def postprocess_frames(frames: torch.Tensor) -> torch.Tensor: + """Converts model output back to ComfyUI standard format. + + Converts images from (Batch, Channels, Height, Width) back to + (Batch, Height, Width, Channels) and moves them to CPU memory. + + Args: + frames (torch.Tensor): Output tensor from the model in BCHW format. + + Returns: + torch.Tensor: Postprocessed tensor in BHWC format on CPU. + """ + return einops.rearrange(frames, "n c h w -> n h w c")[..., :3].cpu() + +def generate_frames_rife( + frames, + clear_cache_after_n_frames, + multiplier, + return_middle_frame_function + ): + """Core interpolation loop for generating intermediate frames. + + Iterates through a sequence of frames and calls the provided inference + function to generate 'n' intermediate frames based on the multiplier. + + Args: + frames (torch.Tensor): The preprocessed frame sequence (BCHW). + clear_cache_after_n_frames (int): Interval for VRAM cache clearing. + multiplier (int): The factor by which to increase the frame count. + return_middle_frame_function (Callable): Function that takes two + frames and a timestep and returns the interpolated frame. + + Returns: + torch.Tensor: The complete sequence including original and + interpolated frames. + """ + output_frames = torch.zeros(multiplier*frames.shape[0], *frames.shape[1:], device="cpu") + out_len = 0 + cache_counter = 0 + + pbar = ProgressBar(len(frames)) + + for frame_itr in range(len(frames) - 1): + + frame_0 = frames[frame_itr:frame_itr+1] + frame_1 = frames[frame_itr+1:frame_itr+2] + + # Store the current original frame + output_frames[out_len] = frame_0 + out_len += 1 + + # Generate intermediate frames + for middle_i in range(1, multiplier): + timestep = middle_i/multiplier + middle_frame = return_middle_frame_function(frame_0, frame_1, timestep).detach().cpu() + + # Copy middle frames to output + output_frames[out_len] = middle_frame + out_len +=1 + + # VRAM Management + cache_counter += 1 + if cache_counter >= clear_cache_after_n_frames: + soft_empty_cache() + cache_counter = 0 + + pbar.update(1) + + # Append final frame + output_frames[out_len] = frames[-1:] + logger(f"done! - {(len(frames) -1) * (multiplier-1)} new frames generated at resolution: {output_frames[0].shape}") + out_len += 1 + + # clear cache for courtesy + soft_empty_cache() + + # Slice to actual length in case of rounding or logic offsets + return output_frames[:out_len]