[Desktop] Support Nvidia Blackwell (#3480)

Co-authored-by: Benjamin Lu <[email protected]>

Author: benceruleanlu

package-lock.json

@@ -71,7 +71,7 @@
   "dependencies": {
     "@alloc/quick-lru": "^5.2.0",
     "@atlaskit/pragmatic-drag-and-drop": "^1.3.1",
-    "@comfyorg/comfyui-electron-types": "^0.4.31",
+    "@comfyorg/comfyui-electron-types": "^0.4.36",
     "@comfyorg/litegraph": "^0.13.3",
     "@primevue/forms": "^4.2.5",
     "@primevue/themes": "^4.2.5",

package.json

@@ -37,7 +37,8 @@
 import {
   CUDA_TORCH_URL,
   NIGHTLY_CPU_TORCH_URL,
-  TorchDeviceType
+  TorchDeviceType,
+  TorchMirrorUrl
 } from '@comfyorg/comfyui-electron-types'
 import Divider from 'primevue/divider'
 import Panel from 'primevue/panel'
@@ -46,6 +47,7 @@ import { ModelRef, computed, onMounted, ref } from 'vue'
 import MirrorItem from '@/components/install/mirror/MirrorItem.vue'
 import { PYPI_MIRROR, PYTHON_MIRROR, UVMirror } from '@/constants/uvMirrors'
 import { t } from '@/i18n'
+import { electronAPI } from '@/utils/envUtil'
 import { isInChina } from '@/utils/networkUtil'
 import { ValidationState, mergeValidationStates } from '@/utils/validationUtil'
 
@@ -55,6 +57,17 @@ const pythonMirror = defineModel<string>('pythonMirror', { required: true })
 const pypiMirror = defineModel<string>('pypiMirror', { required: true })
 const torchMirror = defineModel<string>('torchMirror', { required: true })
 
+const isBlackwellArchitecture = ref(false)
+
+const requiresNightlyPytorch = async (): Promise<boolean> => {
+  try {
+    return await electronAPI().isBlackwell()
+  } catch (error) {
+    console.error('Failed to detect Blackwell architecture:', error)
+    return false
+  }
+}
+
 const getTorchMirrorItem = (device: TorchDeviceType): UVMirror => {
   const settingId = 'Comfy-Desktop.UV.TorchInstallMirror'
   switch (device) {
@@ -65,6 +78,13 @@ const getTorchMirrorItem = (device: TorchDeviceType): UVMirror => {
         fallbackMirror: NIGHTLY_CPU_TORCH_URL
       }
     case 'nvidia':
+      if (isBlackwellArchitecture.value) {
+        return {
+          settingId,
+          mirror: TorchMirrorUrl.NightlyCuda,
+          fallbackMirror: TorchMirrorUrl.NightlyCuda
+        }
+      }
       return {
         settingId,
         mirror: CUDA_TORCH_URL,
@@ -83,6 +103,7 @@ const getTorchMirrorItem = (device: TorchDeviceType): UVMirror => {
 const userIsInChina = ref(false)
 onMounted(async () => {
   userIsInChina.value = await isInChina()
+  isBlackwellArchitecture.value = await requiresNightlyPytorch()
 })
 
 const useFallbackMirror = (mirror: UVMirror) => ({