From 3cdf12ab663dd4796694a5a11b2243969f6e2aa6 Mon Sep 17 00:00:00 2001 From: huchukato Date: Sat, 21 Feb 2026 09:14:05 +0100 Subject: [PATCH] Implement fully automatic CUDA detection and TensorRT installation on node load --- __init__.py | 104 +++++++++++++++++++++++++++++++++++++++++++++++ readme.md | 38 ++++++++++------- requirements.txt | 16 ++++---- 3 files changed, 137 insertions(+), 21 deletions(-) diff --git a/__init__.py b/__init__.py index bc93887..ac33db9 100644 --- a/__init__.py +++ b/__init__.py @@ -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.""" diff --git a/readme.md b/readme.md index 7f77d0f..85822b3 100644 --- a/readme.md +++ b/readme.md @@ -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 diff --git a/requirements.txt b/requirements.txt index 3aff4a0..a838510 100644 --- a/requirements.txt +++ b/requirements.txt @@ -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 \ No newline at end of file +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) \ No newline at end of file