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
This commit is contained in:
+118
-9
@@ -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']
|
||||
|
||||
@@ -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 }
|
||||
}
|
||||
}
|
||||
@@ -2,9 +2,9 @@
|
||||
|
||||
# ComfyUI Rife TensorRT ⚡
|
||||
|
||||
[](https://www.python.org/downloads/release/python-31211/)
|
||||
[](https://developer.nvidia.com/cuda-downloads)
|
||||
[](https://developer.nvidia.com/tensorrt)
|
||||
[](https://www.python.org/downloads/)
|
||||
[](https://developer.nvidia.com/cuda-13-0-2-download-archive)
|
||||
[](https://developer.nvidia.com/tensorrt)
|
||||
[](https://creativecommons.org/licenses/by-nc-sa/4.0/deed.en)
|
||||
|
||||

|
||||
@@ -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
|
||||
|
||||
+6
-4
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user