Handle NVIDIA PyTorch updates

src/constants.ts

@@ -181,6 +181,14 @@ export enum TorchMirrorUrl {
   NightlyCpu = 'https://download.pytorch.org/whl/nightly/cpu',
 }
 
+export type TorchUpdatePolicy = 'auto' | 'defer' | 'pinned';
+
+export type TorchPinnedPackages = {
+  torch?: string;
+  torchaudio?: string;
+  torchvision?: string;
+};
+
 /** Legacy NVIDIA torch mirror used by older installs (CUDA 12.9). */
 export const LEGACY_NVIDIA_TORCH_MIRROR = 'https://download.pytorch.org/whl/cu129';
 
@@ -210,6 +218,8 @@ export const NVIDIA_TORCH_PACKAGES: string[] = [
   `torchaudio==${NVIDIA_TORCH_VERSION}`,
   `torchvision==${NVIDIA_TORCHVISION_VERSION}`,
 ];
+/** Recommended NVIDIA torch package set key (torch/torchaudio/torchvision). */
+export const NVIDIA_TORCH_RECOMMENDED_VERSION = `${NVIDIA_TORCH_VERSION}|${NVIDIA_TORCHVISION_VERSION}`;
 
 /** The log files used by the desktop process. */
 export enum LogFile {

src/install/installationManager.ts

@@ -5,7 +5,14 @@ import { promisify } from 'node:util';
 
 import { strictIpcMain as ipcMain } from '@/infrastructure/ipcChannels';
 
-import { IPC_CHANNELS, InstallStage, ProgressStatus } from '../constants';
+import {
+  IPC_CHANNELS,
+  InstallStage,
+  NVIDIA_TORCHVISION_VERSION,
+  NVIDIA_TORCH_RECOMMENDED_VERSION,
+  NVIDIA_TORCH_VERSION,
+  ProgressStatus,
+} from '../constants';
 import { PythonImportVerificationError } from '../infrastructure/pythonImportVerificationError';
 import { useAppState } from '../main-process/appState';
 import type { AppWindow } from '../main-process/appWindow';
@@ -91,7 +98,7 @@ export class InstallationManager implements HasTelemetry {
     // Convert from old format
     if (state === 'upgraded') installation.upgradeConfig();
 
-    // Install updated manager requirements
+    // Install updated requirements
     if (installation.needsRequirementsUpdate) await this.updatePackages(installation);
 
     // Resolve issues and re-run validation
@@ -382,14 +389,120 @@ export class InstallationManager implements HasTelemetry {
       await installation.virtualEnvironment.installComfyUIRequirements(callbacks);
       await installation.virtualEnvironment.installComfyUIManagerRequirements(callbacks);
       await this.warnIfNvidiaDriverTooOld(installation);
-      await installation.virtualEnvironment.ensureRecommendedNvidiaTorch(callbacks);
+      await this.maybeUpdateNvidiaTorch(installation, callbacks);
       await installation.validate();
     } catch (error) {
       log.error('Error auto-updating packages:', error);
       await this.appWindow.loadPage('server-start');
     }
   }
 
+  private async maybeUpdateNvidiaTorch(installation: ComfyInstallation, callbacks: ProcessCallbacks): Promise<void> {
+    const virtualEnvironment = installation.virtualEnvironment;
+    if (virtualEnvironment.selectedDevice !== 'nvidia') return;
+
+    const config = useDesktopConfig();
+    const updatePolicy = config.get('torchUpdatePolicy');
+    if (updatePolicy === 'pinned') {
+      log.info('Skipping NVIDIA PyTorch update because updates are pinned.');
+      return;
+    }
+
+    const installedVersions = await virtualEnvironment.getInstalledTorchPackageVersions();
+    if (!installedVersions) {
+      log.warn('Skipping NVIDIA PyTorch update because installed versions could not be read.');
+      return;
+    }
+
+    const isOutOfDate = await virtualEnvironment.isNvidiaTorchOutOfDate(installedVersions);
+    if (!isOutOfDate) return;
+
+    const recommendedVersion = NVIDIA_TORCH_RECOMMENDED_VERSION;
+    const lastPromptedVersion = config.get('torchLastPromptedVersion');
+
+    if (config.get('torchOutOfDateRecommendedVersion') !== recommendedVersion) {
+      config.set('torchOutOfDateRecommendedVersion', recommendedVersion);
+      config.set('torchOutOfDatePackages', installedVersions);
+    }
+
+    if (updatePolicy === 'defer' && lastPromptedVersion === recommendedVersion) {
+      log.info('Skipping NVIDIA PyTorch update because updates are deferred for this version.');
+      return;
+    }
+
+    const autoPolicy = updatePolicy === undefined || updatePolicy === 'auto';
+    const updateApproved = autoPolicy && lastPromptedVersion === recommendedVersion;
+
+    if (!updateApproved) {
+      const currentTorch = installedVersions.torch ?? 'unknown';
+      const currentTorchaudio = installedVersions.torchaudio ?? 'unknown';
+      const currentTorchvision = installedVersions.torchvision ?? 'unknown';
+
+      const { response } = await this.appWindow.showMessageBox({
+        type: 'question',
+        title: 'Update PyTorch?',
+        message:
+          'Your NVIDIA PyTorch build is out of date. We can update to the recommended build for performance improvements. This may change memory use and/or compatibility with custom nodes',
+        detail: [
+          `Current: torch ${currentTorch}, torchaudio ${currentTorchaudio}, torchvision ${currentTorchvision}`,
+          `Recommended: torch ${NVIDIA_TORCH_VERSION}, torchaudio ${NVIDIA_TORCH_VERSION}, torchvision ${NVIDIA_TORCHVISION_VERSION}`,
+        ].join('\n'),
+        buttons: ['Update PyTorch', 'Ask again later', 'Silence until next version', 'Silence forever'],
+        defaultId: 0,
+        cancelId: 1,
+      });
+
+      switch (response) {
+        case 1:
+          log.info('Deferring NVIDIA PyTorch update prompt.');
+          return;
+        case 2:
+          config.set('torchLastPromptedVersion', recommendedVersion);
+          config.set('torchUpdatePolicy', 'defer');
+          config.delete('torchPinnedPackages');
+          virtualEnvironment.updateTorchUpdatePolicy('defer', undefined, recommendedVersion);
+          return;
+        case 3:
+          config.set('torchLastPromptedVersion', recommendedVersion);
+          config.set('torchUpdatePolicy', 'pinned');
+          config.set('torchPinnedPackages', installedVersions);
+          virtualEnvironment.updateTorchUpdatePolicy('pinned', installedVersions, recommendedVersion);
+          return;
+        default:
+          config.set('torchLastPromptedVersion', recommendedVersion);
+          config.set('torchUpdatePolicy', 'auto');
+          config.delete('torchPinnedPackages');
+          config.delete('torchUpdateFailureSilencedVersion');
+          virtualEnvironment.updateTorchUpdatePolicy('auto', undefined, recommendedVersion);
+      }
+    } else {
+      virtualEnvironment.updateTorchUpdatePolicy('auto', undefined, recommendedVersion);
+    }
+
+    try {
+      await virtualEnvironment.ensureRecommendedNvidiaTorch(callbacks);
+      config.delete('torchUpdateFailureSilencedVersion');
+    } catch (error) {
+      log.error('Error updating NVIDIA PyTorch packages:', error);
+      if (config.get('torchUpdateFailureSilencedVersion') === recommendedVersion) return;
+
+      const { response } = await this.appWindow.showMessageBox({
+        type: 'warning',
+        title: 'PyTorch update failed',
+        message:
+          'We could not install the recommended NVIDIA PyTorch build. This may be because your configured torch mirror does not provide it.',
+        detail: 'We will retry the update on each startup.',
+        buttons: ['OK', "Don't show again"],
+        defaultId: 0,
+        cancelId: 0,
+      });
+
+      if (response === 1) {
+        config.set('torchUpdateFailureSilencedVersion', recommendedVersion);
+      }
+    }
+  }
+
   /**
    * Warns the user if their NVIDIA driver is too old for the required CUDA build.
    * @param installation The current installation.

src/main-process/comfyInstallation.ts

@@ -72,6 +72,9 @@ export class ComfyInstallation {
       pythonMirror: useComfySettings().get('Comfy-Desktop.UV.PythonInstallMirror'),
       pypiMirror: useComfySettings().get('Comfy-Desktop.UV.PypiInstallMirror'),
       torchMirror: useComfySettings().get('Comfy-Desktop.UV.TorchInstallMirror'),
+      torchUpdatePolicy: useDesktopConfig().get('torchUpdatePolicy'),
+      torchPinnedPackages: useDesktopConfig().get('torchPinnedPackages'),
+      torchUpdateDecisionVersion: useDesktopConfig().get('torchLastPromptedVersion'),
     });
   }
 

src/store/desktopSettings.ts

@@ -1,3 +1,4 @@
+import type { TorchPinnedPackages, TorchUpdatePolicy } from '../constants';
 import type { GpuType, TorchDeviceType } from '../preload';
 
 export type DesktopInstallState = 'started' | 'installed' | 'upgraded';
@@ -35,4 +36,16 @@ export type DesktopSettings = {
   versionConsentedMetrics?: string;
   /** Whether the user has generated an image successfully. */
   hasGeneratedSuccessfully?: boolean;
+  /** How to handle NVIDIA PyTorch updates. */
+  torchUpdatePolicy?: TorchUpdatePolicy;
+  /** The pinned NVIDIA torch package versions when updates are disabled. */
+  torchPinnedPackages?: TorchPinnedPackages;
+  /** The recommended NVIDIA torch version tied to the current update decision. */
+  torchLastPromptedVersion?: string;
+  /** The recommended NVIDIA torch version whose update failure prompt is suppressed. */
+  torchUpdateFailureSilencedVersion?: string;
+  /** The recommended NVIDIA torch version recorded when we first detected an out-of-date torch install. */
+  torchOutOfDateRecommendedVersion?: string;
+  /** The torch package versions recorded when we first detected an out-of-date torch install. */
+  torchOutOfDatePackages?: TorchPinnedPackages;
 };

