305 lines
11 KiB
Python
305 lines
11 KiB
Python
import os
|
|
import sys
|
|
import subprocess
|
|
import platform
|
|
import argparse
|
|
import shutil
|
|
|
|
def run_command(command):
|
|
try:
|
|
subprocess.check_call(command)
|
|
except subprocess.CalledProcessError as e:
|
|
print(f"Error executing command: {e}")
|
|
return False
|
|
return True
|
|
|
|
def is_uv_available():
|
|
try:
|
|
subprocess.check_call([sys.executable, "-m", "uv", "--version"], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
|
|
return True
|
|
except:
|
|
return False
|
|
|
|
def get_env_info():
|
|
info = {}
|
|
|
|
# OS
|
|
if platform.system() == "Windows":
|
|
info['platform'] = "win_amd64"
|
|
else:
|
|
info['platform'] = "manylinux_2_35_x86_64" # Default specialized tag
|
|
|
|
# Python version
|
|
py_ver = f"cp{sys.version_info.major}{sys.version_info.minor}"
|
|
info['python'] = py_ver
|
|
|
|
# Torch and CUDA
|
|
try:
|
|
import torch
|
|
torch_raw = torch.__version__.split('+')[0]
|
|
torch_ver = torch_raw.split('.')
|
|
info['torch'] = f"torch{torch_ver[0]}{torch_ver[1]}"
|
|
info['torch_raw'] = torch_raw
|
|
|
|
if hasattr(torch.version, 'cuda') and torch.version.cuda:
|
|
cuda_ver = torch.version.cuda.split('.')
|
|
info['cuda'] = f"cu{cuda_ver[0]}{cuda_ver[1]}"
|
|
else:
|
|
info['cuda'] = None
|
|
except ImportError:
|
|
print("PyTorch not found. Please install PyTorch with CUDA support first.")
|
|
sys.exit(1)
|
|
|
|
return info
|
|
|
|
def show_recommendations():
|
|
print("\n" + "="*50)
|
|
print("RECOMMENDED ENVIRONMENTS")
|
|
print("="*50)
|
|
print("If you encounter 404 errors, it is likely because your environment")
|
|
print("combination is not yet supported on the wheel repository.")
|
|
print("\nPreferred configurations:")
|
|
print("\nYou can install the preferred Torch version using:")
|
|
print("pip install torch==2.5.1 --index-url https://download.pytorch.org/whl/cu124")
|
|
print("="*50 + "\n")
|
|
|
|
def find_wheel_url(lib_name, env):
|
|
import urllib.request
|
|
import re
|
|
|
|
index_url = f"https://pozzettiandrea.github.io/cuda-wheels/{lib_name}/"
|
|
|
|
try:
|
|
req = urllib.request.Request(index_url, headers={'User-Agent': 'Mozilla/5.0'})
|
|
with urllib.request.urlopen(req) as response:
|
|
html = response.read().decode('utf-8')
|
|
except Exception as e:
|
|
print(f" Warning: Could not access {lib_name} wheel index: {e}")
|
|
return None
|
|
|
|
# Find all wheel links
|
|
links = re.findall(r'https?://[^\s<>"]+\.whl', html)
|
|
if not links:
|
|
return None
|
|
|
|
target_py = env['python']
|
|
target_plat = "win_amd64" if platform.system() == "Windows" else "manylinux"
|
|
target_cuda = env['cuda'] # 'cu128' or 'cu130'
|
|
target_torch = env['torch'] # 'torch28' or 'torch210'
|
|
|
|
best_url = None
|
|
best_score = -1
|
|
|
|
for url in links:
|
|
filename = url.split('/')[-1]
|
|
|
|
# 1. Platform & Python MUST match
|
|
if target_plat not in filename:
|
|
continue
|
|
if target_py not in filename:
|
|
continue
|
|
|
|
# Parse CUDA and Torch from filename
|
|
# Example: flash_attn-2.8.3+cu124torch2.4-cp312-cp312-win_amd64.whl
|
|
cuda_match = re.search(r'cu(\d+)', filename)
|
|
torch_match = re.search(r'torch(\d+\.?\d*)', filename)
|
|
|
|
if not cuda_match or not torch_match:
|
|
continue
|
|
|
|
file_cuda = f"cu{cuda_match.group(1)}"
|
|
raw_torch_ver = torch_match.group(1)
|
|
# Handle cases where torch version might or might not have dots
|
|
file_torch = raw_torch_ver.replace(".", "") # '24' or '210'
|
|
file_torch_str = f"torch{file_torch}"
|
|
|
|
# Score calculation
|
|
score = 0
|
|
|
|
# 1. CUDA Major Compatibility
|
|
if file_cuda[:4] == target_cuda[:4]:
|
|
score += 1000 # Same major version
|
|
elif target_cuda.startswith('cu13') and file_cuda.startswith('cu12'):
|
|
score += 100 # CUDA 13 -> 12 fallback
|
|
else:
|
|
continue # Major mismatch
|
|
|
|
# 2. Torch Major Compatibility
|
|
if not file_torch or len(target_torch) < 6:
|
|
continue
|
|
if target_torch.startswith(f"torch{file_torch[0]}"):
|
|
score += 1000
|
|
else:
|
|
continue # Torch major mismatch
|
|
|
|
# 3. Exact Version Bonuses
|
|
if file_cuda == target_cuda:
|
|
score += 5000 # Massive bonus for perfect CUDA
|
|
else:
|
|
# Low-priority tie-breaker: prefer higher versions within the same major
|
|
try:
|
|
score += int(file_cuda[2:])
|
|
except: pass
|
|
|
|
if file_torch_str == target_torch:
|
|
score += 10000 # Massive bonus for perfect Torch
|
|
else:
|
|
# Prefer closest torch version
|
|
try:
|
|
diff = abs(int(file_torch) - int(target_torch[5:]))
|
|
score += (500 - min(diff * 10, 500))
|
|
except: pass
|
|
|
|
if score > best_score:
|
|
best_score = score
|
|
best_url = url
|
|
|
|
if best_url:
|
|
print(f" Selected best matching wheel: {best_url.split('/')[-1]} (Score: {best_score})")
|
|
return best_url
|
|
|
|
return None
|
|
|
|
def install_flash_attn(pip_base, env, dry_run=False):
|
|
print("\n--- Installing flash-attn ---")
|
|
wheel_url = find_wheel_url("flash-attn", env)
|
|
|
|
if wheel_url:
|
|
print(f" Found matching wheel: {wheel_url}")
|
|
if dry_run:
|
|
print(f" [Dry Run] Would run: {' '.join(pip_base + ['install', wheel_url])}")
|
|
return True
|
|
else:
|
|
return run_command(pip_base + ["install", wheel_url])
|
|
else:
|
|
print(f" Warning: No matching flash-attn wheel found for {env['cuda']}, {env['torch']}, {env['python']} on {env['platform']}.")
|
|
if platform.system() == "Windows":
|
|
print(" Skipping flash-attn on Windows as it is often incompatible and optional for inference.")
|
|
return True
|
|
return False
|
|
|
|
def install_triton_windows(pip_base, env, dry_run=False):
|
|
if platform.system() != "Windows":
|
|
return True
|
|
|
|
print("\n--- Installing triton-windows ---")
|
|
|
|
torch_raw = env.get('torch_raw', '2.0.0')
|
|
try:
|
|
from packaging import version
|
|
except ImportError:
|
|
class VersionFallback:
|
|
def __init__(self, v): self.v = [int(x) for x in v.split('.')[:2]]
|
|
def __ge__(self, other): return self.v >= [int(x) for x in other.split('.')[:2]]
|
|
version = type('obj', (object,), {'parse': lambda v: VersionFallback(v)})
|
|
|
|
v_torch = version.parse(torch_raw)
|
|
triton_spec = "triton-windows"
|
|
|
|
if v_torch >= version.parse("2.10.0"):
|
|
triton_spec = "triton-windows<3.7"
|
|
elif v_torch >= version.parse("2.9.0"):
|
|
triton_spec = "triton-windows<3.6"
|
|
elif v_torch >= version.parse("2.8.0"):
|
|
triton_spec = "triton-windows<3.5"
|
|
elif v_torch >= version.parse("2.7.0"):
|
|
triton_spec = "triton-windows<3.4"
|
|
elif v_torch >= version.parse("2.6.0"):
|
|
triton_spec = "triton-windows<3.3"
|
|
|
|
if dry_run:
|
|
print(f" [Dry Run] Would run: {' '.join(pip_base + ['install', '-U', triton_spec])}")
|
|
return True
|
|
else:
|
|
return run_command(pip_base + ["install", "-U", triton_spec])
|
|
|
|
def install():
|
|
parser = argparse.ArgumentParser(description="ComfyUI-SegviGen Installation")
|
|
parser.add_argument("--dry-run", action="store_true", help="Print the steps without installing")
|
|
args = parser.parse_args()
|
|
|
|
print("--- ComfyUI-SegviGen Installation ---")
|
|
if args.dry_run:
|
|
print("(DRY RUN MODE)")
|
|
|
|
env = get_env_info()
|
|
print(f"Detected Environment:")
|
|
print(f" OS: {platform.system()}")
|
|
print(f" Python: {env['python']}")
|
|
print(f" PyTorch: {env['torch']}")
|
|
print(f" CUDA: {env['cuda']}")
|
|
|
|
if not env['cuda']:
|
|
print("Error: CUDA not detected in PyTorch. These wheels require CUDA.")
|
|
sys.exit(1)
|
|
|
|
# 1. Install standard requirements
|
|
current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
requirements_path = os.path.join(current_dir, "requirements.txt")
|
|
|
|
use_uv = is_uv_available()
|
|
pip_base = [sys.executable, "-m", "uv", "pip"] if use_uv else [sys.executable, "-m", "pip"]
|
|
|
|
if os.path.exists(requirements_path):
|
|
print(f"\nInstalling requirements from {requirements_path}...")
|
|
if args.dry_run:
|
|
print(f" [Dry Run] Would run: {' '.join(pip_base + ['install', '-r', requirements_path])}")
|
|
else:
|
|
run_command(pip_base + ["install", "-r", requirements_path])
|
|
|
|
# 2. Install flash-attn
|
|
install_flash_attn(pip_base, env, dry_run=args.dry_run)
|
|
|
|
# 3. Install correct triton-windows version
|
|
install_triton_windows(pip_base, env, dry_run=args.dry_run)
|
|
|
|
# 4. Define other CUDA wheels
|
|
packages = [
|
|
{"name": "cumesh", "tag": "cumesh-latest"},
|
|
{"name": "flex-gemm", "tag": "flex_gemm-latest"},
|
|
{"name": "nvdiffrast", "tag": "nvdiffrast-latest"},
|
|
{"name": "nvdiffrec-render", "tag": "nvdiffrec_render-latest"},
|
|
{"name": "o-voxel", "tag": "o_voxel-latest"},
|
|
]
|
|
|
|
print("\nInstalling CUDA wheels...")
|
|
|
|
errors_occurred = False
|
|
for pkg in packages:
|
|
print(f"Installing {pkg['name']}...")
|
|
wheel_url = find_wheel_url(pkg['name'], env)
|
|
if wheel_url:
|
|
print(f" Found matching wheel: {wheel_url}")
|
|
if args.dry_run:
|
|
print(f" [Dry Run] Would run: {' '.join(pip_base + ['install', wheel_url])}")
|
|
else:
|
|
if not run_command(pip_base + ["install", wheel_url]):
|
|
print(f"Warning: Failed to install {pkg['name']}.")
|
|
errors_occurred = True
|
|
else:
|
|
print(f" Warning: No matching {pkg['name']} wheel found in the index for {env['platform']}.")
|
|
# Try git fallbacks for known libraries on Windows
|
|
git_fallbacks = {
|
|
"nvdiffrast": "git+https://github.com/NVlabs/nvdiffrast",
|
|
"flex_gemm": "git+https://github.com/NVlabs/flex-gemm",
|
|
"flex-gemm": "git+https://github.com/NVlabs/flex-gemm"
|
|
}
|
|
if pkg['name'] in git_fallbacks:
|
|
target = git_fallbacks[pkg['name']]
|
|
print(f" Attempting source installation for {pkg['name']} as fallback: {target}")
|
|
if args.dry_run:
|
|
print(f" [Dry Run] Would run: {' '.join(pip_base + ['install', target])}")
|
|
elif run_command(pip_base + ["install", target]):
|
|
print(f" Successfully installed {pkg['name']} from source.")
|
|
continue
|
|
|
|
errors_occurred = True
|
|
|
|
if errors_occurred:
|
|
show_recommendations()
|
|
|
|
print("\nInstallation complete!")
|
|
|
|
if __name__ == "__main__":
|
|
install()
|