diff --git a/install.sh b/install.sh index 5010116b0..34bc29dc3 100755 --- a/install.sh +++ b/install.sh @@ -39,6 +39,18 @@ check_install_jq_tar() { fi } +determine_avx_support() { + if grep -q avx512 /proc/cpuinfo; then + echo "-avx512" + elif grep -q avx2 /proc/cpuinfo; then + echo "-avx2" + elif grep -q avx /proc/cpuinfo; then + echo "-avx" + else + echo "" + fi +} + # Function to download and install nitro install_nitro() { rm -rf /tmp/nitro @@ -73,8 +85,10 @@ create_uninstall_script() { # Determine OS and architecture OS=$(uname -s) +ARCH=$(uname -m) VERSION="latest" GPU="" +AVX="" check_install_jq_tar @@ -91,6 +105,18 @@ do shift shift ;; + --avx) + AVX="-avx" + shift + ;; + --avx2) + AVX="-avx2" + shift + ;; + --avx512) + AVX="-avx512" + shift + ;; esac done @@ -126,7 +152,10 @@ fi # Construct download URL based on OS, ARCH, GPU and VERSION case $OS in Linux) - FILE_NAME="nitro-${VERSION}-linux-amd64${GPU}${CUDA_VERSION}.tar.gz" + if [ -z "$AVX" ]; then + AVX=$(determine_avx_support) + fi + FILE_NAME="nitro-${VERSION}-linux-amd64${AVX}${GPU}${CUDA_VERSION}.tar.gz" ;; Darwin) ARCH_FORMAT="mac-universal" @@ -140,6 +169,16 @@ esac DOWNLOAD_URL="https://github.com/janhq/nitro/releases/download/v${VERSION}/${FILE_NAME}" +# Check AVX support +if [ -z "$AVX" ] && [ "$OS" == "Linux" ]; then + echo "AVX is not supported on this system." + exit 1 +fi + +# Remove existing Nitro installation +echo "Removing existing Nitro installation..." +rm -rf /usr/local/bin/nitro + # Download, install, and create uninstall script install_nitro "$DOWNLOAD_URL" create_uninstall_script