Update NVIDIA torch version on updates

src/constants.ts

@@ -201,6 +201,14 @@ export const AMD_TORCH_PACKAGES: string[] = [
   'https://repo.radeon.com/rocm/windows/rocm-rel-7.1.1/torchvision-0.24.0+rocmsdk20251116-cp312-cp312-win_amd64.whl',
 ];
 
+export const NVIDIA_TORCH_VERSION = '2.9.1+cu130';
+export const NVIDIA_TORCHVISION_VERSION = '0.24.1+cu130';
+export const NVIDIA_TORCH_PACKAGES: string[] = [
+  `torch==${NVIDIA_TORCH_VERSION}`,
+  `torchaudio==${NVIDIA_TORCH_VERSION}`,
+  `torchvision==${NVIDIA_TORCHVISION_VERSION}`,
+];
+
 /** The log files used by the desktop process. */
 export enum LogFile {
   /** The ComfyUI server log file. */

src/install/installationManager.ts

@@ -1,5 +1,7 @@
 import { Notification, app, dialog, shell } from 'electron';
 import log from 'electron-log/main';
+import { exec } from 'node:child_process';
+import { promisify } from 'node:util';
 
 import { strictIpcMain as ipcMain } from '@/infrastructure/ipcChannels';
 
@@ -14,12 +16,37 @@ import { CmCli } from '../services/cmCli';
 import { captureSentryException } from '../services/sentry';
 import { type HasTelemetry, ITelemetry, trackEvent } from '../services/telemetry';
 import { type DesktopConfig, useDesktopConfig } from '../store/desktopConfig';
-import { canExecuteShellCommand, validateHardware } from '../utils';
+import { canExecuteShellCommand, compareVersions, validateHardware } from '../utils';
 import type { ProcessCallbacks, VirtualEnvironment } from '../virtualEnvironment';
 import { createProcessCallbacks } from './createProcessCallbacks';
 import { InstallWizard } from './installWizard';
 import { Troubleshooting } from './troubleshooting';
 
+const execAsync = promisify(exec);
+const NVIDIA_DRIVER_MIN_VERSION = '580';
+
+/**
+ * Extracts the NVIDIA driver version from `nvidia-smi` output.
+ * @param output The `nvidia-smi` output to parse.
+ * @returns The driver version, if present.
+ */
+export function parseNvidiaDriverVersionFromSmiOutput(output: string): string | undefined {
+  const match = output.match(/driver version\s*:\s*([\d.]+)/i);
+  return match?.[1];
+}
+
+/**
+ * Returns `true` when the NVIDIA driver version is below the minimum.
+ * @param driverVersion The detected driver version.
+ * @param minimumVersion The minimum required driver version.
+ */
+export function isNvidiaDriverBelowMinimum(
+  driverVersion: string,
+  minimumVersion: string = NVIDIA_DRIVER_MIN_VERSION
+): boolean {
+  return compareVersions(driverVersion, minimumVersion) < 0;
+}
+
 /** High-level / UI control over the installation of ComfyUI server. */
 export class InstallationManager implements HasTelemetry {
   constructor(
@@ -354,13 +381,68 @@ export class InstallationManager implements HasTelemetry {
     try {
       await installation.virtualEnvironment.installComfyUIRequirements(callbacks);
       await installation.virtualEnvironment.installComfyUIManagerRequirements(callbacks);
+      await this.warnIfNvidiaDriverTooOld(installation);
+      await installation.virtualEnvironment.ensureRecommendedNvidiaTorch(callbacks);
       await installation.validate();
     } catch (error) {
       log.error('Error auto-updating packages:', error);
       await this.appWindow.loadPage('server-start');
     }
   }
 
+  /**
+   * Warns the user if their NVIDIA driver is too old for the required CUDA build.
+   * @param installation The current installation.
+   */
+  private async warnIfNvidiaDriverTooOld(installation: ComfyInstallation): Promise<void> {
+    if (process.platform !== 'win32') return;
+    if (installation.virtualEnvironment.selectedDevice !== 'nvidia') return;
+
+    const driverVersion =
+      (await this.getNvidiaDriverVersionFromSmi()) ?? (await this.getNvidiaDriverVersionFromSmiFallback());
+    if (!driverVersion) return;
+
+    if (!isNvidiaDriverBelowMinimum(driverVersion)) return;
+
+    await this.appWindow.showMessageBox({
+      type: 'warning',
+      title: 'Update NVIDIA Driver',
+      message: 'Your NVIDIA driver may be too old for PyTorch 2.9.1 + cu130.',
+      detail: `Detected driver version: ${driverVersion}\nRecommended minimum: ${NVIDIA_DRIVER_MIN_VERSION}\n\nPlease consider updating your NVIDIA drivers and retrying if you run into issues.`,
+      buttons: ['OK'],
+    });
+  }
+
+  /**
+   * Reads the NVIDIA driver version from nvidia-smi query output.
+   */
+  private async getNvidiaDriverVersionFromSmi(): Promise<string | undefined> {
+    try {
+      const { stdout } = await execAsync('nvidia-smi --query-gpu=driver_version --format=csv,noheader');
+      const version = stdout
+        .split(/\r?\n/)
+        .map((line) => line.trim())
+        .find(Boolean);
+      return version || undefined;
+    } catch (error) {
+      log.debug('Failed to read NVIDIA driver version via nvidia-smi query.', error);
+      return undefined;
+    }
+  }
+
+  /**
+   * Reads the NVIDIA driver version from nvidia-smi standard output.
+   */
+  private async getNvidiaDriverVersionFromSmiFallback(): Promise<string | undefined> {
+    try {
+      const { stdout } = await execAsync('nvidia-smi');
+      return parseNvidiaDriverVersionFromSmiOutput(stdout);
+    } catch (error) {
+      log.debug('Failed to read NVIDIA driver version via nvidia-smi output.', error);
+      return undefined;
+    }
+  }
+
   static setReinstallHandler(installation: ComfyInstallation) {
     ipcMain.removeHandler(IPC_CHANNELS.REINSTALL);
     ipcMain.handle(IPC_CHANNELS.REINSTALL, async () => await InstallationManager.reinstall(installation));

src/virtualEnvironment.ts

@@ -7,7 +7,15 @@ import { readdir, rm } from 'node:fs/promises';
 import os, { EOL } from 'node:os';
 import path from 'node:path';
 
-import { AMD_ROCM_SDK_PACKAGES, AMD_TORCH_PACKAGES, InstallStage, TorchMirrorUrl } from './constants';
+import {
+  AMD_ROCM_SDK_PACKAGES,
+  AMD_TORCH_PACKAGES,
+  InstallStage,
+  NVIDIA_TORCHVISION_VERSION,
+  NVIDIA_TORCH_PACKAGES,
+  NVIDIA_TORCH_VERSION,
+  TorchMirrorUrl,
+} from './constants';
 import { PythonImportVerificationError } from './infrastructure/pythonImportVerificationError';
 import { useAppState } from './main-process/appState';
 import { createInstallStageInfo } from './main-process/installStages';
@@ -16,7 +24,7 @@ import { runPythonImportVerifyScript } from './services/pythonImportVerifier';
 import { captureSentryException } from './services/sentry';
 import { HasTelemetry, ITelemetry, trackEvent } from './services/telemetry';
 import { getDefaultShell, getDefaultShellArgs } from './shell/util';
-import { pathAccessible } from './utils';
+import { compareVersions, pathAccessible } from './utils';
 
 export type ProcessCallbacks = {
   onStdout?: (data: string) => void;
@@ -43,6 +51,11 @@ interface PipInstallConfig {
   indexStrategy?: 'compatible' | 'unsafe-best-match';
 }
 
+type TorchPackageName = 'torch' | 'torchaudio' | 'torchvision';
+type TorchPackageVersions = Record<TorchPackageName, string | undefined>;
+
+const TORCH_PACKAGE_NAMES: TorchPackageName[] = ['torch', 'torchaudio', 'torchvision'];
+
 function getPipInstallArgs(config: PipInstallConfig): string[] {
   const installArgs = ['pip', 'install'];
 
@@ -629,6 +642,118 @@ export class VirtualEnvironment implements HasTelemetry, PythonExecutor {
     }
   }
 
+  /**
+   * Ensures NVIDIA installs use the recommended PyTorch packages.
+   * @param callbacks The callbacks to use for the command.
+   */
+  async ensureRecommendedNvidiaTorch(callbacks?: ProcessCallbacks): Promise<void> {
+    if (this.selectedDevice !== 'nvidia') return;
+
+    const installedVersions = await this.getInstalledTorchPackageVersions();
+    if (installedVersions && this.meetsMinimumNvidiaTorchVersions(installedVersions)) {
+      log.info('NVIDIA PyTorch packages already satisfy minimum recommended versions.', installedVersions);
+      return;
+    }
+
+    const torchMirror = this.torchMirror || getDefaultTorchMirror(this.selectedDevice);
+    const config: PipInstallConfig = {
+      packages: NVIDIA_TORCH_PACKAGES,
+      indexUrl: torchMirror,
+      prerelease: torchMirror.includes('nightly'),
+    };
+
+    const installArgs = getPipInstallArgs(config);
+    log.info('Installing recommended NVIDIA PyTorch packages.', { installedVersions });
+    const { exitCode: pinnedExitCode } = await this.runUvCommandAsync(installArgs, callbacks);
+
+    if (pinnedExitCode === 0) return;
+
+    log.warn('Failed to install recommended NVIDIA PyTorch packages. Falling back to unpinned install.', {
+      exitCode: pinnedExitCode,
+    });
+
+    const fallbackConfig: PipInstallConfig = {
+      packages: ['torch', 'torchvision', 'torchaudio'],
+      indexUrl: torchMirror,
+      prerelease: torchMirror.includes('nightly'),
+      upgradePackages: true,
+    };
+    const fallbackArgs = getPipInstallArgs(fallbackConfig);
+    const { exitCode: fallbackExitCode } = await this.runUvCommandAsync(fallbackArgs, callbacks);
+    if (fallbackExitCode !== 0) {
+      throw new Error(
+        `Failed to install NVIDIA PyTorch packages (pinned exit ${pinnedExitCode}, fallback exit ${fallbackExitCode})`
+      );
+    }
+  }
+
+  /**
+   * Reads installed torch package versions using `uv pip list --format=json`.
+   * @returns The torch package versions when available, otherwise `undefined`.
+   */
+  private async getInstalledTorchPackageVersions(): Promise<TorchPackageVersions | undefined> {
+    let stdout = '';
+    let stderr = '';
+    const callbacks: ProcessCallbacks = {
+      onStdout: (data) => {
+        stdout += data;
+      },
+      onStderr: (data) => {
+        stderr += data;
+      },
+    };
+
+    const { exitCode } = await this.runUvAsync(['pip', 'list', '--format=json', '--color=never'], callbacks);
+
+    if (exitCode !== 0) {
+      log.warn('Failed to read torch package versions.', { exitCode, stderr });
+      return undefined;
+    }
+
+    if (!stdout.trim()) {
+      log.warn('Torch package list output was empty.', { stderr });
+      return undefined;
+    }
+
+    let parsed: unknown;
+    try {
+      parsed = JSON.parse(stdout);
+    } catch (error) {
+      log.warn('Failed to parse torch package list output.', { error, stdout, stderr });
+      return undefined;
+    }
+
+    if (!Array.isArray(parsed)) {
+      log.warn('Torch package list output was not an array.', { stdout, stderr });
+      return undefined;
+    }
+
+    const versions: TorchPackageVersions = {
+      torch: undefined,
+      torchaudio: undefined,
+      torchvision: undefined,
+    };
+    let matched = 0;
+    const isRecord = (value: unknown): value is Record<string, unknown> => typeof value === 'object' && value !== null;
+
+    for (const entry of parsed) {
+      if (!isRecord(entry)) continue;
+      const name = entry.name;
+      const version = entry.version;
+      if (typeof name !== 'string' || typeof version !== 'string') continue;
+      const packageName = TORCH_PACKAGE_NAMES.find((pkg) => pkg === name.trim().toLowerCase());
+      if (!packageName) continue;
+      matched += 1;
+      versions[packageName] = version.trim() || undefined;
+    }
+    if (matched === 0) {
+      log.warn('Torch package list did not contain expected packages.', { stdout, stderr });
+      return undefined;
+    }
+
+    return versions;
+  }
+
   /**
    * Installs AMD ROCm SDK packages on Windows.
    * @param callbacks The callbacks to use for the command.
@@ -838,10 +963,52 @@ export class VirtualEnvironment implements HasTelemetry, PythonExecutor {
       return 'package-upgrade';
     }
 
+    if (await this.needsNvidiaTorchUpgrade()) {
+      log.info('NVIDIA PyTorch version out of date. Treating as package upgrade.');
+      return 'package-upgrade';
+    }
+
     log.debug('hasRequirements result:', 'OK');
     return 'OK';
   }
 
+  /**
+   * Returns `true` when NVIDIA PyTorch should be upgraded to the recommended version.
+   * @returns `true` when NVIDIA PyTorch is out of date, otherwise `false`.
+   */
+  private async needsNvidiaTorchUpgrade(): Promise<boolean> {
+    if (this.selectedDevice !== 'nvidia') return false;
+
+    const installedVersions = await this.getInstalledTorchPackageVersions();
+    if (!installedVersions) {
+      log.warn('Unable to read NVIDIA torch package versions. Skipping NVIDIA torch upgrade check.');
+      return false;
+    }
+
+    return !this.meetsMinimumNvidiaTorchVersions(installedVersions);
+  }
+
+  private meetsMinimumNvidiaTorchVersions(installedVersions: TorchPackageVersions): boolean {
+    const torch = installedVersions.torch;
+    const torchaudio = installedVersions.torchaudio;
+    const torchvision = installedVersions.torchvision;
+    if (!torch || !torchaudio || !torchvision) return false;
+
+    const requiredCudaTag = '+cu130';
+    if (
+      !torch.includes(requiredCudaTag) ||
+      !torchaudio.includes(requiredCudaTag) ||
+      !torchvision.includes(requiredCudaTag)
+    )
+      return false;
+
+    if (compareVersions(torch, NVIDIA_TORCH_VERSION) < 0) return false;
+    if (compareVersions(torchaudio, NVIDIA_TORCH_VERSION) < 0) return false;
+    if (compareVersions(torchvision, NVIDIA_TORCHVISION_VERSION) < 0) return false;
+
+    return true;
+  }
+
   /**
    * Verifies that the Python environment can import modules that frequently show up in errors.
    * @returns `true` if the Python environment successfully imports the modules, otherwise `false`.