742 lines
29 KiB
Python
742 lines
29 KiB
Python
"""
|
|
Installation script for ComfyUI-SAM3 optional dependencies.
|
|
Called by ComfyUI Manager during installation/update.
|
|
"""
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import argparse
|
|
|
|
# Global flag for compile-only mode
|
|
COMPILE_ONLY = False
|
|
|
|
|
|
def install_requirements():
|
|
"""
|
|
Install dependencies from requirements.txt.
|
|
"""
|
|
print("[ComfyUI-SAM3] Installing requirements.txt dependencies...")
|
|
|
|
script_dir = os.path.dirname(os.path.abspath(__file__))
|
|
requirements_path = os.path.join(script_dir, "requirements.txt")
|
|
|
|
if not os.path.exists(requirements_path):
|
|
print("[ComfyUI-SAM3] [WARNING] requirements.txt not found, skipping")
|
|
return False
|
|
|
|
try:
|
|
result = subprocess.run(
|
|
[sys.executable, "-m", "pip", "install", "-r", requirements_path],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=600
|
|
)
|
|
|
|
if result.returncode == 0:
|
|
print("[ComfyUI-SAM3] [OK] Requirements installed successfully")
|
|
else:
|
|
print("[ComfyUI-SAM3] [WARNING] Requirements installation had issues")
|
|
if result.stderr:
|
|
print(f"[ComfyUI-SAM3] Error details: {result.stderr[:500]}")
|
|
return False
|
|
|
|
except Exception as e:
|
|
print(f"[ComfyUI-SAM3] [WARNING] Requirements installation error: {e}")
|
|
return False
|
|
|
|
# Install ninja for faster CUDA compilation
|
|
print("[ComfyUI-SAM3] Installing ninja build system...")
|
|
try:
|
|
result = subprocess.run(
|
|
[sys.executable, "-m", "pip", "install", "ninja"],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=300
|
|
)
|
|
if result.returncode == 0:
|
|
print("[ComfyUI-SAM3] [OK] Ninja installed successfully")
|
|
else:
|
|
print("[ComfyUI-SAM3] [WARNING] Ninja installation failed (optional, compilation will be slower)")
|
|
except Exception as e:
|
|
print(f"[ComfyUI-SAM3] [WARNING] Ninja installation error: {e}")
|
|
|
|
return True
|
|
|
|
|
|
def find_cuda_home():
|
|
"""
|
|
Find CUDA installation directory from all possible sources.
|
|
Returns path to CUDA installation or None if not found.
|
|
"""
|
|
import site
|
|
import shutil
|
|
|
|
# Check existing CUDA_HOME environment variable
|
|
if "CUDA_HOME" in os.environ and os.path.exists(os.environ["CUDA_HOME"]):
|
|
cuda_home = os.environ["CUDA_HOME"]
|
|
nvcc_path = os.path.join(cuda_home, "bin", "nvcc")
|
|
if os.path.exists(nvcc_path) or os.path.exists(nvcc_path + ".exe"):
|
|
print(f"[ComfyUI-SAM3] Found CUDA via CUDA_HOME: {cuda_home}")
|
|
return cuda_home
|
|
|
|
# Check if nvcc is in system PATH (e.g., from apt-get install nvidia-cuda-toolkit)
|
|
nvcc_in_path = shutil.which("nvcc")
|
|
if nvcc_in_path:
|
|
# Derive CUDA_HOME from nvcc location
|
|
# /usr/local/cuda-12.4/bin/nvcc -> /usr/local/cuda-12.4 (preferred)
|
|
# /usr/bin/nvcc -> /usr (old CUDA 11.5, only use if no better option)
|
|
nvcc_dir = os.path.dirname(nvcc_in_path) # Get bin directory
|
|
if os.path.basename(nvcc_dir) == "bin":
|
|
cuda_home = os.path.dirname(nvcc_dir) # Get parent directory
|
|
# Only return /usr if it's not the system directory (avoid old CUDA 11.5)
|
|
# Prefer checking explicit CUDA 12.x paths first
|
|
if cuda_home != "/usr" or "cuda-12" in cuda_home or "cuda-11" in nvcc_in_path:
|
|
print(f"[ComfyUI-SAM3] Found CUDA via system PATH: {cuda_home} (nvcc at {nvcc_in_path})")
|
|
return cuda_home
|
|
else:
|
|
# nvcc is at /usr/bin/nvcc (old CUDA 11.5), defer to explicit path checking below
|
|
print(f"[ComfyUI-SAM3] Found nvcc at {nvcc_in_path} but deferring to explicit CUDA 12.x path search...")
|
|
|
|
# Check system CUDA installations (Linux/Mac)
|
|
# Prioritize CUDA 12.x installations
|
|
system_paths = [
|
|
"/usr/local/cuda-12.4",
|
|
"/usr/local/cuda-12.3",
|
|
"/usr/local/cuda-12.2",
|
|
"/usr/local/cuda-12.1",
|
|
"/usr/local/cuda-12.0",
|
|
"/usr/local/cuda",
|
|
"/usr/cuda",
|
|
"/opt/cuda",
|
|
"/usr/local/cuda-12",
|
|
"/usr/local/cuda-11",
|
|
]
|
|
|
|
for cuda_path in system_paths:
|
|
if os.path.exists(cuda_path):
|
|
nvcc_path = os.path.join(cuda_path, "bin", "nvcc")
|
|
if os.path.exists(nvcc_path):
|
|
print(f"[ComfyUI-SAM3] Found system CUDA installation: {cuda_path}")
|
|
return cuda_path
|
|
|
|
# Check Windows common locations
|
|
if sys.platform == "win32":
|
|
program_files = os.environ.get("ProgramFiles", "C:\\Program Files")
|
|
cuda_base = os.path.join(program_files, "NVIDIA GPU Computing Toolkit", "CUDA")
|
|
if os.path.exists(cuda_base):
|
|
# Find latest version
|
|
versions = [d for d in os.listdir(cuda_base) if os.path.isdir(os.path.join(cuda_base, d))]
|
|
if versions:
|
|
versions.sort(reverse=True)
|
|
cuda_path = os.path.join(cuda_base, versions[0])
|
|
nvcc_path = os.path.join(cuda_path, "bin", "nvcc.exe")
|
|
if os.path.exists(nvcc_path):
|
|
print(f"[ComfyUI-SAM3] Found Windows CUDA installation: {cuda_path}")
|
|
return cuda_path
|
|
|
|
# Check pip-installed CUDA in site-packages
|
|
site_packages = site.getsitepackages()
|
|
if isinstance(site_packages, str):
|
|
site_packages = [site_packages]
|
|
|
|
for sp in site_packages:
|
|
# Check nvidia/cuda_nvcc structure
|
|
cuda_nvcc_paths = [
|
|
os.path.join(sp, "nvidia", "cuda_nvcc"),
|
|
os.path.join(sp, "nvidia_cuda_nvcc"),
|
|
os.path.join(sp, "cuda_nvcc"),
|
|
]
|
|
|
|
for cuda_path in cuda_nvcc_paths:
|
|
if os.path.exists(cuda_path):
|
|
# Check if nvcc binary exists
|
|
nvcc_path = os.path.join(cuda_path, "bin", "nvcc")
|
|
if os.path.exists(nvcc_path) or os.path.exists(nvcc_path + ".exe"):
|
|
print(f"[ComfyUI-SAM3] Found pip-installed CUDA: {cuda_path}")
|
|
return cuda_path
|
|
|
|
# Check conda environment (if running in one)
|
|
if "CONDA_PREFIX" in os.environ:
|
|
conda_prefix = os.environ["CONDA_PREFIX"]
|
|
nvcc_path = os.path.join(conda_prefix, "bin", "nvcc")
|
|
if os.path.exists(nvcc_path):
|
|
print(f"[ComfyUI-SAM3] Found conda CUDA installation: {conda_prefix}")
|
|
return conda_prefix
|
|
|
|
return None
|
|
|
|
|
|
def check_nvcc_available():
|
|
"""
|
|
Check if CUDA compiler (nvcc) is available.
|
|
Returns True if nvcc is found, False otherwise.
|
|
"""
|
|
try:
|
|
result = subprocess.run(
|
|
["nvcc", "--version"],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=10
|
|
)
|
|
if result.returncode == 0:
|
|
# Extract CUDA version from nvcc output
|
|
output = result.stdout
|
|
if "release" in output:
|
|
version = output.split('release')[1].split(',')[0].strip()
|
|
print(f"[ComfyUI-SAM3] [OK] CUDA compiler found: release {version}")
|
|
else:
|
|
print(f"[ComfyUI-SAM3] [OK] CUDA compiler found")
|
|
print(f"[ComfyUI-SAM3] nvcc output: {output.strip()}")
|
|
return True
|
|
else:
|
|
print(f"[ComfyUI-SAM3] [WARNING] nvcc command failed with return code {result.returncode}")
|
|
if result.stderr:
|
|
print(f"[ComfyUI-SAM3] Error: {result.stderr[:200]}")
|
|
return False
|
|
except FileNotFoundError:
|
|
print("[ComfyUI-SAM3] [WARNING] nvcc command not found in PATH")
|
|
return False
|
|
except subprocess.TimeoutExpired:
|
|
print("[ComfyUI-SAM3] [WARNING] nvcc command timed out")
|
|
return False
|
|
except Exception as e:
|
|
print(f"[ComfyUI-SAM3] [WARNING] Error checking nvcc: {e}")
|
|
return False
|
|
|
|
|
|
def get_cuda_arch_list():
|
|
"""
|
|
Detect GPU compute capability and return TORCH_CUDA_ARCH_LIST string.
|
|
Returns None if detection fails.
|
|
"""
|
|
try:
|
|
import torch
|
|
if not torch.cuda.is_available():
|
|
print("[ComfyUI-SAM3] [WARNING] CUDA not available in PyTorch")
|
|
return None
|
|
|
|
# Get compute capability of current device
|
|
major, minor = torch.cuda.get_device_capability(0)
|
|
compute_cap = f"{major}.{minor}"
|
|
|
|
# Get device name for informational purposes
|
|
device_name = torch.cuda.get_device_name(0)
|
|
print(f"[ComfyUI-SAM3] Detected GPU: {device_name} (compute capability {compute_cap})")
|
|
|
|
# Common architecture list covering most modern GPUs
|
|
# Include current GPU and common architectures
|
|
arch_list = []
|
|
|
|
# Add current GPU architecture
|
|
arch_list.append(compute_cap)
|
|
|
|
# Add common architectures for compatibility
|
|
common_archs = ["7.5", "8.0", "8.6", "8.9", "9.0"]
|
|
for arch in common_archs:
|
|
if arch not in arch_list:
|
|
arch_list.append(arch)
|
|
|
|
arch_string = " ".join(arch_list)
|
|
print(f"[ComfyUI-SAM3] Using CUDA architectures: {arch_string}")
|
|
|
|
return arch_string
|
|
|
|
except Exception as e:
|
|
print(f"[ComfyUI-SAM3] [WARNING] Could not detect GPU compute capability: {e}")
|
|
# Fallback to common architectures
|
|
return "7.5 8.0 8.6 8.9 9.0"
|
|
|
|
|
|
def setup_cuda_environment():
|
|
"""
|
|
Setup CUDA environment variables (CUDA_HOME and PATH) for compilation.
|
|
Searches for CUDA installation from all possible sources.
|
|
Returns dict with environment variables if CUDA found, None otherwise.
|
|
"""
|
|
# First try to find existing CUDA installation
|
|
cuda_home = find_cuda_home()
|
|
|
|
if cuda_home:
|
|
# Create a copy of current environment
|
|
env = os.environ.copy()
|
|
|
|
# Set CUDA_HOME environment variable
|
|
env["CUDA_HOME"] = cuda_home
|
|
os.environ["CUDA_HOME"] = cuda_home # Also set in current process
|
|
print(f"[ComfyUI-SAM3] Set CUDA_HOME={cuda_home}")
|
|
|
|
# Add CUDA bin directory to PATH
|
|
cuda_bin = os.path.join(cuda_home, "bin")
|
|
if cuda_bin not in env.get("PATH", ""):
|
|
env["PATH"] = cuda_bin + os.pathsep + env.get("PATH", "")
|
|
os.environ["PATH"] = env["PATH"] # Also set in current process
|
|
print(f"[ComfyUI-SAM3] Added to PATH: {cuda_bin}")
|
|
|
|
# Verify nvcc is accessible
|
|
if check_nvcc_available():
|
|
# Print final environment summary for debugging
|
|
print(f"[ComfyUI-SAM3] CUDA environment configured:")
|
|
print(f"[ComfyUI-SAM3] CUDA_HOME: {env.get('CUDA_HOME', 'not set')}")
|
|
print(f"[ComfyUI-SAM3] PATH includes: {cuda_bin}")
|
|
return env
|
|
else:
|
|
print(f"[ComfyUI-SAM3] [WARNING] nvcc not accessible after setting CUDA_HOME={cuda_home}")
|
|
|
|
print("[ComfyUI-SAM3] [WARNING] Could not setup CUDA environment")
|
|
return None
|
|
|
|
|
|
def try_install_cuda_toolkit():
|
|
"""
|
|
Attempt to install CUDA toolkit if nvcc is missing.
|
|
Supports Linux, Windows, and Mac (though Mac doesn't support CUDA anymore).
|
|
Returns True if successful or already installed, False otherwise.
|
|
"""
|
|
print("[ComfyUI-SAM3] Attempting to install CUDA toolkit...")
|
|
|
|
try:
|
|
import torch
|
|
|
|
cuda_version = torch.version.cuda
|
|
if not cuda_version:
|
|
print("[ComfyUI-SAM3] [WARNING] PyTorch not built with CUDA support")
|
|
return False
|
|
|
|
# Extract major version from CUDA version (e.g., "12" from "12.8")
|
|
cuda_major = cuda_version.split('.')[0]
|
|
print(f"[ComfyUI-SAM3] PyTorch CUDA version: {cuda_version} (major: cu{cuda_major})")
|
|
|
|
# Try pip install first (works on all platforms but may be incomplete)
|
|
package_name = f"nvidia-cuda-nvcc-cu{cuda_major}"
|
|
print(f"[ComfyUI-SAM3] Trying to install {package_name} via pip...")
|
|
|
|
result = subprocess.run(
|
|
[sys.executable, "-m", "pip", "install", package_name],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=300
|
|
)
|
|
|
|
if result.returncode == 0:
|
|
print("[ComfyUI-SAM3] [OK] CUDA compiler package installed via pip")
|
|
|
|
# Try to setup environment with newly installed CUDA
|
|
if setup_cuda_environment():
|
|
print("[ComfyUI-SAM3] [OK] CUDA environment configured successfully")
|
|
return True
|
|
else:
|
|
print("[ComfyUI-SAM3] [WARNING] pip package installed but nvcc not accessible")
|
|
print("[ComfyUI-SAM3] Note: pip nvidia-cuda-nvcc packages are often incomplete")
|
|
print("[ComfyUI-SAM3] Trying system package manager as fallback...")
|
|
else:
|
|
print(f"[ComfyUI-SAM3] [WARNING] pip install of {package_name} failed")
|
|
if result.stderr:
|
|
print(f"[ComfyUI-SAM3] Error: {result.stderr[:300]}")
|
|
|
|
# Platform-specific installation fallbacks
|
|
if sys.platform.startswith('linux'):
|
|
return try_install_cuda_linux()
|
|
elif sys.platform == 'win32':
|
|
return try_install_cuda_windows(cuda_version)
|
|
elif sys.platform == 'darwin':
|
|
return try_install_cuda_mac()
|
|
else:
|
|
print(f"[ComfyUI-SAM3] [WARNING] Unsupported platform: {sys.platform}")
|
|
return False
|
|
|
|
except Exception as e:
|
|
print(f"[ComfyUI-SAM3] [WARNING] CUDA toolkit installation error: {e}")
|
|
return False
|
|
|
|
|
|
def try_install_cuda_linux():
|
|
"""Try to install CUDA toolkit on Linux using package managers."""
|
|
print("[ComfyUI-SAM3] Attempting Linux package manager installation...")
|
|
|
|
# Try apt-get (Debian/Ubuntu)
|
|
try:
|
|
print("[ComfyUI-SAM3] Trying apt-get (Debian/Ubuntu)...")
|
|
|
|
# First, ensure build-essential is installed (includes g++ and gcc)
|
|
print("[ComfyUI-SAM3] Ensuring build-essential and g++ are installed...")
|
|
result = subprocess.run(
|
|
["sudo", "apt-get", "install", "-y", "build-essential", "g++"],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=300
|
|
)
|
|
if result.returncode == 0:
|
|
print("[ComfyUI-SAM3] [OK] Build tools installed successfully")
|
|
else:
|
|
print("[ComfyUI-SAM3] [WARNING] Failed to install build-essential/g++")
|
|
|
|
# Update package lists
|
|
result = subprocess.run(
|
|
["sudo", "apt-get", "update"],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=120
|
|
)
|
|
if result.returncode == 0:
|
|
# Install CUDA 12.x toolkit (try 12.4, 12.3, 12.2, etc.)
|
|
cuda_packages = ["cuda-toolkit-12-4", "cuda-toolkit-12-3", "cuda-toolkit-12-2", "cuda-toolkit-12-1", "cuda-toolkit-12-0"]
|
|
|
|
for cuda_pkg in cuda_packages:
|
|
print(f"[ComfyUI-SAM3] Trying to install {cuda_pkg}...")
|
|
result = subprocess.run(
|
|
["sudo", "apt-get", "install", "-y", "--fix-missing", cuda_pkg],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=600
|
|
)
|
|
if result.returncode == 0:
|
|
print(f"[ComfyUI-SAM3] [OK] CUDA compiler installed via apt-get: {cuda_pkg}")
|
|
if setup_cuda_environment():
|
|
return True
|
|
else:
|
|
print(f"[ComfyUI-SAM3] {cuda_pkg} not available, trying next...")
|
|
|
|
except (FileNotFoundError, subprocess.TimeoutExpired) as e:
|
|
print(f"[ComfyUI-SAM3] apt-get not available or failed: {e}")
|
|
|
|
# Try yum (RedHat/CentOS/Fedora)
|
|
try:
|
|
print("[ComfyUI-SAM3] Trying yum (RedHat/CentOS)...")
|
|
result = subprocess.run(
|
|
["sudo", "yum", "install", "-y", "cuda-toolkit"],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=600
|
|
)
|
|
if result.returncode == 0:
|
|
print("[ComfyUI-SAM3] [OK] CUDA compiler installed via yum")
|
|
if setup_cuda_environment():
|
|
return True
|
|
except (FileNotFoundError, subprocess.TimeoutExpired):
|
|
pass
|
|
|
|
# Try dnf (Fedora)
|
|
try:
|
|
print("[ComfyUI-SAM3] Trying dnf (Fedora)...")
|
|
result = subprocess.run(
|
|
["sudo", "dnf", "install", "-y", "cuda-toolkit"],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=600
|
|
)
|
|
if result.returncode == 0:
|
|
print("[ComfyUI-SAM3] [OK] CUDA compiler installed via dnf")
|
|
if setup_cuda_environment():
|
|
return True
|
|
except (FileNotFoundError, subprocess.TimeoutExpired):
|
|
pass
|
|
|
|
print("[ComfyUI-SAM3] [WARNING] Could not install CUDA compiler via package managers")
|
|
print("[ComfyUI-SAM3] Manual installation required:")
|
|
print("[ComfyUI-SAM3] https://developer.nvidia.com/cuda-downloads")
|
|
print("[ComfyUI-SAM3] Select your Linux distribution and follow the instructions")
|
|
return False
|
|
|
|
|
|
def try_install_cuda_windows(cuda_version):
|
|
"""Try to install CUDA toolkit on Windows."""
|
|
print("[ComfyUI-SAM3] Attempting Windows CUDA installation...")
|
|
|
|
# Try chocolatey if available
|
|
try:
|
|
print("[ComfyUI-SAM3] Checking for Chocolatey package manager...")
|
|
result = subprocess.run(
|
|
["choco", "--version"],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=10
|
|
)
|
|
if result.returncode == 0:
|
|
print("[ComfyUI-SAM3] Chocolatey found, attempting installation...")
|
|
result = subprocess.run(
|
|
["choco", "install", "cuda", "-y"],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=600
|
|
)
|
|
if result.returncode == 0:
|
|
print("[ComfyUI-SAM3] [OK] CUDA installed via Chocolatey")
|
|
if setup_cuda_environment():
|
|
return True
|
|
except (FileNotFoundError, subprocess.TimeoutExpired):
|
|
print("[ComfyUI-SAM3] Chocolatey not available")
|
|
|
|
# Windows requires manual installation - provide instructions
|
|
print("[ComfyUI-SAM3] [WARNING] Automatic CUDA installation not available on Windows")
|
|
print("[ComfyUI-SAM3] Manual installation required:")
|
|
print("[ComfyUI-SAM3]")
|
|
print("[ComfyUI-SAM3] 1. Visit: https://developer.nvidia.com/cuda-downloads")
|
|
print("[ComfyUI-SAM3] 2. Select: Windows > x86_64 > your Windows version")
|
|
print("[ComfyUI-SAM3] 3. Download and run the installer")
|
|
print(f"[ComfyUI-SAM3] 4. Make sure to install CUDA {cuda_version} to match your PyTorch version")
|
|
print("[ComfyUI-SAM3] 5. After installation, restart your terminal and run install.py again")
|
|
return False
|
|
|
|
|
|
def try_install_cuda_mac():
|
|
"""Handle Mac CUDA installation (not supported)."""
|
|
print("[ComfyUI-SAM3] [WARNING] CUDA is not supported on macOS")
|
|
print("[ComfyUI-SAM3] Apple Silicon Macs do not support NVIDIA GPUs")
|
|
print("[ComfyUI-SAM3] Intel Macs with NVIDIA GPUs are no longer officially supported by NVIDIA")
|
|
print("[ComfyUI-SAM3]")
|
|
print("[ComfyUI-SAM3] GPU-accelerated video tracking is not available on Mac")
|
|
print("[ComfyUI-SAM3] CPU fallback will be used (slower but functional)")
|
|
return False
|
|
|
|
|
|
def try_install_torch_generic_nms():
|
|
"""
|
|
Attempt to install torch_generic_nms for GPU-accelerated video tracking.
|
|
This provides 5-10x faster NMS compared to CPU fallback.
|
|
"""
|
|
print("[ComfyUI-SAM3] Checking for torch_generic_nms (GPU-accelerated NMS for video)...")
|
|
|
|
# Check if already installed
|
|
try:
|
|
import torch_generic_nms
|
|
print("[ComfyUI-SAM3] [OK] torch_generic_nms already installed")
|
|
return True
|
|
except ImportError:
|
|
pass
|
|
|
|
# Check if CUDA is available in PyTorch
|
|
try:
|
|
import torch
|
|
if not COMPILE_ONLY and not torch.cuda.is_available():
|
|
print("[ComfyUI-SAM3] [INFO] CUDA not available, skipping GPU NMS compilation")
|
|
print("[ComfyUI-SAM3] Video tracking will use CPU fallback (slower but functional)")
|
|
return False
|
|
elif COMPILE_ONLY:
|
|
print("[ComfyUI-SAM3] [INFO] Compile-only mode: checking PyTorch CUDA support...")
|
|
if not torch.version.cuda:
|
|
print("[ComfyUI-SAM3] [ERROR] PyTorch was not built with CUDA support")
|
|
print("[ComfyUI-SAM3] Cannot compile CUDA extensions without CUDA-enabled PyTorch")
|
|
return False
|
|
print(f"[ComfyUI-SAM3] [OK] PyTorch built with CUDA {torch.version.cuda}")
|
|
except ImportError:
|
|
print("[ComfyUI-SAM3] [WARNING] PyTorch not found, cannot install torch_generic_nms")
|
|
return False
|
|
|
|
# Setup CUDA environment (find CUDA from all possible sources)
|
|
cuda_env = setup_cuda_environment()
|
|
if not cuda_env:
|
|
print("[ComfyUI-SAM3] [WARNING] CUDA compiler not found in system")
|
|
print("[ComfyUI-SAM3] Attempting to install CUDA toolkit...")
|
|
|
|
if not try_install_cuda_toolkit():
|
|
print("[ComfyUI-SAM3] [WARNING] Could not install CUDA compiler")
|
|
print("[ComfyUI-SAM3] Video tracking will use CPU NMS fallback")
|
|
print("[ComfyUI-SAM3] To install manually, see: https://developer.nvidia.com/cuda-downloads")
|
|
return False
|
|
|
|
# Try setup again after installation
|
|
cuda_env = setup_cuda_environment()
|
|
if not cuda_env:
|
|
print("[ComfyUI-SAM3] [WARNING] CUDA environment setup failed after toolkit installation")
|
|
return False
|
|
|
|
# Get CUDA architecture list
|
|
cuda_arch_list = get_cuda_arch_list()
|
|
|
|
if cuda_arch_list:
|
|
cuda_env["TORCH_CUDA_ARCH_LIST"] = cuda_arch_list
|
|
os.environ["TORCH_CUDA_ARCH_LIST"] = cuda_arch_list # Also set in current process
|
|
print(f"[ComfyUI-SAM3] Set TORCH_CUDA_ARCH_LIST={cuda_arch_list}")
|
|
|
|
# Set C/C++ compiler paths explicitly (needed for CUDA compilation)
|
|
import shutil
|
|
gcc_path = shutil.which("gcc-11") or shutil.which("gcc")
|
|
gxx_path = shutil.which("g++-11") or shutil.which("g++")
|
|
if gcc_path:
|
|
cuda_env["CC"] = gcc_path
|
|
os.environ["CC"] = gcc_path
|
|
print(f"[ComfyUI-SAM3] Set CC={gcc_path}")
|
|
if gxx_path:
|
|
cuda_env["CXX"] = gxx_path
|
|
os.environ["CXX"] = gxx_path
|
|
print(f"[ComfyUI-SAM3] Set CXX={gxx_path}")
|
|
|
|
# Install torch_generic_nms from GitHub
|
|
print("[ComfyUI-SAM3] Compiling torch_generic_nms from source (may take 1-2 minutes)...")
|
|
|
|
try:
|
|
result = subprocess.run(
|
|
[
|
|
sys.executable, "-m", "pip", "install",
|
|
"--no-build-isolation", # Use current environment's torch
|
|
"git+https://github.com/ronghanghu/torch_generic_nms"
|
|
],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=180, # Increased timeout for compilation
|
|
env=cuda_env # Pass CUDA environment explicitly to subprocess
|
|
)
|
|
|
|
if result.returncode == 0:
|
|
print("[ComfyUI-SAM3] [OK] torch_generic_nms compiled and installed successfully")
|
|
print("[ComfyUI-SAM3] Video tracking will use GPU-accelerated NMS (5-10x faster)")
|
|
return True
|
|
else:
|
|
print("[ComfyUI-SAM3] [WARNING] torch_generic_nms compilation failed")
|
|
print("[ComfyUI-SAM3] This is optional - video tracking will use CPU NMS fallback")
|
|
if result.stderr:
|
|
print(f"[ComfyUI-SAM3] Error details: {result.stderr[:500]}")
|
|
return False
|
|
|
|
except subprocess.TimeoutExpired:
|
|
print("[ComfyUI-SAM3] [WARNING] torch_generic_nms compilation timed out")
|
|
print("[ComfyUI-SAM3] Video tracking will use CPU NMS fallback")
|
|
return False
|
|
except Exception as e:
|
|
print(f"[ComfyUI-SAM3] [WARNING] torch_generic_nms installation error: {e}")
|
|
print("[ComfyUI-SAM3] Continuing without GPU NMS (video tracking will still work)")
|
|
return False
|
|
|
|
|
|
def try_install_cc_torch():
|
|
"""
|
|
Attempt to install cc_torch for GPU-accelerated connected components.
|
|
Used for hole filling in video tracking masks.
|
|
"""
|
|
print("[ComfyUI-SAM3] Checking for cc_torch (GPU-accelerated connected components)...")
|
|
|
|
# Check if already installed
|
|
try:
|
|
import cc_torch
|
|
print("[ComfyUI-SAM3] [OK] cc_torch already installed")
|
|
return True
|
|
except ImportError:
|
|
pass
|
|
|
|
# Check if CUDA is available in PyTorch
|
|
try:
|
|
import torch
|
|
if not COMPILE_ONLY and not torch.cuda.is_available():
|
|
print("[ComfyUI-SAM3] [INFO] CUDA not available, skipping GPU connected components")
|
|
print("[ComfyUI-SAM3] Video tracking will use CPU fallback (slower but functional)")
|
|
return False
|
|
elif COMPILE_ONLY:
|
|
print("[ComfyUI-SAM3] [INFO] Compile-only mode: checking PyTorch CUDA support...")
|
|
if not torch.version.cuda:
|
|
print("[ComfyUI-SAM3] [ERROR] PyTorch was not built with CUDA support")
|
|
print("[ComfyUI-SAM3] Cannot compile CUDA extensions without CUDA-enabled PyTorch")
|
|
return False
|
|
print(f"[ComfyUI-SAM3] [OK] PyTorch built with CUDA {torch.version.cuda}")
|
|
except ImportError:
|
|
print("[ComfyUI-SAM3] [WARNING] PyTorch not found, cannot install cc_torch")
|
|
return False
|
|
|
|
# Setup CUDA environment (find CUDA from all possible sources)
|
|
cuda_env = setup_cuda_environment()
|
|
if not cuda_env:
|
|
print("[ComfyUI-SAM3] [WARNING] CUDA compiler not found in system")
|
|
print("[ComfyUI-SAM3] Attempting to install CUDA toolkit...")
|
|
|
|
if not try_install_cuda_toolkit():
|
|
print("[ComfyUI-SAM3] [WARNING] Could not install CUDA compiler")
|
|
print("[ComfyUI-SAM3] Video tracking will use CPU connected components fallback")
|
|
return False
|
|
|
|
# Try setup again after installation
|
|
cuda_env = setup_cuda_environment()
|
|
if not cuda_env:
|
|
print("[ComfyUI-SAM3] [WARNING] CUDA environment setup failed after toolkit installation")
|
|
return False
|
|
|
|
# Get CUDA architecture list
|
|
cuda_arch_list = get_cuda_arch_list()
|
|
|
|
if cuda_arch_list:
|
|
cuda_env["TORCH_CUDA_ARCH_LIST"] = cuda_arch_list
|
|
os.environ["TORCH_CUDA_ARCH_LIST"] = cuda_arch_list # Also set in current process
|
|
print(f"[ComfyUI-SAM3] Set TORCH_CUDA_ARCH_LIST={cuda_arch_list}")
|
|
|
|
# Set C/C++ compiler paths explicitly (needed for CUDA compilation)
|
|
import shutil
|
|
gcc_path = shutil.which("gcc-11") or shutil.which("gcc")
|
|
gxx_path = shutil.which("g++-11") or shutil.which("g++")
|
|
if gcc_path:
|
|
cuda_env["CC"] = gcc_path
|
|
os.environ["CC"] = gcc_path
|
|
print(f"[ComfyUI-SAM3] Set CC={gcc_path}")
|
|
if gxx_path:
|
|
cuda_env["CXX"] = gxx_path
|
|
os.environ["CXX"] = gxx_path
|
|
print(f"[ComfyUI-SAM3] Set CXX={gxx_path}")
|
|
|
|
# Install cc_torch from GitHub
|
|
print("[ComfyUI-SAM3] Compiling cc_torch from source (may take 1-2 minutes)...")
|
|
|
|
try:
|
|
result = subprocess.run(
|
|
[
|
|
sys.executable, "-m", "pip", "install",
|
|
"--no-build-isolation", # Use current environment's torch
|
|
"git+https://github.com/ronghanghu/cc_torch.git"
|
|
],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=180, # Increased timeout for compilation
|
|
env=cuda_env # Pass CUDA environment explicitly to subprocess
|
|
)
|
|
|
|
if result.returncode == 0:
|
|
print("[ComfyUI-SAM3] [OK] cc_torch compiled and installed successfully")
|
|
print("[ComfyUI-SAM3] Video tracking will use GPU-accelerated connected components")
|
|
return True
|
|
else:
|
|
print("[ComfyUI-SAM3] [WARNING] cc_torch compilation failed")
|
|
print("[ComfyUI-SAM3] This is optional - video tracking will use CPU fallback")
|
|
if result.stderr:
|
|
print(f"[ComfyUI-SAM3] Error details: {result.stderr[:500]}")
|
|
return False
|
|
|
|
except subprocess.TimeoutExpired:
|
|
print("[ComfyUI-SAM3] [WARNING] cc_torch compilation timed out")
|
|
print("[ComfyUI-SAM3] Video tracking will use CPU fallback")
|
|
return False
|
|
except Exception as e:
|
|
print(f"[ComfyUI-SAM3] [WARNING] cc_torch installation error: {e}")
|
|
print("[ComfyUI-SAM3] Continuing without GPU connected components")
|
|
return False
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# Parse command line arguments
|
|
parser = argparse.ArgumentParser(description="Install ComfyUI-SAM3 dependencies")
|
|
parser.add_argument(
|
|
"--compile-only",
|
|
action="store_true",
|
|
help="Compile CUDA extensions even without GPU present (for CI testing)"
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
# Set global flag
|
|
COMPILE_ONLY = args.compile_only
|
|
|
|
if COMPILE_ONLY:
|
|
print("[ComfyUI-SAM3] Running in COMPILE-ONLY mode (will skip GPU runtime checks)")
|
|
|
|
print("[ComfyUI-SAM3] Running installation script...")
|
|
print("="*80)
|
|
|
|
# Install requirements.txt first
|
|
install_requirements()
|
|
|
|
print("="*80)
|
|
|
|
# Then try to install GPU-accelerated dependencies
|
|
try_install_torch_generic_nms()
|
|
|
|
print("="*80)
|
|
|
|
try_install_cc_torch()
|
|
|
|
print("="*80)
|
|
print("[ComfyUI-SAM3] Installation script completed")
|