Ethanfel
c08fe58fe7
cupy is now a best-effort install for NVIDIA users. Non-CUDA setups (ROCm, CPU) skip cupy and use PyTorch fallback kernels instead. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
47 lines
1.4 KiB
Python
47 lines
1.4 KiB
Python
import subprocess
|
|
import sys
|
|
import os
|
|
|
|
|
|
def get_cupy_package():
|
|
"""Detect PyTorch's CUDA version and return the matching cupy package name."""
|
|
try:
|
|
import torch
|
|
if not torch.cuda.is_available():
|
|
return None
|
|
cuda_version = torch.version.cuda
|
|
if cuda_version is None:
|
|
return None
|
|
major = int(cuda_version.split(".")[0])
|
|
cupy_pkg = f"cupy-cuda{major}x"
|
|
return cupy_pkg
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def install():
|
|
# Install core requirements first
|
|
requirements_path = os.path.join(os.path.dirname(__file__), "requirements.txt")
|
|
subprocess.check_call([
|
|
sys.executable, "-m", "pip", "install", "-r", requirements_path
|
|
])
|
|
|
|
# Try to install cupy for NVIDIA users (optional, improves performance)
|
|
cupy_pkg = get_cupy_package()
|
|
if cupy_pkg:
|
|
try:
|
|
subprocess.check_call([
|
|
sys.executable, "-m", "pip", "install", cupy_pkg
|
|
])
|
|
print(f"[Tween] cupy installed ({cupy_pkg}) — fast CUDA kernels enabled")
|
|
except subprocess.CalledProcessError:
|
|
print(f"[Tween] WARNING: Could not install {cupy_pkg}. "
|
|
f"BIM-VFI, SGM-VFI, and GIMM-VFI will use slower PyTorch fallback.")
|
|
else:
|
|
print("[Tween] cupy skipped (no NVIDIA CUDA). "
|
|
"BIM-VFI, SGM-VFI, and GIMM-VFI will use PyTorch fallback.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
install()
|