From 2c8c8718e97f69beb642255fb02bb47105baa130 Mon Sep 17 00:00:00 2001
From: Chenlei Hu <huchenlei@proton.me>
Date: Tue, 28 Jan 2025 18:00:08 -0800
Subject: [PATCH] [Desktop] Apply default pypi mirror when installing CPU torch
 (#2367)

---
 .../install/MirrorsConfiguration.vue          | 49 +++++++++++++++----
 src/constants/uvMirrors.ts                    | 45 ++++++-----------
 src/views/InstallView.vue                     |  1 +
 3 files changed, 55 insertions(+), 40 deletions(-)

diff --git a/src/components/install/MirrorsConfiguration.vue b/src/components/install/MirrorsConfiguration.vue
index 510acf1d5..93eb07dc4 100644
--- a/src/components/install/MirrorsConfiguration.vue
+++ b/src/components/install/MirrorsConfiguration.vue
@@ -5,10 +5,7 @@
     :collapsed="!showMirrorInputs"
     pt:root="bg-neutral-800 border-none w-[600px]"
   >
-    <template
-      v-for="([item, modelValue], index) in mirrors"
-      :key="item.settingId"
-    >
+    <template v-for="([item, modelValue], index) in mirrors" :key="item.mirror">
       <Divider v-if="index > 0" />
 
       <MirrorItem
@@ -34,25 +31,59 @@
 </template>
 
 <script setup lang="ts">
-import _ from 'lodash'
+import {
+  CUDA_TORCH_URL,
+  NIGHTLY_CPU_TORCH_URL,
+  TorchDeviceType
+} from '@comfyorg/comfyui-electron-types'
 import Divider from 'primevue/divider'
 import Panel from 'primevue/panel'
-import { computed, ref } from 'vue'
+import { ModelRef, computed, ref } from 'vue'
 
 import MirrorItem from '@/components/install/mirror/MirrorItem.vue'
-import { UV_MIRRORS } from '@/constants/uvMirrors'
+import { PYPI_MIRROR, PYTHON_MIRROR, UVMirror } from '@/constants/uvMirrors'
 import { t } from '@/i18n'
 import { ValidationState, mergeValidationStates } from '@/utils/validationUtil'
 
 const showMirrorInputs = ref(false)
+const { device } = defineProps<{ device: TorchDeviceType }>()
 const pythonMirror = defineModel<string>('pythonMirror', { required: true })
 const pypiMirror = defineModel<string>('pypiMirror', { required: true })
 const torchMirror = defineModel<string>('torchMirror', { required: true })
 
-const mirrors = _.zip(UV_MIRRORS, [pythonMirror, pypiMirror, torchMirror])
+const getTorchMirrorItem = (device: TorchDeviceType): UVMirror => {
+  const settingId = 'Comfy-Desktop.UV.TorchInstallMirror'
+  switch (device) {
+    case 'mps':
+      return {
+        settingId,
+        mirror: NIGHTLY_CPU_TORCH_URL,
+        fallbackMirror: NIGHTLY_CPU_TORCH_URL
+      }
+    case 'nvidia':
+      return {
+        settingId,
+        mirror: CUDA_TORCH_URL,
+        fallbackMirror: CUDA_TORCH_URL
+      }
+    case 'cpu':
+    default:
+      return {
+        settingId,
+        mirror: PYPI_MIRROR.mirror,
+        fallbackMirror: PYPI_MIRROR.fallbackMirror
+      }
+  }
+}
+
+const mirrors = computed<[UVMirror, ModelRef<string>][]>(() => [
+  [PYTHON_MIRROR, pythonMirror],
+  [PYPI_MIRROR, pypiMirror],
+  [getTorchMirrorItem(device), torchMirror]
+])
 
 const validationStates = ref<ValidationState[]>(
-  mirrors.map(() => ValidationState.IDLE)
+  mirrors.value.map(() => ValidationState.IDLE)
 )
 const validationState = computed(() => {
   return mergeValidationStates(validationStates.value)
diff --git a/src/constants/uvMirrors.ts b/src/constants/uvMirrors.ts
index 21ed0ff2f..6631aef7c 100644
--- a/src/constants/uvMirrors.ts
+++ b/src/constants/uvMirrors.ts
@@ -1,8 +1,3 @@
-import { CUDA_TORCH_URL } from '@comfyorg/comfyui-electron-types'
-import { NIGHTLY_CPU_TORCH_URL } from '@comfyorg/comfyui-electron-types'
-
-import { electronAPI, isElectron } from '@/utils/envUtil'
-
 export interface UVMirror {
   /**
    * The setting id defined for the mirror.
@@ -22,30 +17,18 @@ export interface UVMirror {
   validationPathSuffix?: string
 }
 
-const DEFAULT_TORCH_MIRROR = isElectron()
-  ? electronAPI().getPlatform() === 'darwin'
-    ? NIGHTLY_CPU_TORCH_URL
-    : CUDA_TORCH_URL
-  : ''
+export const PYTHON_MIRROR: UVMirror = {
+  settingId: 'Comfy-Desktop.UV.PythonInstallMirror',
+  mirror:
+    'https://github.com/astral-sh/python-build-standalone/releases/download',
+  fallbackMirror:
+    'https://ghfast.top/https://github.com/astral-sh/python-build-standalone/releases/download',
+  validationPathSuffix:
+    '/cpython-3.12.4%2B20240713-aarch64-apple-darwin-install_only.tar.gz'
+}
 
-export const UV_MIRRORS: UVMirror[] = [
-  {
-    settingId: 'Comfy-Desktop.UV.PythonInstallMirror',
-    mirror:
-      'https://github.com/astral-sh/python-build-standalone/releases/download',
-    fallbackMirror:
-      'https://ghfast.top/https://github.com/astral-sh/python-build-standalone/releases/download',
-    validationPathSuffix:
-      '/cpython-3.12.4%2B20240713-aarch64-apple-darwin-install_only.tar.gz'
-  },
-  {
-    settingId: 'Comfy-Desktop.UV.PypiInstallMirror',
-    mirror: 'https://pypi.org/simple/',
-    fallbackMirror: 'https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple'
-  },
-  {
-    settingId: 'Comfy-Desktop.UV.TorchInstallMirror',
-    mirror: DEFAULT_TORCH_MIRROR,
-    fallbackMirror: DEFAULT_TORCH_MIRROR
-  }
-]
+export const PYPI_MIRROR: UVMirror = {
+  settingId: 'Comfy-Desktop.UV.PypiInstallMirror',
+  mirror: 'https://pypi.org/simple/',
+  fallbackMirror: 'https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple'
+}
diff --git a/src/views/InstallView.vue b/src/views/InstallView.vue
index c0c0754e7..9e1b6bbbf 100644
--- a/src/views/InstallView.vue
+++ b/src/views/InstallView.vue
@@ -82,6 +82,7 @@
             v-model:allowMetrics="allowMetrics"
           />
           <MirrorsConfiguration
+            :device="device"
             v-model:pythonMirror="pythonMirror"
             v-model:pypiMirror="pypiMirror"
             v-model:torchMirror="torchMirror"