From 8d74f966ae5cc03465aff9959be2424226891ab7 Mon Sep 17 00:00:00 2001
From: Robin Huang <robin.j.huang@gmail.com>
Date: Fri, 15 Nov 2024 18:02:40 -0800
Subject: [PATCH] Fixed: Install pytorch separately in fallback installation.

---
 src/virtualEnvironment.ts | 34 ++++++++++++++++++++++++++++++----
 1 file changed, 30 insertions(+), 4 deletions(-)

diff --git a/src/virtualEnvironment.ts b/src/virtualEnvironment.ts
index a36bf25e..4ff78c8e 100644
--- a/src/virtualEnvironment.ts
+++ b/src/virtualEnvironment.ts
@@ -214,13 +214,39 @@ export class VirtualEnvironment {
   }
 
   private async installComfyUIRequirements(callbacks?: ProcessCallbacks): Promise<void> {
-    log.info(`Installing ComfyUI requirements from ${this.comfyUIRequirementsPath}`);
-    const installCmd = ['pip', 'install', '-r', this.comfyUIRequirementsPath];
-
     if (process.platform === 'win32') {
-      installCmd.push('--index-url', 'https://download.pytorch.org/whl/cu121');
+      log.info('Installing PyTorch CUDA 12.1');
+      await this.runUvCommandAsync(
+        [
+          'pip',
+          'install',
+          'torch',
+          'torchvision',
+          'torchaudio',
+          '--index-url',
+          'https://download.pytorch.org/whl/cu121',
+        ],
+        callbacks
+      );
     }
 
+    if (process.platform === 'darwin') {
+      log.info('Installing PyTorch Nightly for macOS.');
+      await this.runUvCommandAsync(
+        [
+          'pip',
+          'install',
+          'torch',
+          'torchvision',
+          'torchaudio',
+          '--extra-index-url',
+          'https://download.pytorch.org/whl/nightly/cpu',
+        ],
+        callbacks
+      );
+    }
+    log.info(`Installing ComfyUI requirements from ${this.comfyUIRequirementsPath}`);
+    const installCmd = ['pip', 'install', '-r', this.comfyUIRequirementsPath];
     const { exitCode } = await this.runUvCommandAsync(installCmd, callbacks);
     if (exitCode !== 0) {
       throw new Error(`Failed to install requirements.txt: exit code ${exitCode}`);