src/virtualEnvironment.ts

@@ -14,8 +14,11 @@ import {
   LEGACY_NVIDIA_TORCH_MIRROR,
   NVIDIA_TORCHVISION_VERSION,
   NVIDIA_TORCH_PACKAGES,
+  NVIDIA_TORCH_RECOMMENDED_VERSION,
   NVIDIA_TORCH_VERSION,
   TorchMirrorUrl,
+  TorchPinnedPackages,
+  TorchUpdatePolicy,
 } from './constants';
 import { PythonImportVerificationError } from './infrastructure/pythonImportVerificationError';
 import { useAppState } from './main-process/appState';
@@ -89,13 +92,7 @@ function getPipInstallArgs(config: PipInstallConfig): string[] {
   return installArgs;
 }
 
-/**
- * Returns the default torch mirror for the given device.
- * @param device The device type
- * @returns The default torch mirror
- */
-function getDefaultTorchMirror(device: TorchDeviceType): string {
-  log.debug('Falling back to default torch mirror');
+function getDeviceDefaultTorchMirror(device: TorchDeviceType): string {
   switch (device) {
     case 'mps':
       return TorchMirrorUrl.NightlyCpu;
@@ -140,6 +137,9 @@ export class VirtualEnvironment implements HasTelemetry, PythonExecutor {
   readonly pythonMirror?: string;
   readonly pypiMirror?: string;
   readonly torchMirror?: string;
+  torchUpdatePolicy?: TorchUpdatePolicy;
+  torchPinnedPackages?: TorchPinnedPackages;
+  torchUpdateDecisionVersion?: string;
   uvPty: pty.IPty | undefined;
 
   /** The environment variables to set for uv. */
@@ -184,13 +184,19 @@ export class VirtualEnvironment implements HasTelemetry, PythonExecutor {
       pythonMirror,
       pypiMirror,
       torchMirror,
+      torchUpdatePolicy,
+      torchPinnedPackages,
+      torchUpdateDecisionVersion,
     }: {
       telemetry: ITelemetry;
       selectedDevice?: TorchDeviceType;
       pythonVersion?: string;
       pythonMirror?: string;
       pypiMirror?: string;
       torchMirror?: string;
+      torchUpdatePolicy?: TorchUpdatePolicy;
+      torchPinnedPackages?: TorchPinnedPackages;
+      torchUpdateDecisionVersion?: string;
     }
   ) {
     this.basePath = basePath;
@@ -200,6 +206,9 @@ export class VirtualEnvironment implements HasTelemetry, PythonExecutor {
     this.pythonMirror = pythonMirror;
     this.pypiMirror = pypiMirror;
     this.torchMirror = fixDeviceMirrorMismatch(selectedDevice!, torchMirror);
+    this.torchUpdatePolicy = torchUpdatePolicy;
+    this.torchPinnedPackages = torchPinnedPackages;
+    this.torchUpdateDecisionVersion = torchUpdateDecisionVersion;
 
     // uv defaults to .venv
     this.venvPath = path.join(basePath, '.venv');
@@ -263,6 +272,34 @@ export class VirtualEnvironment implements HasTelemetry, PythonExecutor {
     return primary;
   }
 
+  updateTorchUpdatePolicy(
+    policy: TorchUpdatePolicy | undefined,
+    pinnedPackages?: TorchPinnedPackages,
+    decisionVersion?: string
+  ) {
+    this.torchUpdatePolicy = policy;
+    if (pinnedPackages !== undefined) {
+      this.torchPinnedPackages = pinnedPackages;
+    } else if (policy !== 'pinned') {
+      this.torchPinnedPackages = undefined;
+    }
+    if (decisionVersion !== undefined) {
+      this.torchUpdateDecisionVersion = decisionVersion;
+    }
+  }
+
+  isUsingCustomTorchMirror(): boolean {
+    if (!this.torchMirror) return false;
+    return this.torchMirror !== getDeviceDefaultTorchMirror(this.selectedDevice);
+  }
+
+  private shouldSkipNvidiaTorchUpgrade(): boolean {
+    if (this.torchUpdatePolicy === 'pinned') return true;
+    if (this.torchUpdatePolicy === 'defer' && this.torchUpdateDecisionVersion === NVIDIA_TORCH_RECOMMENDED_VERSION)
+      return true;
+    return false;
+  }
+
   public async create(callbacks?: ProcessCallbacks): Promise<void> {
     try {
       await this.createEnvironment(callbacks);
@@ -629,7 +666,11 @@ export class VirtualEnvironment implements HasTelemetry, PythonExecutor {
       return;
     }
 
-    const torchMirror = this.torchMirror || getDefaultTorchMirror(this.selectedDevice);
+    let torchMirror = this.torchMirror;
+    if (!torchMirror) {
+      log.info('Falling back to default torch mirror');
+      torchMirror = getDeviceDefaultTorchMirror(this.selectedDevice);
+    }
     const config: PipInstallConfig = {
       packages: ['torch', 'torchvision', 'torchaudio'],
       indexUrl: torchMirror,
@@ -652,14 +693,22 @@ export class VirtualEnvironment implements HasTelemetry, PythonExecutor {
    */
   async ensureRecommendedNvidiaTorch(callbacks?: ProcessCallbacks): Promise<void> {
     if (this.selectedDevice !== 'nvidia') return;
+    if (this.shouldSkipNvidiaTorchUpgrade()) {
+      log.info('Skipping NVIDIA PyTorch upgrade due to pinned policy or deferred updates.');
+      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);
+    let torchMirror = this.torchMirror;
+    if (!torchMirror) {
+      log.info('Falling back to default torch mirror');
+      torchMirror = getDeviceDefaultTorchMirror(this.selectedDevice);
+    }
     const config: PipInstallConfig = {
       packages: NVIDIA_TORCH_PACKAGES,
       indexUrl: torchMirror,
@@ -695,7 +744,7 @@ export class VirtualEnvironment implements HasTelemetry, PythonExecutor {
    * 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> {
+  async getInstalledTorchPackageVersions(): Promise<TorchPackageVersions | undefined> {
     let stdout = '';
     let stderr = '';
     const callbacks: ProcessCallbacks = {
@@ -758,6 +807,15 @@ export class VirtualEnvironment implements HasTelemetry, PythonExecutor {
     return versions;
   }
 
+  async isNvidiaTorchOutOfDate(installedVersions?: TorchPackageVersions): Promise<boolean> {
+    if (this.selectedDevice !== 'nvidia') return false;
+
+    const resolvedVersions = installedVersions ?? (await this.getInstalledTorchPackageVersions());
+    if (!resolvedVersions) return false;
+
+    return !this.meetsMinimumNvidiaTorchVersions(resolvedVersions);
+  }
+
   /**
    * Installs AMD ROCm SDK packages on Windows.
    * @param callbacks The callbacks to use for the command.
@@ -982,6 +1040,7 @@ export class VirtualEnvironment implements HasTelemetry, PythonExecutor {
    */
   private async needsNvidiaTorchUpgrade(): Promise<boolean> {
     if (this.selectedDevice !== 'nvidia') return false;
+    if (this.shouldSkipNvidiaTorchUpgrade()) return false;
 
     const installedVersions = await this.getInstalledTorchPackageVersions();
     if (!installedVersions) {