From 52b58b4a80a35da539aac619174f2f34fc463baa Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Wed, 3 Apr 2024 20:06:40 +1100
Subject: [PATCH] feat(installer): remove extra GPU options

- Remove `CUDA_AND_DML`. This was for onnx, which we have since removed.
- Remove `AUTODETECT`. This option causes problems for windows users, as it falls back on default pypi index resulting in a non-CUDA torch being installed.
- Add more explicit settings for extra index URL, based on the torch website
- Fix bug where `xformers` wasn't installed on linux and/or windows when autodetect was selected
---
 installer/lib/installer.py | 19 +++++++++++++------
 installer/lib/messages.py  | 21 +--------------------
 2 files changed, 14 insertions(+), 26 deletions(-)

diff --git a/installer/lib/installer.py b/installer/lib/installer.py
index 45bc764023..5297366ccb 100644
--- a/installer/lib/installer.py
+++ b/installer/lib/installer.py
@@ -404,22 +404,29 @@ def get_torch_source() -> Tuple[str | None, str | None]:
     # device can be one of: "cuda", "rocm", "cpu", "cuda_and_dml, autodetect"
     device = select_gpu()
 
+    # The correct extra index URLs for torch are inconsistent, see https://pytorch.org/get-started/locally/#start-locally
+
     url = None
-    optional_modules = "[onnx]"
+    optional_modules: str | None = None
     if OS == "Linux":
         if device.value == "rocm":
             url = "https://download.pytorch.org/whl/rocm5.6"
         elif device.value == "cpu":
             url = "https://download.pytorch.org/whl/cpu"
-
+        elif device.value == "cuda":
+            # CUDA uses the default PyPi index
+            optional_modules = "[xformers,onnx-cuda]"
     elif OS == "Windows":
         if device.value == "cuda":
             url = "https://download.pytorch.org/whl/cu121"
             optional_modules = "[xformers,onnx-cuda]"
-        if device.value == "cuda_and_dml":
-            url = "https://download.pytorch.org/whl/cu121"
-            optional_modules = "[xformers,onnx-directml]"
+        elif device.value == "cpu":
+            # CPU  uses the default PyPi index, no optional modules
+            pass
+    elif OS == "Darwin":
+        # macOS uses the default PyPi index, no optional modules
+        pass
 
-    # in all other cases, Torch wheels should be coming from PyPi as of Torch 1.13
+    # Fall back to defaults
 
     return (url, optional_modules)
diff --git a/installer/lib/messages.py b/installer/lib/messages.py
index 257a587d9c..dcd65a9813 100644
--- a/installer/lib/messages.py
+++ b/installer/lib/messages.py
@@ -207,10 +207,8 @@ def dest_path(dest: Optional[str | Path] = None) -> Path | None:
 
 class GpuType(Enum):
     CUDA = "cuda"
-    CUDA_AND_DML = "cuda_and_dml"
     ROCM = "rocm"
     CPU = "cpu"
-    AUTODETECT = "autodetect"
 
 
 def select_gpu() -> GpuType:
@@ -226,10 +224,6 @@ def select_gpu() -> GpuType:
         "an [gold1 b]NVIDIA[/] GPU (using CUDA™)",
         GpuType.CUDA,
     )
-    nvidia_with_dml = (
-        "an [gold1 b]NVIDIA[/] GPU (using CUDA™, and DirectML™ for ONNX) -- ALPHA",
-        GpuType.CUDA_AND_DML,
-    )
     amd = (
         "an [gold1 b]AMD[/] GPU (using ROCm™)",
         GpuType.ROCM,
@@ -238,27 +232,19 @@ def select_gpu() -> GpuType:
         "Do not install any GPU support, use CPU for generation (slow)",
         GpuType.CPU,
     )
-    autodetect = (
-        "I'm not sure what to choose",
-        GpuType.AUTODETECT,
-    )
 
     options = []
     if OS == "Windows":
-        options = [nvidia, nvidia_with_dml, cpu]
+        options = [nvidia, cpu]
     if OS == "Linux":
         options = [nvidia, amd, cpu]
     elif OS == "Darwin":
         options = [cpu]
-        # future CoreML?
 
     if len(options) == 1:
         print(f'Your platform [gold1]{OS}-{ARCH}[/] only supports the "{options[0][1]}" driver. Proceeding with that.')
         return options[0][1]
 
-    # "I don't know" is always added the last option
-    options.append(autodetect)  # type: ignore
-
     options = {str(i): opt for i, opt in enumerate(options, 1)}
 
     console.rule(":space_invader: GPU (Graphics Card) selection :space_invader:")
@@ -292,11 +278,6 @@ def select_gpu() -> GpuType:
         ),
     )
 
-    if options[choice][1] is GpuType.AUTODETECT:
-        console.print(
-            "No problem. We will install CUDA support first :crossed_fingers: If Invoke does not detect a GPU, please re-run the installer and select one of the other GPU types."
-        )
-
     return options[choice][1]