103 lines
3.4 KiB
Python
Executable File
103 lines
3.4 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""
|
|
Auto-install script for ComfyUI RIFE TensorRT
|
|
Detects CUDA version and installs appropriate requirements
|
|
"""
|
|
|
|
import subprocess
|
|
import sys
|
|
import os
|
|
import re
|
|
|
|
def run_command(cmd):
|
|
"""Run command and return output"""
|
|
try:
|
|
result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
|
|
return result.stdout.strip(), result.returncode
|
|
except Exception as e:
|
|
return "", 1
|
|
|
|
def detect_cuda_version():
|
|
"""Detect CUDA version"""
|
|
print("🔍 Detecting CUDA version...")
|
|
|
|
# Try nvcc command
|
|
stdout, returncode = run_command("nvcc --version")
|
|
if returncode == 0:
|
|
match = re.search(r"release (\d+\.\d+)", stdout)
|
|
if match:
|
|
version = match.group(1)
|
|
print(f"✅ Found CUDA version: {version}")
|
|
return version
|
|
|
|
# Try CUDA_PATH
|
|
cuda_path = os.environ.get("CUDA_PATH")
|
|
if cuda_path:
|
|
nvcc_path = os.path.join(cuda_path, "bin", "nvcc")
|
|
if os.path.exists(nvcc_path):
|
|
stdout, returncode = run_command(f"{nvcc_path} --version")
|
|
if returncode == 0:
|
|
match = re.search(r"release (\d+\.\d+)", stdout)
|
|
if match:
|
|
version = match.group(1)
|
|
print(f"✅ Found CUDA version via CUDA_PATH: {version}")
|
|
return version
|
|
|
|
# Try CUDA_HOME
|
|
cuda_home = os.environ.get("CUDA_HOME")
|
|
if cuda_home:
|
|
nvcc_path = os.path.join(cuda_home, "bin", "nvcc")
|
|
if os.path.exists(nvcc_path):
|
|
stdout, returncode = run_command(f"{nvcc_path} --version")
|
|
if returncode == 0:
|
|
match = re.search(r"release (\d+\.\d+)", stdout)
|
|
if match:
|
|
version = match.group(1)
|
|
print(f"✅ Found CUDA version via CUDA_HOME: {version}")
|
|
return version
|
|
|
|
print("⚠️ Could not detect CUDA version automatically")
|
|
print("Please ensure CUDA is installed and nvcc is in your PATH")
|
|
print("Or set CUDA_PATH or CUDA_HOME environment variables")
|
|
return None
|
|
|
|
def install_requirements(cuda_version):
|
|
"""Install appropriate requirements based on CUDA version"""
|
|
major_version = int(cuda_version.split('.')[0])
|
|
|
|
print(f"📦 Installing requirements for CUDA {major_version}...")
|
|
|
|
if major_version == 13:
|
|
print("🚀 Using CUDA 13 requirements (RTX 50 series)")
|
|
requirements_file = "requirements.txt"
|
|
elif major_version == 12:
|
|
print("🔧 Using CUDA 12 requirements (RTX 30/40 series)")
|
|
requirements_file = "requirements_cu12.txt"
|
|
else:
|
|
print(f"❌ Unsupported CUDA version: {cuda_version}")
|
|
print("Supported versions: CUDA 12.x, CUDA 13.x")
|
|
return False
|
|
|
|
# Install requirements
|
|
returncode = subprocess.run([sys.executable, "-m", "pip", "install", "-r", requirements_file]).returncode
|
|
|
|
if returncode == 0:
|
|
print("✅ Installation completed successfully!")
|
|
print("🎯 You can now use the ComfyUI RIFE TensorRT node")
|
|
return True
|
|
else:
|
|
print("❌ Installation failed!")
|
|
return False
|
|
|
|
def main():
|
|
"""Main installation function"""
|
|
cuda_version = detect_cuda_version()
|
|
if not cuda_version:
|
|
sys.exit(1)
|
|
|
|
success = install_requirements(cuda_version)
|
|
sys.exit(0 if success else 1)
|
|
|
|
if __name__ == "__main__":
|
|
main()
|