From 6d423b383b340b4fd5e99ab03c1e22ad8a59c56d Mon Sep 17 00:00:00 2001
From: Jeffrey Morgan <jmorganca@gmail.com>
Date: Tue, 28 May 2024 14:41:50 -0700
Subject: [PATCH] Improve install experience on WSL2 and Linux (#4653)

---
 scripts/install.sh | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/scripts/install.sh b/scripts/install.sh
index 04b18796..c30de483 100644
--- a/scripts/install.sh
+++ b/scripts/install.sh
@@ -35,7 +35,7 @@ esac
 
 KERN=$(uname -r)
 case "$KERN" in
-    *icrosoft*WSL2 | *icrosoft*wsl2) ;;
+    *icrosoft*WSL2 | *icrosoft*wsl2) IS_WSL2=true;;
     *icrosoft) error "Microsoft WSL1 is not currently supported. Please upgrade to WSL2 with 'wsl --set-version <distro> 2'" ;;
     *) ;;
 esac
@@ -131,6 +131,17 @@ if available systemctl; then
     configure_systemd
 fi
 
+# WSL2 only supports GPUs via nvidia passthrough
+# so check for nvidia-smi to determine if GPU is available
+if [ "$IS_WSL2" = true ]; then
+    if available nvidia-smi && [ -z "$(nvidia-smi | grep -o "CUDA Version: [0-9]*\.[0-9]*")" ]; then
+        status "Nvidia GPU detected."
+    fi
+    install_success
+    exit 0
+fi
+
+# Install GPU dependencies on Linux
 if ! available lspci && ! available lshw; then
     warning "Unable to detect NVIDIA/AMD GPU. Install lspci or lshw to automatically detect and install GPU dependencies."
     exit 0
@@ -310,3 +321,4 @@ if command -v nvidia-persistenced > /dev/null 2>&1; then
 fi
 
 status "NVIDIA GPU ready."
+install_success