Add NVIDIA hardware info display to settings with PyTorch version switching capability. Auto-select optimal PyTorch version during initial setup based on detected hardware.

josStorer · josStorer · commit c213873dcf4f · 2025-06-25T00:04:21.000+08:00
diff --git a/backend-golang/hw_info.go b/backend-golang/hw_info.go
@@ -0,0 +1,101 @@
+package backend_golang
+
+import (
+	"errors"
+	"os/exec"
+	"strconv"
+	"strings"
+)
+
+func (a *App) GetNvidiaGpuCount() (int, error) {
+	// temperature.gpu,utilization.gpu,utilization.memory,memory.total,memory.free,memory.used
+	// gpu_name,gpu_bus_id,driver_version
+	// nvidia-smi --help-query-gpu
+	output, err := exec.Command("nvidia-smi", "--query-gpu=count", "--format=csv,noheader,nounits").CombinedOutput()
+	if err != nil {
+		return 0, err
+	}
+	return strconv.Atoi(strings.TrimSpace(string(output)))
+}
+
+func (a *App) GetCudaComputeCapability(index int) (string, error) {
+	output, err := exec.Command("nvidia-smi", "-i="+strconv.Itoa(index), "--query-gpu=compute_cap", "--format=csv,noheader,nounits").CombinedOutput()
+	if err != nil {
+		return "", err
+	}
+
+	computeCap := strings.TrimSpace(string(output))
+	if computeCap == "" {
+		return "", errors.New("compute capability is empty")
+	}
+
+	return computeCap, nil
+}
+
+func (a *App) GetMaxCudaComputeCapability() (string, error) {
+	gpuCount, err := a.GetNvidiaGpuCount()
+	if err != nil {
+		return "", err
+	}
+	maxComputeCap := "0.0"
+	for i := 0; i < gpuCount; i++ {
+		computeCap, err := a.GetCudaComputeCapability(i)
+		if err != nil {
+			return "", err
+		}
+		computeCapFloat, err := strconv.ParseFloat(computeCap, 64)
+		if err != nil {
+			return "", err
+		}
+		maxComputeCapFloat, err := strconv.ParseFloat(maxComputeCap, 64)
+		if err != nil {
+			return "", err
+		}
+		if computeCapFloat > maxComputeCapFloat {
+			maxComputeCap = computeCap
+		}
+	}
+	if maxComputeCap == "0.0" {
+		return "", errors.New("no cuda compute capability")
+	}
+	return maxComputeCap, nil
+}
+
+func (a *App) GetSupportedCudaVersion() (string, error) {
+	output, err := exec.Command("nvidia-smi", "--query").CombinedOutput()
+	if err != nil {
+		return "", err
+	}
+
+	lines := strings.Split(string(output), "\n")
+
+	for _, line := range lines {
+		if strings.Contains(line, "CUDA Version") {
+			return strings.TrimSpace(strings.Split(line, ":")[1]), nil
+		}
+	}
+
+	return "", errors.New("cuda version is empty")
+}
+
+func (a *App) GetTorchVersion(python string) (string, error) {
+	var err error
+	if python == "" {
+		python, err = a.GetPython()
+		if err != nil {
+			return "", err
+		}
+	}
+
+	output, err := exec.Command(python, "-c", "import torch; print(torch.__version__)").CombinedOutput()
+	if err != nil {
+		return "", err
+	}
+
+	version := strings.TrimSpace(string(output))
+	if version == "" {
+		return "", errors.New("torch version is empty")
+	}
+
+	return version, nil
+}
diff --git a/backend-golang/rwkv.go b/backend-golang/rwkv.go
@@ -4,6 +4,7 @@ package backend_golang
 import (
 	"encoding/json"
 	"errors"
+	"fmt"
 	"os"
 	"os/exec"
 	"runtime"
@@ -221,13 +222,49 @@ func (a *App) MergeLora(python string, useGpu bool, loraAlpha int, baseModel str
 	return Cmd(args...)
 }
 
-func (a *App) InstallPyDep(python string, cnMirror bool) (string, error) {
+func (a *App) InstallTorch(python string, cnMirror bool, torchVersion string, cuSourceVersion string) (string, error) {
+	if runtime.GOOS != "windows" {
+		return "", errors.New("only support windows")
+	}
+
+	var err error
+	cuSourceVersion = strings.Replace(cuSourceVersion, ".", "", 1)
+	torchWhlUrl := fmt.Sprintf("torch==%s --index-url https://download.pytorch.org/whl/cu%s", torchVersion, cuSourceVersion)
+	if python == "" {
+		python, err = a.GetPython()
+		if cnMirror && python == "py310/python.exe" {
+			torchWhlUrl = fmt.Sprintf("https://mirrors.aliyun.com/pytorch-wheels/cu%s/torch-%s+cu%s-cp310-cp310-win_amd64.whl", cuSourceVersion, torchVersion, cuSourceVersion)
+		}
+		if runtime.GOOS == "windows" {
+			python = `"%CD%/` + python + `"`
+		}
+	}
+	if err != nil {
+		return "", err
+	}
+
+	a.ChangeFileLine("./py310/python310._pth", 3, "Lib\\site-packages")
+	installScript := python + " ./backend-python/get-pip.py -i https://mirrors.aliyun.com/pypi/simple --no-warn-script-location\n" +
+		python + " -m pip install " + torchWhlUrl + " --no-warn-script-location\n" +
+		"exit"
+	if !cnMirror {
+		installScript = strings.Replace(installScript, " -i https://mirrors.aliyun.com/pypi/simple", "", -1)
+	}
+	err = os.WriteFile(a.exDir+"install-py-dep.bat", []byte(installScript), 0644)
+	if err != nil {
+		return "", err
+	}
+	return Cmd("install-py-dep.bat")
+}
+
+func (a *App) InstallPyDep(python string, cnMirror bool, torchVersion string, cuSourceVersion string) (string, error) {
 	var err error
-	torchWhlUrl := "torch==1.13.1 torchvision==0.14.1 torchaudio==0.13.1 --index-url https://download.pytorch.org/whl/cu117"
+	cuSourceVersion = strings.Replace(cuSourceVersion, ".", "", 1)
+	torchWhlUrl := fmt.Sprintf("torch==%s --index-url https://download.pytorch.org/whl/cu%s", torchVersion, cuSourceVersion)
 	if python == "" {
 		python, err = a.GetPython()
 		if cnMirror && python == "py310/python.exe" {
-			torchWhlUrl = "https://mirrors.aliyun.com/pytorch-wheels/cu117/torch-1.13.1+cu117-cp310-cp310-win_amd64.whl"
+			torchWhlUrl = fmt.Sprintf("https://mirrors.aliyun.com/pytorch-wheels/cu%s/torch-%s+cu%s-cp310-cp310-win_amd64.whl", cuSourceVersion, torchVersion, cuSourceVersion)
 		}
 		if runtime.GOOS == "windows" {
 			python = `"%CD%/` + python + `"`
diff --git a/frontend/src/_locales/ja/main.json b/frontend/src/_locales/ja/main.json
@@ -379,5 +379,9 @@
   "Download File": "ファイルをダウンロード",
   "Force Enable Deep Think, Currently Only DeepSeek API and RWKV Runner Server Support": "深い思考を強制的に有効化、現在はDeepSeek APIとRWKV Runnerサーバーのみサポート",
   "DeepThink": "深い思考",
-  "Reasoning": "推論思考"
+  "Reasoning": "推論思考",
+  "Pytorch Version": "Pytorchバージョン",
+  "Not Installed": "未インストール",
+  "Driver CUDA Version": "ドライバーCUDAバージョン",
+  "CUDA Compute Capability": "CUDA計算互換性"
 }
diff --git a/frontend/src/_locales/zh-hans/main.json b/frontend/src/_locales/zh-hans/main.json
@@ -379,5 +379,9 @@
   "Download File": "下载文件",
   "Force Enable Deep Think, Currently Only DeepSeek API and RWKV Runner Server Support": "强制启用深度思考, 目前仅 DeepSeek API 及 RWKV Runner 服务器支持",
   "DeepThink": "深度思考",
-  "Reasoning": "推理思考"
+  "Reasoning": "推理思考",
+  "Pytorch Version": "Pytorch版本",
+  "Not Installed": "未安装",
+  "Driver CUDA Version": "驱动CUDA版本",
+  "CUDA Compute Capability": "CUDA计算兼容性"
 }
diff --git a/frontend/src/pages/Settings.tsx b/frontend/src/pages/Settings.tsx
@@ -6,22 +6,38 @@ import {
   AccordionPanel,
   Dropdown,
   Input,
+  Label,
   Option,
   Switch,
 } from '@fluentui/react-components'
+import { compare } from 'compare-versions'
 import { observer } from 'mobx-react-lite'
 import { useTranslation } from 'react-i18next'
 import { toast } from 'react-toastify'
-import { RestartApp } from '../../wailsjs/go/backend_golang/App'
+import {
+  GetTorchVersion,
+  InstallTorch,
+  RestartApp,
+} from '../../wailsjs/go/backend_golang/App'
 import { Labeled } from '../components/Labeled'
 import { Page } from '../components/Page'
 import commonStore from '../stores/commonStore'
 import { Language, Languages } from '../types/settings'
 import { checkUpdate, toastWithButton } from '../utils'
+import {
+  getAvailableTorchCuVersion,
+  torchVersions,
+} from '../utils/get-available-torch-cu-version'
 
 export const GeneralSettings: FC = observer(() => {
   const { t } = useTranslation()
 
+  useEffect(() => {
+    if (commonStore.platform === 'windows' && !commonStore.torchVersion) {
+      commonStore.refreshTorchVersion()
+    }
+  }, [])
+
   return (
     <div className="flex flex-col gap-2">
       <Labeled
@@ -89,6 +105,63 @@ export const GeneralSettings: FC = observer(() => {
           }
         />
       )}
+      {commonStore.platform === 'windows' && (
+        <Labeled
+          label={t('Pytorch Version')}
+          flex
+          spaceBetween
+          content={
+            <Dropdown
+              style={{ minWidth: 0 }}
+              listbox={{ style: { minWidth: 'fit-content' } }}
+              value={commonStore.torchVersion || t('Not Installed')!}
+              selectedOptions={[
+                commonStore.torchVersion
+                  ? torchVersions.find((v) =>
+                      commonStore.torchVersion.includes(v)
+                    ) || ''
+                  : '',
+              ]}
+              onOptionSelect={(_, data) => {
+                const selectedVersion = data.optionValue
+                if (selectedVersion) {
+                  const isSelectingCurrent =
+                    commonStore.torchVersion.includes(selectedVersion)
+                  if (!isSelectingCurrent) {
+                    const { torchVersion, cuSourceVersion } =
+                      getAvailableTorchCuVersion(
+                        selectedVersion,
+                        commonStore.driverCudaVersion || '11.7'
+                      )
+                    InstallTorch(
+                      commonStore.settings.customPythonPath,
+                      commonStore.settings.cnMirror,
+                      torchVersion,
+                      cuSourceVersion
+                    )
+                      .then(() => {
+                        commonStore.refreshTorchVersion()
+                      })
+                      .catch((e) => {
+                        toast.error(e)
+                      })
+                  }
+                }
+              }}
+            >
+              {torchVersions.map((v) => (
+                <Option key={v} value={v}>
+                  {v}
+                </Option>
+              ))}
+            </Dropdown>
+          }
+        />
+      )}
+      {commonStore.platform === 'windows' &&
+        commonStore.cudaComputeCapability && (
+          <Label size="small">{`${t('Driver CUDA Version')}: ${commonStore.driverCudaVersion} - ${t('CUDA Compute Capability')}: ${commonStore.cudaComputeCapability}`}</Label>
+        )}
       <Labeled
         label={t('Dark Mode')}
         flex
diff --git a/frontend/src/startup.ts b/frontend/src/startup.ts
@@ -5,7 +5,10 @@ import { toast } from 'react-toastify'
 import manifest from '../../manifest.json'
 import {
   FileExists,
+  GetMaxCudaComputeCapability,
   GetPlatform,
+  GetSupportedCudaVersion,
+  GetTorchVersion,
   ListDirFiles,
   ReadJson,
 } from '../wailsjs/go/backend_golang/App'
@@ -75,6 +78,15 @@ export async function startup() {
       //   })
       // }
     })
+    setTimeout(() => {
+      GetMaxCudaComputeCapability().then((c) => {
+        commonStore.setCudaComputeCapability(c)
+      })
+      GetSupportedCudaVersion().then((v) => {
+        commonStore.setDriverCudaVersion(v)
+      })
+      commonStore.refreshTorchVersion()
+    }, 1000)
   }
 }
 
diff --git a/frontend/src/stores/commonStore.ts b/frontend/src/stores/commonStore.ts
@@ -2,6 +2,7 @@ import { ChartData } from 'chart.js'
 import i18n from 'i18next'
 import { makeAutoObservable } from 'mobx'
 import manifest from '../../../manifest.json'
+import { GetTorchVersion } from '../../wailsjs/go/backend_golang/App'
 import { WindowSetDarkTheme, WindowSetLightTheme } from '../../wailsjs/runtime'
 import {
   defaultCompositionPrompt,
@@ -72,6 +73,9 @@ class CommonStore {
   monitorData: MonitorData | null = null
   depComplete: boolean = false
   platform: Platform = 'windows'
+  cudaComputeCapability: string = ''
+  driverCudaVersion: string = ''
+  torchVersion: string = ''
   proxyPort: number = 0
   lastModelName: string = ''
   stateModels: string[] = []
@@ -371,6 +375,20 @@ class CommonStore {
     this.platform = value
   }
 
+  setCudaComputeCapability(value: string) {
+    this.cudaComputeCapability = value
+  }
+
+  setDriverCudaVersion(value: string) {
+    this.driverCudaVersion = value
+  }
+
+  refreshTorchVersion() {
+    GetTorchVersion(this.settings.customPythonPath).then((v) => {
+      this.torchVersion = v
+    })
+  }
+
   setProxyPort(value: number) {
     this.proxyPort = value
   }
diff --git a/frontend/src/utils/get-available-torch-cu-version.ts b/frontend/src/utils/get-available-torch-cu-version.ts
@@ -0,0 +1,30 @@
+import { compare } from 'compare-versions'
+
+export const torchVersions = ['1.13.1', '2.7.1']
+
+export function getAvailableTorchCuVersion(
+  torchVersion: string,
+  driverCudaVersion: string
+) {
+  let retTorchVersion = ''
+  let retCuSourceVersion = ''
+  const targetTorchVersion = torchVersion.split('+')[0]
+  if (compare(targetTorchVersion, '2.7.1', '>=')) {
+    retTorchVersion = '2.7.1'
+    if (compare(driverCudaVersion, '12.8', '>=')) {
+      retCuSourceVersion = '12.8'
+    } else if (compare(driverCudaVersion, '12.6', '>=')) {
+      retCuSourceVersion = '12.6'
+    } else {
+      retCuSourceVersion = '11.8'
+    }
+  } else {
+    retTorchVersion = '1.13.1'
+    if (compare(driverCudaVersion, '11.7', '>=')) {
+      retCuSourceVersion = '11.7'
+    } else {
+      retCuSourceVersion = '11.6'
+    }
+  }
+  return { torchVersion: retTorchVersion, cuSourceVersion: retCuSourceVersion }
+}
diff --git a/frontend/src/utils/index.tsx b/frontend/src/utils/index.tsx
diff --git a/frontend/wailsjs/go/backend_golang/App.d.ts b/frontend/wailsjs/go/backend_golang/App.d.ts
diff --git a/frontend/wailsjs/go/backend_golang/App.js b/frontend/wailsjs/go/backend_golang/App.js
