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

Implement fully automatic CUDA detection and TensorRT installation on node load

Browse Source
This commit is contained in:
huchukato
2026-02-21 09:14:05 +01:00
parent 6eb23462d9
commit 3cdf12ab66
3 changed files with 137 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+104
View File
@@ -1,6 +1,8 @@
import torch
import os
import sys
import subprocess
import re
from pathlib import Path
from comfy.model_management import get_torch_device
from .vfi_utilities import preprocess_frames, postprocess_frames, generate_frames_rife, logger
@@ -9,6 +11,108 @@ from .utilities import download_file, ColoredLogger
import folder_paths
import time
# Auto-detect CUDA and install appropriate TensorRT packages
def _auto_install_tensorrt():
"""Auto-detect CUDA version and install appropriate TensorRT packages if needed"""
try:
# Check if TensorRT is already installed
try:
import tensorrt
print("✅ TensorRT already installed")
return True
except ImportError:
print("🔍 TensorRT not found, detecting CUDA version...")
# Detect CUDA version
cuda_version = None
# Try nvcc command
try:
result = subprocess.run("nvcc --version", shell=True, capture_output=True, text=True)
if result.returncode == 0:
match = re.search(r"release (\d+\.\d+)", result.stdout)
if match:
cuda_version = match.group(1)
print(f"✅ Detected CUDA version: {cuda_version}")
except:
pass
# Try CUDA_PATH
if not cuda_version and os.environ.get("CUDA_PATH"):
nvcc_path = os.path.join(os.environ["CUDA_PATH"], "bin", "nvcc")
if os.path.exists(nvcc_path):
try:
result = subprocess.run(f"{nvcc_path} --version", shell=True, capture_output=True, text=True)
if result.returncode == 0:
match = re.search(r"release (\d+\.\d+)", result.stdout)
if match:
cuda_version = match.group(1)
print(f"✅ Detected CUDA via CUDA_PATH: {cuda_version}")
except:
pass
# Try CUDA_HOME
if not cuda_version and os.environ.get("CUDA_HOME"):
nvcc_path = os.path.join(os.environ["CUDA_HOME"], "bin", "nvcc")
if os.path.exists(nvcc_path):
try:
result = subprocess.run(f"{nvcc_path} --version", shell=True, capture_output=True, text=True)
if result.returncode == 0:
match = re.search(r"release (\d+\.\d+)", result.stdout)
if match:
cuda_version = match.group(1)
print(f"✅ Detected CUDA via CUDA_HOME: {cuda_version}")
except:
pass
if not cuda_version:
print("⚠️ Could not detect CUDA version automatically")
print("Please run 'python install.py' manually to install TensorRT")
return False
# Install appropriate TensorRT packages
major_version = int(cuda_version.split('.')[0])
if major_version == 13:
print("🚀 Installing CUDA 13 TensorRT packages (RTX 50 series)")
packages = [
"tensorrt_cu13==10.15.1.29",
"tensorrt_cu13_bindings==10.15.1.29",
"tensorrt_cu13_libs==10.15.1.29",
"cuda-toolkit>=13.0.0,<13.1.0"
]
elif major_version == 12:
print("🔧 Installing CUDA 12 TensorRT packages (RTX 30/40 series)")
packages = [
"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"
]
else:
print(f"❌ Unsupported CUDA version: {cuda_version}")
return False
# Install packages
for package in packages:
print(f"Installing {package}...")
result = subprocess.run([sys.executable, "-m", "pip", "install", package], capture_output=True)
if result.returncode != 0:
print(f"❌ Failed to install {package}")
print(f"Error: {result.stderr.decode()}")
return False
print("✅ TensorRT installation completed successfully!")
return True
except Exception as e:
print(f"❌ Auto-installation failed: {e}")
print("Please run 'python install.py' manually to install TensorRT")
return False
# Run auto-install on module import
_auto_install_tensorrt()
# 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."""
readme.md
+24 -14
View File
@@ -39,31 +39,41 @@ git clone https://github.com/huchukato/ComfyUI-RIFE-TensorRT-Auto
cd ./ComfyUI-RIFE-TensorRT-Auto
```
### 🎯 Auto-Install (Recommended)
### 🎯 Fully Automatic Installation (Recommended)
The node includes auto-install scripts that detect your CUDA version and install the appropriate requirements:
**This node features fully automatic CUDA detection and TensorRT installation!**
**Linux/macOS:**
When ComfyUI loads the node for the first time, it will:
1. **Auto-detect your CUDA version** (12 or 13)
2. **Install the appropriate TensorRT packages** automatically
3. **Configure everything** for seamless operation
**No manual steps required!** Just clone the repo and restart ComfyUI.
### 📦 Manual Installation Options
If you prefer manual installation or encounter issues:
**Auto-install scripts:**
```bash
# Linux/macOS
./install.sh
```
**Windows:**
```cmd
# Windows
install.bat
```
**Python (cross-platform):**
```bash
# Python (cross-platform)
python install.py
```
The scripts automatically:
- Detect CUDA version via `nvcc`, `CUDA_PATH`, or `CUDA_HOME`
- Install CUDA 13 requirements for RTX 50 series
- Install CUDA 12 requirements for RTX 30/40 series
**Manual requirements:**
```bash
# For CUDA 13 (RTX 50 series)
pip install -r requirements.txt
### 📦 Manual Installation
# For CUDA 12 (RTX 30/40 series)
pip install -r requirements_cu12.txt
```
### ⚠️ CUDA Version Selection
requirements.txt
+9 -7
View File
@@ -1,12 +1,14 @@
# Default: CUDA 13 (RTX 50 series, CUDA 13.0+)
# For CUDA 12, use requirements_cu12.txt instead
tensorrt_cu13==10.15.1.29
tensorrt_cu13_bindings==10.15.1.29
tensorrt_cu13_libs==10.15.1.29
cuda-toolkit>=13.0.0,<13.1.0
# Base dependencies - TensorRT will be auto-installed based on CUDA detection
# The node will automatically detect CUDA 12/13 and install appropriate TensorRT packages
# Base dependencies that work with any CUDA version
einops
colored
polygraphy
cuda-python
requests
tqdm
tqdm
# Note: TensorRT packages will be auto-installed when the node loads
# CUDA 13: tensorrt_cu13==10.15.1.29 (RTX 50 series)
# CUDA 12: tensorrt-cu12==10.13.3.9 (RTX 30/40 series)