Update setup.py

setup.py

@@ -1,341 +1,476 @@
-#!/usr/bin/env python3
-"""
-Setup script for the Image Tagger application.
-This script checks and installs all required dependencies.
-"""
-
-# Python 3.12+ compatibility patch for pkgutil.ImpImporter
-import sys
-if sys.version_info >= (3, 12):
-    import pkgutil
-    import importlib.machinery
-    
-    # Add ImpImporter as a compatibility shim for older packages
-    if not hasattr(pkgutil, 'ImpImporter'):
-        class ImpImporter:
-            def __init__(self, path=None):
-                self.path = path
-            
-            def find_module(self, fullname, path=None):
-                return None
-        
-        pkgutil.ImpImporter = ImpImporter
-
-import os
-import sys
-import subprocess
-import platform
-from pathlib import Path
-import re
-import urllib.request
-import shutil
-import tempfile
-import time
-import webbrowser
-
-# Define the required packages
-SETUPTOOLS_PACKAGES = [
-    "setuptools>=58.0.0",
-    "setuptools-distutils>=0.3.0", 
-    "wheel>=0.38.0",
-]
-
-REQUIRED_PACKAGES = [
-    "streamlit>=1.21.0",
-    "pillow>=9.0.0",
-    # CRITICAL: Pin NumPy to 1.24.x and prevent 2.x installation
-    "numpy>=1.24.0,<2.0.0",
-    "ninja>=1.10.0",
-    "packaging>=20.0",
-    "matplotlib>=3.5.0",
-    "tqdm>=4.62.0",
-    "scipy>=1.7.0",
-    "safetensors>=0.3.0",
-    "timm>=0.9.0",  # Add this line for PyTorch Image Models
-]
-
-# Packages to install after PyTorch
-POST_TORCH_PACKAGES = [
-    "einops>=0.6.1",
-]
-
-CUDA_PACKAGES = {
-    "11.8": "torch==2.0.1+cu118 torchvision==0.15.2+cu118 --index-url https://download.pytorch.org/whl/cu118",
-    "11.7": "torch==2.0.1+cu117 torchvision==0.15.2+cu117 --index-url https://download.pytorch.org/whl/cu117", 
-    "11.6": "torch==2.0.1+cu116 torchvision==0.15.2+cu116 --index-url https://download.pytorch.org/whl/cu116",
-    "cpu": "torch==2.0.1+cpu torchvision==0.15.2+cpu --index-url https://download.pytorch.org/whl/cpu"
-}
-
-# ONNX and acceleration packages
-ONNX_PACKAGES = [
-    "onnx>=1.14.0",
-    "onnxruntime>=1.15.0",
-    "onnxruntime-gpu>=1.15.0;platform_system!='Darwin'",
-]
-
-# Colors for terminal output
-class Colors:
-    HEADER = '\033[95m'
-    BLUE = '\033[94m'
-    GREEN = '\033[92m'
-    WARNING = '\033[93m'
-    FAIL = '\033[91m'
-    ENDC = '\033[0m'
-    BOLD = '\033[1m'
-
-def print_colored(text, color):
-    """Print text in color"""
-    if sys.platform == "win32":
-        print(text)
-    else:
-        print(f"{color}{text}{Colors.ENDC}")
-
-def check_and_fix_numpy():
-    """Check for NumPy 2.x and fix compatibility issues"""
-    print_colored("\nChecking NumPy compatibility...", Colors.BLUE)
-    
-    pip_path = get_venv_pip()
-    
-    try:
-        result = subprocess.run([pip_path, "show", "numpy"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
-        if result.returncode == 0:
-            version_match = re.search(r"Version: ([\d\.]+)", result.stdout)
-            if version_match:
-                current_version = version_match.group(1)
-                print_colored(f"Found NumPy {current_version}", Colors.BLUE)
-                
-                if current_version.startswith("2."):
-                    print_colored(f"ERROR: NumPy {current_version} is incompatible with PyTorch!", Colors.FAIL)
-                    print_colored("Uninstalling NumPy 2.x and installing compatible version...", Colors.BLUE)
-                    
-                    # Force uninstall NumPy 2.x
-                    subprocess.run([pip_path, "uninstall", "-y", "numpy"], check=True)
-                    
-                    # Install compatible NumPy version with constraints
-                    subprocess.run([pip_path, "install", "numpy>=1.24.0,<2.0.0"], check=True)
-                    
-                    print_colored("[OK] NumPy downgraded to compatible version", Colors.GREEN)
-                    return True
-                    
-                elif current_version.startswith("1.24."):
-                    print_colored(f"[OK] NumPy {current_version} is compatible", Colors.GREEN)
-                    return False
-                else:
-                    print_colored(f"Updating NumPy to recommended version...", Colors.BLUE)
-                    subprocess.run([pip_path, "install", "numpy>=1.24.0,<2.0.0"], check=True)
-                    return True
-    except Exception as e:
-        print_colored(f"Error checking NumPy: {e}", Colors.WARNING)
-    
-    return False
-
-def install_packages(cuda_version):
-    """Install required packages using pip"""
-    print_colored("\nInstalling required packages...", Colors.BLUE)
-    
-    pip_path = get_venv_pip()
-    
-    # Upgrade pip first
-    try:
-        subprocess.run([pip_path, "install", "--upgrade", "pip"], check=True)
-        print_colored("[OK] Pip upgraded successfully", Colors.GREEN)
-    except subprocess.CalledProcessError:
-        print_colored("Warning: Failed to upgrade pip", Colors.WARNING)
-    
-    # Install setuptools packages first
-    print_colored("\nInstalling setuptools...", Colors.BLUE)
-    for package in SETUPTOOLS_PACKAGES:
-        try:
-            subprocess.run([pip_path, "install", package], check=True)
-            print_colored(f"[OK] Installed {package}", Colors.GREEN)
-        except subprocess.CalledProcessError as e:
-            print_colored(f"Warning: Issue installing {package}: {e}", Colors.WARNING)
-    
-    # Check and fix NumPy compatibility before installing other packages
-    numpy_was_updated = check_and_fix_numpy()
-    
-    # Install base packages
-    for package in REQUIRED_PACKAGES:
-        try:
-            print_colored(f"Installing {package}...", Colors.BLUE)
-            subprocess.run([pip_path, "install", package], check=True)
-            print_colored(f"[OK] Installed {package}", Colors.GREEN)
-        except subprocess.CalledProcessError as e:
-            print_colored(f"Error installing {package}: {e}", Colors.FAIL)
-            return False
-    
-    # If NumPy was updated, we need to reinstall PyTorch to ensure compatibility
-    if numpy_was_updated:
-        print_colored("\nNumPy was updated, ensuring PyTorch compatibility...", Colors.BLUE)
-        try:
-            # Uninstall existing PyTorch
-            subprocess.run([pip_path, "uninstall", "-y", "torch", "torchvision"], check=False)
-        except:
-            pass  # Ignore errors if not installed
-    
-    # Install PyTorch with appropriate CUDA version
-    print_colored(f"\nInstalling PyTorch {'with CUDA support' if cuda_version != 'cpu' else '(CPU version)'}...", Colors.BLUE)
-    torch_command = CUDA_PACKAGES[cuda_version].split()
-    try:
-        subprocess.run([pip_path, "install"] + torch_command, check=True)
-        print_colored("[OK] PyTorch installed successfully", Colors.GREEN)
-    except subprocess.CalledProcessError as e:
-        print_colored(f"Error installing PyTorch: {e}", Colors.FAIL)
-        return False
-    
-    # Install post-PyTorch packages
-    for package in POST_TORCH_PACKAGES:
-        try:
-            subprocess.run([pip_path, "install", package], check=True)
-            print_colored(f"[OK] Installed {package}", Colors.GREEN)
-        except subprocess.CalledProcessError as e:
-            print_colored(f"Error installing {package}: {e}", Colors.FAIL)
-            return False
-    
-    # Final NumPy compatibility check
-    print_colored("\nPerforming final compatibility check...", Colors.BLUE)
-    try:
-        # Test import in the virtual environment
-        python_path = get_venv_python()
-        test_cmd = [python_path, "-c", "import torch; import torchvision; import numpy; print('All imports successful')"]
-        result = subprocess.run(test_cmd, capture_output=True, text=True)
-        
-        if result.returncode == 0:
-            print_colored("[OK] All packages are compatible", Colors.GREEN)
-        else:
-            print_colored(f"Warning: Compatibility test failed: {result.stderr}", Colors.WARNING)
-            # Try to fix by reinstalling with --force-reinstall
-            print_colored("Attempting to fix with force reinstall...", Colors.BLUE)
-            subprocess.run([pip_path, "install", "--force-reinstall", "numpy>=1.24.0,<2.0.0"], check=True)
-            
-    except Exception as e:
-        print_colored(f"Warning: Could not perform compatibility check: {e}", Colors.WARNING)
-    
-    return True
-
-def check_python_version():
-    """Check if Python version is 3.8 or higher"""
-    print_colored("Checking Python version...", Colors.BLUE)
-    
-    version = sys.version_info
-    if version.major < 3 or (version.major == 3 and version.minor < 8):
-        print_colored("Error: Python 3.8 or higher is required. You have " + sys.version, Colors.FAIL)
-        return False
-    
-    print_colored(f"[OK] Python {version.major}.{version.minor}.{version.micro} detected", Colors.GREEN)
-    return True
-
-def create_virtual_env():
-    """Create a virtual environment if one doesn't exist"""
-    print_colored("\nChecking for virtual environment...", Colors.BLUE)
-    
-    venv_path = Path("venv")
-    if venv_path.exists():
-        print_colored("[OK] Virtual environment already exists", Colors.GREEN)
-        return True
-    
-    print_colored("Creating a new virtual environment...", Colors.BLUE)
-    try:
-        subprocess.run([sys.executable, "-m", "venv", "venv"], check=True)
-        print_colored("[OK] Virtual environment created successfully", Colors.GREEN)
-        return True
-    except subprocess.CalledProcessError:
-        print_colored("Error: Failed to create virtual environment", Colors.FAIL)
-        return False
-
-def get_venv_python():
-    """Get path to Python in the virtual environment"""
-    if sys.platform == "win32":
-        return os.path.join("venv", "Scripts", "python.exe")
-    else:
-        return os.path.join("venv", "bin", "python")
-
-def get_venv_pip():
-    """Get path to pip in the virtual environment"""
-    if sys.platform == "win32":
-        return os.path.join("venv", "Scripts", "pip.exe")
-    else:
-        return os.path.join("venv", "bin", "pip")
-
-def check_cuda():
-    """Check CUDA availability and version"""
-    print_colored("\nChecking for CUDA...", Colors.BLUE)
-    
-    cuda_available = False
-    cuda_version = None
-    
-    try:
-        if sys.platform == "win32":
-            process = subprocess.run(["where", "nvidia-smi"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
-        else:
-            process = subprocess.run(["which", "nvidia-smi"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
-        
-        if process.returncode == 0:
-            nvidia_smi = subprocess.run(["nvidia-smi"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
-            if nvidia_smi.returncode == 0:
-                cuda_available = True
-                match = re.search(r"CUDA Version: (\d+\.\d+)", nvidia_smi.stdout)
-                if match:
-                    cuda_version = match.group(1)
-    except Exception as e:
-        print_colored(f"Error checking CUDA: {str(e)}", Colors.WARNING)
-    
-    if cuda_available and cuda_version:
-        print_colored(f"[OK] CUDA {cuda_version} detected", Colors.GREEN)
-        for supported_version in CUDA_PACKAGES.keys():
-            if supported_version != "cpu" and float(supported_version) <= float(cuda_version):
-                return supported_version
-    
-    print_colored("No CUDA detected, using CPU-only version", Colors.WARNING)
-    return "cpu"
-
-def install_onnx_packages(cuda_version):
-    """Install ONNX packages"""
-    print_colored("\nInstalling ONNX packages...", Colors.BLUE)
-    
-    pip_path = get_venv_pip()
-    
-    try:
-        subprocess.run([pip_path, "install", "onnx>=1.14.0"], check=True)
-        
-        if cuda_version != "cpu":
-            subprocess.run([pip_path, "install", "onnxruntime-gpu>=1.15.0"], check=True)
-        else:
-            subprocess.run([pip_path, "install", "onnxruntime>=1.15.0"], check=True)
-            
-        print_colored("[OK] ONNX packages installed", Colors.GREEN)
-    except subprocess.CalledProcessError as e:
-        print_colored(f"Warning: ONNX installation issues: {e}", Colors.WARNING)
-    
-    return True
-
-def main():
-    """Main setup function"""
-    print_colored("=" * 60, Colors.HEADER)
-    print_colored("      Image Tagger - Setup Script", Colors.HEADER)
-    print_colored("=" * 60, Colors.HEADER)
-    
-    if not check_python_version():
-        return False
-    
-    if not create_virtual_env():
-        return False
-    
-    cuda_version = check_cuda()
-    
-    if not install_packages(cuda_version):
-        return False
-
-    if not install_onnx_packages(cuda_version):
-        print_colored("Warning: ONNX packages had issues", Colors.WARNING)
-    
-    print_colored("\n" + "=" * 60, Colors.HEADER)
-    print_colored("           Setup completed successfully!", Colors.GREEN)
-    print_colored("=" * 60, Colors.HEADER)
-    
-    return True
-
-if __name__ == "__main__":
-    success = main()
-    if not success:
+#!/usr/bin/env python3
+"""
+Setup script for the Image Tagger application.
+This script checks and installs all required dependencies.
+"""
+
+# Python 3.12+ compatibility patch for pkgutil.ImpImporter
+import sys
+if sys.version_info >= (3, 12):
+    import pkgutil
+    import importlib.machinery
+    
+    # Add ImpImporter as a compatibility shim for older packages
+    if not hasattr(pkgutil, 'ImpImporter'):
+        class ImpImporter:
+            def __init__(self, path=None):
+                self.path = path
+            
+            def find_module(self, fullname, path=None):
+                return None
+        
+        pkgutil.ImpImporter = ImpImporter
+
+import os
+import sys
+import subprocess
+import platform
+from pathlib import Path
+import re
+import urllib.request
+import shutil
+import tempfile
+import time
+import webbrowser
+
+# Define the required packages
+SETUPTOOLS_PACKAGES = [
+    "setuptools>=58.0.0",
+    "setuptools-distutils>=0.3.0", 
+    "wheel>=0.38.0",
+]
+
+REQUIRED_PACKAGES = [
+    "streamlit>=1.21.0",
+    "pillow>=9.0.0",
+    "ninja>=1.10.0",
+    "packaging>=20.0",
+    "matplotlib>=3.5.0",
+    "tqdm>=4.62.0",
+    "scipy>=1.7.0",
+    "safetensors>=0.3.0",
+]
+
+
+# Packages to install after PyTorch
+POST_TORCH_PACKAGES = [
+    "einops>=0.6.1",
+    "timm>=0.9.0",
+]
+
+
+CUDA_PACKAGES = {
+    "11.8": "torch==2.0.1+cu118 torchvision==0.15.2+cu118 --index-url https://download.pytorch.org/whl/cu118",
+    "11.7": "torch==2.0.1+cu117 torchvision==0.15.2+cu117 --index-url https://download.pytorch.org/whl/cu117", 
+    "11.6": "torch==2.0.1+cu116 torchvision==0.15.2+cu116 --index-url https://download.pytorch.org/whl/cu116",
+    "cpu": "torch==2.0.1+cpu torchvision==0.15.2+cpu --index-url https://download.pytorch.org/whl/cpu"
+}
+
+# ONNX and acceleration packages
+ONNX_PACKAGES = [
+    "onnx>=1.14.0",
+    "onnxruntime>=1.15.0",
+    "onnxruntime-gpu>=1.15.0;platform_system!='Darwin'",
+]
+
+# Colors for terminal output
+class Colors:
+    HEADER = '\033[95m'
+    BLUE = '\033[94m'
+    GREEN = '\033[92m'
+    WARNING = '\033[93m'
+    FAIL = '\033[91m'
+    ENDC = '\033[0m'
+    BOLD = '\033[1m'
+
+def print_colored(text, color):
+    """Print text in color"""
+    if sys.platform == "win32":
+        print(text)
+    else:
+        print(f"{color}{text}{Colors.ENDC}")
+
+def ensure_compatible_numpy() -> bool:
+    """
+    Install a NumPy version compatible with the current Python version.
+
+    - Python < 3.13: NumPy 1.24.x–1.x (we keep <2 to avoid ecosystem surprises)
+    - Python >= 3.13: NumPy 2.1+ (since 1.26 wheels don't support 3.13)
+    """
+    print_colored("\nEnsuring NumPy compatibility...", Colors.BLUE)
+
+    pip_path = get_venv_pip()
+
+    if sys.version_info >= (3, 13):
+        desired_spec = "numpy>=2.1.0,<3.0.0"
+        desired_major = 2
+    else:
+        desired_spec = "numpy>=1.24.0,<2.0.0"
+        desired_major = 1
+
+    def parse_ver(text: str):
+        m = re.search(r"Version:\s*([0-9]+)\.([0-9]+)\.([0-9]+)", text)
+        if not m:
+            return None
+        return (int(m.group(1)), int(m.group(2)), int(m.group(3)))
+
+    try:
+        show = subprocess.run(
+            [pip_path, "show", "numpy"],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True
+        )
+
+        if show.returncode == 0:
+            ver = parse_ver(show.stdout)
+            if ver:
+                major, minor, patch = ver
+                print_colored(f"Found NumPy {major}.{minor}.{patch}", Colors.BLUE)
+
+                ok = (major == desired_major) and (major != 1 or minor >= 24)
+                if ok:
+                    print_colored("[OK] NumPy version is compatible", Colors.GREEN)
+                    return False
+
+                print_colored("NumPy version not compatible for this Python; reinstalling...", Colors.BLUE)
+                subprocess.run([pip_path, "uninstall", "-y", "numpy"], check=False)
+
+        print_colored(f"Installing {desired_spec}...", Colors.BLUE)
+        subprocess.run([pip_path, "install", desired_spec], check=True)
+        print_colored(f"[OK] Installed {desired_spec}", Colors.GREEN)
+        return True
+
+    except subprocess.CalledProcessError as e:
+        print_colored(f"Error installing NumPy: {e}", Colors.FAIL)
+        return False
+
+def install_packages(cuda_version):
+    """Install required packages using pip"""
+    print_colored("\nInstalling required packages...", Colors.BLUE)
+    
+    pip_path = get_venv_pip()
+    
+    # Upgrade pip first
+    try:
+        subprocess.run([pip_path, "install", "--upgrade", "pip"], check=True)
+        print_colored("[OK] Pip upgraded successfully", Colors.GREEN)
+    except subprocess.CalledProcessError:
+        print_colored("Warning: Failed to upgrade pip", Colors.WARNING)
+    
+    # Install setuptools packages first
+    print_colored("\nInstalling setuptools...", Colors.BLUE)
+    for package in SETUPTOOLS_PACKAGES:
+        try:
+            subprocess.run([pip_path, "install", package], check=True)
+            print_colored(f"[OK] Installed {package}", Colors.GREEN)
+        except subprocess.CalledProcessError as e:
+            print_colored(f"Warning: Issue installing {package}: {e}", Colors.WARNING)
+    
+    # Check and fix NumPy compatibility before installing other packages
+    numpy_was_updated  = ensure_compatible_numpy()
+    
+    # Install base packages
+    for package in REQUIRED_PACKAGES:
+        try:
+            print_colored(f"Installing {package}...", Colors.BLUE)
+            subprocess.run([pip_path, "install", package], check=True)
+            print_colored(f"[OK] Installed {package}", Colors.GREEN)
+        except subprocess.CalledProcessError as e:
+            print_colored(f"Error installing {package}: {e}", Colors.FAIL)
+            return False
+    
+    # Install PyTorch with appropriate CUDA version
+    print_colored(
+        f"\nInstalling PyTorch ({'GPU '+cuda_version if cuda_version != 'cpu' else 'CPU-only'})...",
+        Colors.BLUE
+    )
+    if not install_pytorch(cuda_version):
+        return False
+    
+    # Install post-PyTorch packages
+    for package in POST_TORCH_PACKAGES:
+        try:
+            subprocess.run([pip_path, "install", package], check=True)
+            print_colored(f"[OK] Installed {package}", Colors.GREEN)
+        except subprocess.CalledProcessError as e:
+            print_colored(f"Error installing {package}: {e}", Colors.FAIL)
+            return False
+    
+    # Final NumPy compatibility check
+    print_colored("\nPerforming final compatibility check...", Colors.BLUE)
+    try:
+        # Test import in the virtual environment
+        python_path = get_venv_python()
+        test_cmd = [python_path, "-c", "import torch; import torchvision; import numpy; print('All imports successful')"]
+        result = subprocess.run(test_cmd, capture_output=True, text=True)
+        
+        if result.returncode == 0:
+            print_colored("[OK] All packages import successfully", Colors.GREEN)
+        else:
+            print_colored(f"Warning: Import test failed:\n{result.stderr}", Colors.WARNING)
+            
+    except Exception as e:
+        print_colored(f"Warning: Could not perform compatibility check: {e}", Colors.WARNING)
+    
+    return True
+
+def check_python_version():
+    v = sys.version_info
+    if (v.major, v.minor) < (3, 10) or (v.major, v.minor) >= (3, 14):
+        print_colored(
+            f"Error: Python 3.10–3.13 required. You have {v.major}.{v.minor}.",
+            Colors.FAIL
+        )
+        return False
+    return True
+
+def create_virtual_env():
+    """Create a virtual environment if one doesn't exist"""
+    print_colored("\nChecking for virtual environment...", Colors.BLUE)
+    
+    venv_path = Path("venv")
+    if venv_path.exists():
+        print_colored("[OK] Virtual environment already exists", Colors.GREEN)
+        return True
+    
+    print_colored("Creating a new virtual environment...", Colors.BLUE)
+    try:
+        subprocess.run([sys.executable, "-m", "venv", "venv"], check=True)
+        print_colored("[OK] Virtual environment created successfully", Colors.GREEN)
+        return True
+    except subprocess.CalledProcessError:
+        print_colored("Error: Failed to create virtual environment", Colors.FAIL)
+        return False
+
+def get_venv_python():
+    """Get path to Python in the virtual environment"""
+    if sys.platform == "win32":
+        return os.path.join("venv", "Scripts", "python.exe")
+    else:
+        return os.path.join("venv", "bin", "python")
+
+def get_venv_pip():
+    """Get path to pip in the virtual environment"""
+    if sys.platform == "win32":
+        return os.path.join("venv", "Scripts", "pip.exe")
+    else:
+        return os.path.join("venv", "bin", "pip")
+
+def check_cuda():
+    """Check CUDA availability and version"""
+    print_colored("\nChecking for CUDA...", Colors.BLUE)
+    
+    cuda_available = False
+    cuda_version = None
+    
+    try:
+        if sys.platform == "win32":
+            process = subprocess.run(["where", "nvidia-smi"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
+        else:
+            process = subprocess.run(["which", "nvidia-smi"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
+        
+        if process.returncode == 0:
+            nvidia_smi = subprocess.run(["nvidia-smi"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
+            if nvidia_smi.returncode == 0:
+                cuda_available = True
+                match = re.search(r"CUDA Version: (\d+\.\d+)", nvidia_smi.stdout)
+                if match:
+                    cuda_version = match.group(1)
+    except Exception as e:
+        print_colored(f"Error checking CUDA: {str(e)}", Colors.WARNING)
+    
+    if cuda_available and cuda_version:
+        print_colored(f"[OK] CUDA {cuda_version} detected", Colors.GREEN)
+        return _pick_pytorch_tag_from_driver(cuda_version)
+
+    print_colored("No CUDA detected, using CPU-only version", Colors.WARNING)
+    return "cpu"
+
+
+# Order matters: we try the newest runtime your driver supports, then fall back.
+PYTORCH_CUDA_TAGS = [
+    (13.0, "cu130"),
+    (12.9, "cu129"),
+    (12.8, "cu128"),
+    (12.6, "cu126"),
+    (12.4, "cu124"),
+    (12.1, "cu121"),
+    (11.8, "cu118"),
+]
+
+def _pick_pytorch_tag_from_driver(cuda_version_str: str) -> str:
+    try:
+        v = float(cuda_version_str)
+    except Exception:
+        return "cpu"
+
+    for min_v, tag in PYTORCH_CUDA_TAGS:
+        if v >= min_v:
+            return tag
+    return "cpu"
+
+def _torch_install_attempt(pip_path: str, index_url: str, pre: bool = False) -> bool:
+    # Always add PyPI as extra index so dependencies are resolvable.
+    cmd = [pip_path, "install", "--upgrade"]
+    if pre:
+        cmd.append("--pre")
+    cmd += [
+        "torch", "torchvision",
+        "--index-url", index_url,
+        "--extra-index-url", "https://pypi.org/simple",
+    ]
+    try:
+        subprocess.run(cmd, check=True)
+        return True
+    except subprocess.CalledProcessError:
+        return False
+
+def install_pytorch(cuda_tag: str) -> bool:
+    """
+    Install torch/torchvision.
+    - On macOS: install from PyPI (MPS/CPU handled by upstream)
+    - On Windows/Linux: install from the PyTorch CUDA/CPU index matching the driver.
+    - If Python 3.13+ and stable fails, try nightly as a fallback.
+    """
+    pip_path = get_venv_pip()
+
+    print_colored("\nInstalling PyTorch...", Colors.BLUE)
+
+    # Clean out any prior torch installs (helps reruns / partial installs)
+    subprocess.run([pip_path, "uninstall", "-y", "torch", "torchvision", "torchaudio"], check=False)
+
+    if sys.platform == "darwin":
+        try:
+            subprocess.run([pip_path, "install", "--upgrade", "torch", "torchvision"], check=True)
+            print_colored("[OK] PyTorch installed (macOS)", Colors.GREEN)
+            return True
+        except subprocess.CalledProcessError as e:
+            print_colored(f"Error installing PyTorch on macOS: {e}", Colors.FAIL)
+            return False
+
+    # Build a fallback list from the detected tag downwards, then CPU last.
+    all_tags = [tag for _, tag in PYTORCH_CUDA_TAGS]
+    if cuda_tag == "cpu":
+        candidates = ["cpu"]
+    else:
+        if cuda_tag in all_tags:
+            start = all_tags.index(cuda_tag)
+            candidates = all_tags[start:] + ["cpu"]
+        else:
+            candidates = ["cu118", "cpu"]
+
+    # 1) Try stable
+    for tag in candidates:
+        index_url = f"https://download.pytorch.org/whl/{tag}"
+        print_colored(f"Trying stable PyTorch from {index_url} ...", Colors.BLUE)
+        if _torch_install_attempt(pip_path, index_url=index_url, pre=False):
+            print_colored(f"[OK] PyTorch installed successfully ({tag})", Colors.GREEN)
+            return True
+
+    # 2) If on 3.13+, try nightly (optional but helps early adopters)
+    if sys.version_info >= (3, 13):
+        for tag in candidates:
+            index_url = f"https://download.pytorch.org/whl/nightly/{tag}"
+            print_colored(f"Trying nightly PyTorch from {index_url} ...", Colors.BLUE)
+            if _torch_install_attempt(pip_path, index_url=index_url, pre=True):
+                print_colored(f"[OK] PyTorch nightly installed successfully ({tag})", Colors.GREEN)
+                return True
+
+    print_colored(
+        "Error: Could not install a compatible PyTorch build for this Python/CUDA combo.\n"
+        "Tip: Python 3.12 is the most broadly supported choice if you hit this.",
+        Colors.FAIL
+    )
+    return False
+
+def _try_pip_install(pip_path: str, spec: str) -> bool:
+    try:
+        subprocess.run([pip_path, "install", spec], check=True)
+        return True
+    except subprocess.CalledProcessError:
+        return False
+
+def install_onnx_packages(cuda_version):
+    """Install ONNX + an available ONNX Runtime for this OS/Python."""
+    print_colored("\nInstalling ONNX packages...", Colors.BLUE)
+    pip_path = get_venv_pip()
+
+    # ONNX itself is fine
+    if not _try_pip_install(pip_path, "onnx>=1.14.0"):
+        print_colored("Error: failed to install onnx", Colors.FAIL)
+        return False
+
+    # Python 3.14+ currently often has no official onnxruntime* wheels on PyPI
+    # (PyPI classifiers are up to 3.13 for onnxruntime/onnxruntime-gpu).
+    if sys.version_info >= (3, 14):
+        print_colored(
+            "Warning: ONNX Runtime wheels are typically not available on PyPI for Python 3.14+ yet.\n"
+            "ONNX installed, but runtime install is skipped. Use Python 3.12/3.13 for ONNX Runtime.",
+            Colors.WARNING
+        )
+        
+    # Decide what to try for runtime
+    is_windows = (sys.platform == "win32")
+    has_cuda = (cuda_version != "cpu")  # in your script this means NVIDIA/CUDA detected
+
+    # Preferred order:
+    # - If CUDA detected: onnxruntime-gpu
+    # - On Windows without CUDA: onnxruntime-directml (GPU via DirectML)
+    # - Fallback: onnxruntime (CPU)
+    tried = []
+
+    if has_cuda:
+        tried.append("onnxruntime-gpu>=1.15.0")
+        if _try_pip_install(pip_path, tried[-1]):
+            print_colored("[OK] Installed onnxruntime-gpu", Colors.GREEN)
+            return True
+
+    if is_windows:
+        tried.append("onnxruntime-directml>=1.15.0")
+        if _try_pip_install(pip_path, tried[-1]):
+            print_colored("[OK] Installed onnxruntime-directml", Colors.GREEN)
+            return True
+
+    tried.append("onnxruntime>=1.15.0")
+    if _try_pip_install(pip_path, tried[-1]):
+        print_colored("[OK] Installed onnxruntime (CPU)", Colors.GREEN)
+        return True
+
+    print_colored(
+        "Warning: Could not install any ONNX Runtime variant.\n"
+        f"Tried: {', '.join(tried)}",
+        Colors.WARNING
+    )
+    return False
+
+def main():
+    """Main setup function"""
+    print_colored("=" * 60, Colors.HEADER)
+    print_colored("      Image Tagger - Setup Script", Colors.HEADER)
+    print_colored("=" * 60, Colors.HEADER)
+    
+    if not check_python_version():
+        return False
+    
+    if not create_virtual_env():
+        return False
+    
+    cuda_version = check_cuda()
+    
+    if not install_packages(cuda_version):
+        return False
+
+    if not install_onnx_packages(cuda_version):
+        print_colored("Warning: ONNX packages had issues", Colors.WARNING)
+    
+    print_colored("\n" + "=" * 60, Colors.HEADER)
+    print_colored("           Setup completed successfully!", Colors.GREEN)
+    print_colored("=" * 60, Colors.HEADER)
+    
+    return True
+
+if __name__ == "__main__":
+    success = main()
+    if not success:
         sys.exit(1)