61 lines
1.9 KiB
Python
61 lines
1.9 KiB
Python
import os
|
|
import subprocess
|
|
import venv
|
|
import sys
|
|
|
|
import importlib.metadata
|
|
|
|
venv_dir = "separator_venv"
|
|
requirements_file = "requirements.txt"
|
|
|
|
if not os.path.exists(venv_dir):
|
|
print(f"Creating virtual environment '{venv_dir}'...")
|
|
venv.create(venv_dir, with_pip=True)
|
|
else:
|
|
print(f"Virtual environment '{venv_dir}' already exists.")
|
|
|
|
if os.name == "nt":
|
|
pip_executable = os.path.join(venv_dir, "Scripts", "pip.exe")
|
|
else:
|
|
pip_executable = os.path.join(venv_dir, "bin", "pip")
|
|
|
|
if os.path.exists(requirements_file):
|
|
print(f"Installing packages from '{requirements_file}'...")
|
|
subprocess.check_call([pip_executable, "install", "-r", requirements_file])
|
|
print("Packages installed successfully.")
|
|
else:
|
|
print(f"'{requirements_file}' not found.")
|
|
|
|
|
|
|
|
def install_pytorch():
|
|
subprocess.check_call([sys.executable, "-m", "pip", "install", "--upgrade",
|
|
"torch", "torchaudio", "torchvision", "--index-url", "https://download.pytorch.org/whl/cu129"])
|
|
|
|
def is_installed(package: str, version: str | None = None, strict: bool = True) -> bool:
|
|
try:
|
|
installed_version = importlib.metadata.version(package)
|
|
if version is not None:
|
|
if strict:
|
|
return installed_version == version
|
|
else:
|
|
return installed_version.startswith(version)
|
|
return True
|
|
|
|
except importlib.metadata.PackageNotFoundError:
|
|
return False
|
|
except Exception as e:
|
|
print(f"[Warning] Failed to check package '{package}': {e}")
|
|
return False
|
|
|
|
if is_installed("torch"):
|
|
print("🔍 Checking for CUDA support and installing PyTorch...")
|
|
|
|
import torch
|
|
if not torch.cuda.is_available():
|
|
print("CUDA not available. Installing CUDA-enabled PyTorch...")
|
|
install_pytorch()
|
|
else:
|
|
print("✅ CUDA is available.")
|
|
|
|
print("\n✅ Setup complete!") |