500 lines
18 KiB
Python
500 lines
18 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Common shared functions for GPU acceleration setup in ComfyUI-SAM3.
|
|
|
|
This module contains shared code used by both speedup.py and speedup_blackwell.py
|
|
to avoid duplication and ensure consistent behavior.
|
|
"""
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
|
|
|
|
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
|
|
nvcc_dir = os.path.dirname(nvcc_in_path)
|
|
if os.path.basename(nvcc_dir) == "bin":
|
|
cuda_home = os.path.dirname(nvcc_dir)
|
|
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:
|
|
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)
|
|
system_paths = [
|
|
"/usr/local/cuda-12.8",
|
|
"/usr/local/cuda-12.7",
|
|
"/usr/local/cuda-12.6",
|
|
"/usr/local/cuda-12.5",
|
|
"/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):
|
|
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:
|
|
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):
|
|
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:
|
|
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 setup_cuda_environment():
|
|
"""
|
|
Setup CUDA environment variables (CUDA_HOME and PATH) for compilation.
|
|
Returns dict with environment variables if CUDA found, None otherwise.
|
|
"""
|
|
cuda_home = find_cuda_home()
|
|
|
|
if cuda_home:
|
|
env = os.environ.copy()
|
|
|
|
env["CUDA_HOME"] = cuda_home
|
|
os.environ["CUDA_HOME"] = cuda_home
|
|
print(f"[ComfyUI-SAM3] Set CUDA_HOME={cuda_home}")
|
|
|
|
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"]
|
|
print(f"[ComfyUI-SAM3] Added to PATH: {cuda_bin}")
|
|
|
|
if check_nvcc_available():
|
|
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 ensure_micromamba_installed():
|
|
"""
|
|
Ensure micromamba is available. If not found, download and install it.
|
|
Returns True if micromamba is available, False otherwise.
|
|
"""
|
|
import shutil
|
|
|
|
# Check if micromamba, mamba, or conda already exists
|
|
for cmd in ["micromamba", "mamba", "conda"]:
|
|
if shutil.which(cmd):
|
|
print(f"[ComfyUI-SAM3] Found {cmd}")
|
|
return True
|
|
|
|
print("[ComfyUI-SAM3] No conda-based package manager found")
|
|
print("[ComfyUI-SAM3] Installing micromamba...")
|
|
|
|
try:
|
|
import tempfile
|
|
import tarfile
|
|
import urllib.request
|
|
|
|
# Determine platform
|
|
if sys.platform.startswith('linux'):
|
|
platform = 'linux-64'
|
|
elif sys.platform == 'darwin':
|
|
import platform as plat
|
|
if plat.machine() == 'arm64':
|
|
platform = 'osx-arm64'
|
|
else:
|
|
platform = 'osx-64'
|
|
else:
|
|
print(f"[ComfyUI-SAM3] [WARNING] Unsupported platform for micromamba: {sys.platform}")
|
|
return False
|
|
|
|
# Download micromamba
|
|
url = f"https://micro.mamba.pm/api/micromamba/{platform}/latest"
|
|
print(f"[ComfyUI-SAM3] Downloading micromamba from {url}...")
|
|
|
|
with tempfile.NamedTemporaryFile(delete=False, suffix='.tar.bz2') as tmp_file:
|
|
urllib.request.urlretrieve(url, tmp_file.name)
|
|
tmp_path = tmp_file.name
|
|
|
|
# Extract to a local directory
|
|
script_dir = os.path.dirname(os.path.abspath(__file__))
|
|
micromamba_dir = os.path.join(script_dir, "bin")
|
|
os.makedirs(micromamba_dir, exist_ok=True)
|
|
|
|
print(f"[ComfyUI-SAM3] Extracting micromamba to {micromamba_dir}...")
|
|
with tarfile.open(tmp_path, 'r:bz2') as tar:
|
|
for member in tar.getmembers():
|
|
if member.name == 'bin/micromamba':
|
|
member.name = 'micromamba'
|
|
tar.extract(member, micromamba_dir)
|
|
break
|
|
|
|
# Make executable
|
|
micromamba_path = os.path.join(micromamba_dir, "micromamba")
|
|
os.chmod(micromamba_path, 0o755)
|
|
|
|
# Add to PATH for current session
|
|
os.environ["PATH"] = micromamba_dir + os.pathsep + os.environ.get("PATH", "")
|
|
|
|
# Clean up
|
|
os.unlink(tmp_path)
|
|
|
|
# Verify installation
|
|
if shutil.which("micromamba"):
|
|
print(f"[ComfyUI-SAM3] [OK] micromamba installed successfully at {micromamba_path}")
|
|
return True
|
|
else:
|
|
print("[ComfyUI-SAM3] [WARNING] micromamba installed but not found in PATH")
|
|
return False
|
|
|
|
except Exception as e:
|
|
print(f"[ComfyUI-SAM3] [WARNING] Failed to install micromamba: {e}")
|
|
return False
|
|
|
|
|
|
def try_install_cuda_micromamba(cuda_version=None):
|
|
"""
|
|
Try to install CUDA toolkit using micromamba/conda/mamba.
|
|
Returns True if successful, False otherwise.
|
|
"""
|
|
import shutil
|
|
|
|
print("[ComfyUI-SAM3] Attempting conda-based installation...")
|
|
|
|
# Ensure micromamba is available
|
|
if not ensure_micromamba_installed():
|
|
print("[ComfyUI-SAM3] Could not setup micromamba, skipping conda-based installation")
|
|
return False
|
|
|
|
# Check for micromamba, mamba, or conda (in preference order)
|
|
installer = None
|
|
for cmd in ["micromamba", "mamba", "conda"]:
|
|
if shutil.which(cmd):
|
|
installer = cmd
|
|
print(f"[ComfyUI-SAM3] Using {cmd} for CUDA toolkit installation...")
|
|
break
|
|
|
|
if not installer:
|
|
print("[ComfyUI-SAM3] No conda-based package manager found after installation attempt")
|
|
return False
|
|
|
|
# Only proceed if we're in a conda environment
|
|
if "CONDA_PREFIX" not in os.environ:
|
|
print("[ComfyUI-SAM3] Not running in conda environment, skipping conda installation")
|
|
return False
|
|
|
|
try:
|
|
# Determine package to install
|
|
package_spec = "cuda-toolkit"
|
|
if cuda_version:
|
|
try:
|
|
parts = cuda_version.split('.')
|
|
cuda_major = int(parts[0])
|
|
cuda_minor = int(parts[1]) if len(parts) > 1 else 0
|
|
|
|
package_spec = f"cuda-toolkit={cuda_major}.{cuda_minor}"
|
|
print(f"[ComfyUI-SAM3] Targeting CUDA version: {package_spec}")
|
|
except (ValueError, IndexError):
|
|
print(f"[ComfyUI-SAM3] Using generic cuda-toolkit package")
|
|
|
|
# Install CUDA toolkit
|
|
print(f"[ComfyUI-SAM3] Running: {installer} install -y {package_spec} -c conda-forge")
|
|
result = subprocess.run(
|
|
[installer, "install", "-y", package_spec, "-c", "conda-forge"],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=600
|
|
)
|
|
|
|
if result.returncode == 0:
|
|
print(f"[ComfyUI-SAM3] [OK] CUDA toolkit installed via {installer}")
|
|
|
|
if setup_cuda_environment():
|
|
return True
|
|
else:
|
|
print("[ComfyUI-SAM3] [WARNING] Installation succeeded but nvcc not accessible")
|
|
conda_prefix = os.environ.get("CONDA_PREFIX", "")
|
|
nvcc_path = os.path.join(conda_prefix, "bin", "nvcc")
|
|
print(f"[ComfyUI-SAM3] Expected nvcc at: {nvcc_path}")
|
|
print(f"[ComfyUI-SAM3] Exists: {os.path.exists(nvcc_path)}")
|
|
return False
|
|
else:
|
|
print(f"[ComfyUI-SAM3] [WARNING] {installer} installation failed")
|
|
if result.stderr:
|
|
print(f"[ComfyUI-SAM3] Error: {result.stderr[:300]}")
|
|
|
|
# Try fallback to major version only
|
|
if cuda_version and "=" in package_spec:
|
|
cuda_major = cuda_version.split('.')[0]
|
|
fallback_spec = f"cuda-toolkit={cuda_major}.*"
|
|
print(f"[ComfyUI-SAM3] Trying fallback: {fallback_spec}")
|
|
result = subprocess.run(
|
|
[installer, "install", "-y", fallback_spec, "-c", "conda-forge"],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=600
|
|
)
|
|
if result.returncode == 0 and setup_cuda_environment():
|
|
print(f"[ComfyUI-SAM3] [OK] CUDA toolkit installed via {installer} with fallback spec")
|
|
return True
|
|
|
|
return False
|
|
|
|
except (FileNotFoundError, subprocess.TimeoutExpired) as e:
|
|
print(f"[ComfyUI-SAM3] {installer} installation failed: {e}")
|
|
return False
|
|
except Exception as e:
|
|
print(f"[ComfyUI-SAM3] Unexpected error during {installer} installation: {e}")
|
|
return False
|
|
|
|
|
|
def try_install_cuda_toolkit():
|
|
"""
|
|
Attempt to install CUDA toolkit if nvcc is missing.
|
|
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
|
|
|
|
cuda_major = cuda_version.split('.')[0]
|
|
print(f"[ComfyUI-SAM3] PyTorch CUDA version: {cuda_version} (major: cu{cuda_major})")
|
|
|
|
# Try micromamba/conda first
|
|
if try_install_cuda_micromamba(cuda_version):
|
|
return True
|
|
|
|
# If that fails, provide manual instructions
|
|
print("[ComfyUI-SAM3] [WARNING] Could not install CUDA compiler automatically")
|
|
print("[ComfyUI-SAM3] Manual installation required:")
|
|
print("[ComfyUI-SAM3] https://developer.nvidia.com/cuda-downloads")
|
|
print("[ComfyUI-SAM3] Select your platform and follow the instructions")
|
|
return False
|
|
|
|
except Exception as e:
|
|
print(f"[ComfyUI-SAM3] [WARNING] CUDA toolkit installation error: {e}")
|
|
return False
|
|
|
|
|
|
def install_extension(name, repo_url, cuda_arch_list, compile_only=False, timeout=600, verbose=False):
|
|
"""
|
|
Unified function to install a CUDA extension from GitHub.
|
|
|
|
Args:
|
|
name: Extension name (e.g., "torch_generic_nms")
|
|
repo_url: GitHub repository URL
|
|
cuda_arch_list: TORCH_CUDA_ARCH_LIST string
|
|
compile_only: Compile even without GPU present
|
|
timeout: Compilation timeout in seconds
|
|
verbose: Show real-time compilation output
|
|
|
|
Returns:
|
|
True if successful, False otherwise
|
|
"""
|
|
print(f"[ComfyUI-SAM3] Checking for {name}...")
|
|
|
|
# Check if already installed
|
|
try:
|
|
__import__(name)
|
|
print(f"[ComfyUI-SAM3] [OK] {name} 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(f"[ComfyUI-SAM3] [INFO] CUDA not available, skipping {name} compilation")
|
|
return False
|
|
elif compile_only:
|
|
print(f"[ComfyUI-SAM3] [INFO] Compile-only mode: checking PyTorch CUDA support...")
|
|
if not torch.version.cuda:
|
|
print(f"[ComfyUI-SAM3] [ERROR] PyTorch was not built with CUDA support")
|
|
return False
|
|
print(f"[ComfyUI-SAM3] [OK] PyTorch built with CUDA {torch.version.cuda}")
|
|
except ImportError:
|
|
print(f"[ComfyUI-SAM3] [WARNING] PyTorch not found, cannot install {name}")
|
|
return False
|
|
|
|
# Setup CUDA environment
|
|
cuda_env = setup_cuda_environment()
|
|
if not cuda_env:
|
|
print(f"[ComfyUI-SAM3] [WARNING] CUDA compiler not found in system")
|
|
print(f"[ComfyUI-SAM3] Attempting to install CUDA toolkit...")
|
|
|
|
if not try_install_cuda_toolkit():
|
|
print(f"[ComfyUI-SAM3] [WARNING] Could not install CUDA compiler")
|
|
return False
|
|
|
|
cuda_env = setup_cuda_environment()
|
|
if not cuda_env:
|
|
print(f"[ComfyUI-SAM3] [WARNING] CUDA environment setup failed after toolkit installation")
|
|
return False
|
|
|
|
# Set CUDA architecture list
|
|
if cuda_arch_list:
|
|
cuda_env["TORCH_CUDA_ARCH_LIST"] = cuda_arch_list
|
|
os.environ["TORCH_CUDA_ARCH_LIST"] = cuda_arch_list
|
|
print(f"[ComfyUI-SAM3] Set TORCH_CUDA_ARCH_LIST={cuda_arch_list}")
|
|
|
|
# Set C/C++ compiler paths explicitly
|
|
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 extension from GitHub
|
|
print(f"[ComfyUI-SAM3] Compiling {name} from source (estimated time: ~60 seconds per architecture)...")
|
|
|
|
pip_args = [
|
|
sys.executable, "-m", "pip", "install",
|
|
"--no-build-isolation",
|
|
]
|
|
|
|
if verbose:
|
|
pip_args.append("--verbose")
|
|
|
|
pip_args.append(repo_url)
|
|
|
|
try:
|
|
result = subprocess.run(
|
|
pip_args,
|
|
capture_output=not verbose,
|
|
text=True,
|
|
timeout=timeout,
|
|
env=cuda_env
|
|
)
|
|
|
|
if result.returncode == 0:
|
|
print(f"[ComfyUI-SAM3] [OK] {name} compiled and installed successfully")
|
|
return True
|
|
else:
|
|
print(f"[ComfyUI-SAM3] [WARNING] {name} compilation failed")
|
|
if result.stderr and not verbose:
|
|
print(f"[ComfyUI-SAM3] Error details: {result.stderr[:500]}")
|
|
return False
|
|
|
|
except subprocess.TimeoutExpired:
|
|
print(f"[ComfyUI-SAM3] [WARNING] {name} compilation timed out after {timeout}s")
|
|
return False
|
|
except Exception as e:
|
|
print(f"[ComfyUI-SAM3] [WARNING] {name} installation error: {e}")
|
|
return False
|