From 5cff7f7dc7cd85052dbe50d53d8296d9c145f642 Mon Sep 17 00:00:00 2001 From: silveroxides Date: Fri, 9 Jan 2026 22:44:36 +0100 Subject: [PATCH] feat: CUDA 13 default, auto-detection, and resolution profiles Requirements updates: - Default to CUDA 13 with TensorRT 10.14.1.48 (requirements.txt) - Add cuda-toolkit version constraints (>=13.0.0,<13.1.0) - Rename requirements_cu13.txt -> requirements_cu12.txt for legacy support - Add cuda-toolkit constraint for CUDA 12 (>=12.8.0,<13.0.0) CUDA toolkit detection: - Add _setup_cuda_dll_path() to auto-detect CUDA installation - Check CUDA_PATH, CUDA_HOME, and default Windows install paths - Handle CUDA 13+ bin/x64 subdirectory for cudart64.dll - Display download link if CUDA not found Engine building updates: - Add resolution profiles (small: 480-896px, medium: 720-1280px) - Add 'custom' profile with CustomResolutionConfig node - Include profile name in engine filename for separate caching - Reduce default VRAM usage from ~15GB to ~2-3GB --- __init__.py | 127 +++++++++++++++++++++++++++++++++++++++--- load_rife_config.json | 9 +++ readme.md | 42 +++++++++----- requirements.txt | 10 ++-- requirements_cu12.txt | 12 ++++ requirements_cu13.txt | 9 --- 6 files changed, 174 insertions(+), 35 deletions(-) create mode 100644 requirements_cu12.txt delete mode 100644 requirements_cu13.txt diff --git a/__init__.py b/__init__.py index 6ee3281..9c890e5 100644 --- a/__init__.py +++ b/__init__.py @@ -1,11 +1,56 @@ import torch import os +import sys +from pathlib import Path 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 + +# Auto-detect CUDA toolkit and add DLL path before importing polygraphy +def _setup_cuda_dll_path(): + """Auto-detect CUDA toolkit and add cudart64 DLL path on Windows.""" + 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") + +_setup_cuda_dll_path() + from polygraphy import cuda import comfy.model_management as mm import tensorrt @@ -13,10 +58,11 @@ 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 +# Default resolution profiles (fallback if config is missing) +DEFAULT_RESOLUTION_PROFILES = { + "small": {"min": 480, "opt": 512, "max": 896}, + "medium": {"min": 720, "opt": 1024, "max": 1280} +} # Logger for this module rife_logger = ColoredLogger("ComfyUI-Rife-Tensorrt") @@ -58,6 +104,34 @@ def load_node_config(config_filename="load_rife_config.json"): # Load the configuration once when the module is imported LOAD_RIFE_NODE_CONFIG = load_node_config() + +class CustomResolutionConfig: + """Node to configure custom resolution dimensions for TensorRT engine building.""" + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "min_dim": ("INT", {"default": 480, "min": 64, "max": 4096, "step": 8, "tooltip": "Minimum resolution dimension"}), + "opt_dim": ("INT", {"default": 512, "min": 64, "max": 4096, "step": 8, "tooltip": "Optimal resolution dimension (most common)"}), + "max_dim": ("INT", {"default": 896, "min": 64, "max": 4096, "step": 8, "tooltip": "Maximum resolution dimension"}), + } + } + + RETURN_TYPES = ("RIFE_RESOLUTION_CONFIG",) + RETURN_NAMES = ("resolution_config",) + FUNCTION = "configure" + CATEGORY = "tensorrt" + DESCRIPTION = "Configure custom resolution dimensions for RIFE TensorRT engine." + + def configure(self, min_dim, opt_dim, max_dim): + config = { + "min": min_dim, + "opt": opt_dim, + "max": max_dim, + } + return (config,) + + class AutoLoadRifeTensorrtModel: @classmethod def INPUT_TYPES(cls): @@ -74,10 +148,23 @@ class AutoLoadRifeTensorrtModel: precision_default = precision_config.get("default", "fp16") precision_tooltip = precision_config.get("tooltip", "Select precision.") + # Resolution profile configuration + profile_config = LOAD_RIFE_NODE_CONFIG.get("resolution_profile", {}) + profile_options = profile_config.get("options", ["small", "medium"]) + # Ensure 'custom' is always available + if "custom" not in profile_options: + profile_options = profile_options + ["custom"] + profile_default = profile_config.get("default", "small") + profile_tooltip = profile_config.get("tooltip", "Resolution range for TensorRT engine. Use 'custom' with the INT inputs below.") + return { "required": { "model": (model_options, {"default": model_default, "tooltip": model_tooltip}), "precision": (precision_options, {"default": precision_default, "tooltip": precision_tooltip}), + "resolution_profile": (profile_options, {"default": profile_default, "tooltip": profile_tooltip}), + }, + "optional": { + "custom_config": ("RIFE_RESOLUTION_CONFIG", {"tooltip": "Custom resolution config (used when profile='custom')"}), } } @@ -87,7 +174,7 @@ class AutoLoadRifeTensorrtModel: 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): + def load_rife_tensorrt_model(self, model, precision, resolution_profile, custom_config=None): tensorrt_models_dir = os.path.join(folder_paths.models_dir, "tensorrt", "rife") onnx_models_dir = os.path.join(folder_paths.models_dir, "onnx") @@ -96,12 +183,32 @@ class AutoLoadRifeTensorrtModel: onnx_model_path = os.path.join(onnx_models_dir, f"{model}.onnx") - # Build tensorrt model path with detailed naming + # Get resolution dimensions based on profile + if resolution_profile == "custom": + if custom_config is None: + rife_logger.warning("Custom profile selected but no custom_config provided. Using defaults.") + dim_min, dim_opt, dim_max = 480, 512, 896 + else: + dim_min = custom_config.get("min", 480) + dim_opt = custom_config.get("opt", 512) + dim_max = custom_config.get("max", 896) + # Use dimensions in profile name for custom engines + profile_name = f"custom_{dim_min}_{dim_opt}_{dim_max}" + else: + profiles = LOAD_RIFE_NODE_CONFIG.get("resolution_profiles", DEFAULT_RESOLUTION_PROFILES) + profile = profiles.get(resolution_profile, DEFAULT_RESOLUTION_PROFILES["small"]) + dim_min = profile.get("min", 480) + dim_opt = profile.get("opt", 512) + dim_max = profile.get("max", 896) + profile_name = resolution_profile + rife_logger.info(f"Using resolution profile '{profile_name}': min={dim_min}, opt={dim_opt}, max={dim_max}") + + # Build tensorrt model path with detailed naming (includes profile) 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") + engine_min_h, engine_opt_h, engine_max_h = dim_min, dim_opt, dim_max + engine_min_w, engine_opt_w, engine_max_w = dim_min, dim_opt, dim_max + tensorrt_model_path = os.path.join(tensorrt_models_dir, f"{model}_{precision}_{profile_name}_{engine_min_batch}x{engine_channel}x{engine_min_h}x{engine_min_w}_{engine_opt_batch}x{engine_channel}x{engine_opt_h}x{engine_opt_w}_{engine_max_batch}x{engine_channel}x{engine_max_h}x{engine_max_w}_{tensorrt.__version__}.trt") if not os.path.exists(tensorrt_model_path): if not os.path.exists(onnx_model_path): @@ -204,11 +311,13 @@ class AutoRifeTensorrt: NODE_CLASS_MAPPINGS = { "AutoRifeTensorrt": AutoRifeTensorrt, "AutoLoadRifeTensorrtModel": AutoLoadRifeTensorrtModel, + "CustomResolutionConfig": CustomResolutionConfig, } NODE_DISPLAY_NAME_MAPPINGS = { "AutoRifeTensorrt": "Auto RIFE TensorRT", "AutoLoadRifeTensorrtModel": "(Down)load RIFE TensorRT Model", + "CustomResolutionConfig": "RIFE Custom Resolution Config", } __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] diff --git a/load_rife_config.json b/load_rife_config.json index 3c01bd0..1701d6c 100644 --- a/load_rife_config.json +++ b/load_rife_config.json @@ -12,5 +12,14 @@ "options": ["fp16", "fp32"], "default": "fp16", "tooltip": "Precision to build the tensorrt engines. Loaded from config." + }, + "resolution_profile": { + "options": ["small", "medium"], + "default": "small", + "tooltip": "Resolution range for TensorRT engine. Small: 480-896px, Medium: 720-1280px. Smaller ranges use less VRAM." + }, + "resolution_profiles": { + "small": { "min": 480, "opt": 512, "max": 896 }, + "medium": { "min": 720, "opt": 1024, "max": 1280 } } } \ No newline at end of file diff --git a/readme.md b/readme.md index c93c16d..8709c42 100644 --- a/readme.md +++ b/readme.md @@ -2,9 +2,9 @@ # ComfyUI Rife 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) +[![python](https://img.shields.io/badge/python-3.12-green)](https://www.python.org/downloads/) +[![cuda](https://img.shields.io/badge/cuda-13.0-green)](https://developer.nvidia.com/cuda-13-0-2-download-archive) +[![trt](https://img.shields.io/badge/TRT-10.14.1.48-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,21 +40,32 @@ cd ./ComfyUI-Rife-Tensorrt pip install -r requirements.txt ``` -### ⚠️ Installation Issues (CUDA Versions) -This node defaults to **CUDA 12** libraries (`requirements.txt`). -If you are using **CUDA 13**, or if you experience dependency errors, please install the specific requirements file for your system: +### ⚠️ CUDA Version Selection -**For CUDA 13:** +This node defaults to **CUDA 13** (RTX 50 series, driver 580+). + +**For CUDA 12 (RTX 30/40 series):** ```bash -pip install -r requirements_cu13.txt +pip install -r requirements_cu12.txt ``` -**For CUDA 12 (Default):** +**For CUDA 13 (Default):** ```bash pip install -r requirements.txt ``` -This ensures the correct `tensorrt-libs` and `tensorrt-bindings` are installed to match your system. +### 📦 CUDA Toolkit Required + +The node automatically detects your CUDA installation via `CUDA_PATH` or `CUDA_HOME` environment variables. + +If CUDA is not detected, download from: https://developer.nvidia.com/cuda-13-0-2-download-archive + +### 🎯 Resolution Profiles + +The node supports resolution profiles to optimize VRAM usage: +- **small**: 480-896px (recommended for most video generation) +- **medium**: 720-1280px (for higher resolution videos) +- **custom**: Connect a "RIFE Custom Resolution Config" node for manual control The following RIFE models are supported and will be automatically downloaded and built: - **rife49_ensemble_True_scale_1_sim** (default) - Latest and most accurate @@ -68,21 +79,26 @@ Models are automatically downloaded from [HuggingFace](https://huggingface.co/yu 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) + - Select resolution profile (small, medium, or custom) - 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 -- 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) +- Windows 11, CUDA 13.0, TensorRT 10.14.1.48, Python 3.12, RTX 5070 Ti +- WSL Ubuntu 24.04.03 LTS, CUDA 12.9, TensorRT 10.13.3.9, Python 3.12.11, RTX 5080 ## 🚨 Updates +### January 2026 +- **CUDA 13 Default**: Updated to CUDA 13.0 and TensorRT 10.14.1.48 +- **Auto CUDA Detection**: Automatically finds CUDA toolkit and DLL paths +- **Resolution Profiles**: Added small/medium/custom profiles to reduce VRAM usage + ### 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 diff --git a/requirements.txt b/requirements.txt index 3c743e0..4fc9c05 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,7 +1,9 @@ -# Default to CUDA 12 as it is the standard for TRT 10.13 -tensorrt-cu12==10.13.3.9 -tensorrt-cu12-libs==10.13.3.9 -tensorrt-cu12-bindings==10.13.3.9 +# Default: CUDA 13 (RTX 50 series, CUDA 13.0+) +# For CUDA 12, use requirements_cu12.txt instead +tensorrt_cu13==10.14.1.48 +tensorrt_cu13_bindings==10.14.1.48 +tensorrt_cu13_libs==10.14.1.48 +cuda-toolkit>=13.0.0,<13.1.0 einops colored polygraphy diff --git a/requirements_cu12.txt b/requirements_cu12.txt new file mode 100644 index 0000000..cbf070d --- /dev/null +++ b/requirements_cu12.txt @@ -0,0 +1,12 @@ +# Legacy: CUDA 12 (RTX 30/40 series, CUDA 12.x) +# For CUDA 13, use requirements.txt instead +tensorrt-cu12==10.13.3.9 +tensorrt-cu12-libs==10.13.3.9 +tensorrt-cu12-bindings==10.13.3.9 +cuda-toolkit>=12.8.0,<13.0.0 +einops +colored +polygraphy +cuda-python +requests +tqdm diff --git a/requirements_cu13.txt b/requirements_cu13.txt deleted file mode 100644 index b23cf54..0000000 --- a/requirements_cu13.txt +++ /dev/null @@ -1,9 +0,0 @@ -tensorrt-cu13==10.13.3.9 -tensorrt-cu13-libs==10.13.3.9 -tensorrt-cu13-bindings==10.13.3.9 -einops -colored -polygraphy -cuda-python -requests -tqdm