Merge pull request LykosAI#996 from ionite34/more-blackwell-updates

Use cu128 nightly torch index instead of w-e-w wheels

(cherry picked from commit 7ab289c)

# Conflicts:
#	StabilityMatrix.Core/Models/Packages/SDWebForge.cs

Author: mohnjiles

StabilityMatrix.Core/Models/Packages/ComfyUI.cs

@@ -210,15 +210,23 @@ public override async Task InstallPackage(
         var isBlackwell =
             SettingsManager.Settings.PreferredGpu?.IsBlackwellGpu() ?? HardwareHelper.HasBlackwellGpu();
 
-        if (isBlackwell)
+        if (isBlackwell && torchVersion is TorchIndex.Cuda)
         {
             pipArgs = pipArgs
-                .AddArg(
-                    "https://huggingface.co/w-e-w/torch-2.6.0-cu128.nv/resolve/main/torch-2.6.0%2Bcu128.nv-cp310-cp310-win_amd64.whl"
-                )
-                .AddArg(
-                    "https://huggingface.co/w-e-w/torch-2.6.0-cu128.nv/resolve/main/torchvision-0.20.0a0%2Bcu128.nv-cp310-cp310-win_amd64.whl"
-                );
+                .AddArg("--upgrade")
+                .AddArg("--pre")
+                .WithTorch()
+                .WithTorchVision()
+                .WithTorchExtraIndex("nightly/cu128");
+
+            if (installedPackage.PipOverrides != null)
+            {
+                pipArgs = pipArgs.WithUserOverrides(installedPackage.PipOverrides);
+            }
+            progress?.Report(
+                new ProgressReport(-1f, "Installing Torch for your shiny new GPU...", isIndeterminate: true)
+            );
+            await venvRunner.PipInstall(pipArgs, onConsoleOutput).ConfigureAwait(false);
         }
         else
         {
@@ -248,6 +256,11 @@ public override async Task InstallPackage(
             };
         }
 
+        if (isBlackwell && torchVersion is TorchIndex.Cuda)
+        {
+            pipArgs = new PipInstallArgs();
+        }
+
         var requirements = new FilePath(installLocation, "requirements.txt");
 
         pipArgs = pipArgs.WithParsedFromRequirementsTxt(

StabilityMatrix.Core/Models/Packages/SDWebForge.cs

@@ -149,7 +149,25 @@ public override async Task InstallPackage(
         var torchVersion = options.PythonOptions.TorchIndex ?? GetRecommendedTorchVersion();
         if (torchVersion is TorchIndex.DirectMl)
         {
+<<<<<<< HEAD
             pipArgs = pipArgs.WithTorchDirectML();
+=======
+            pipArgs = pipArgs
+                .AddArg("--upgrade")
+                .AddArg("--pre")
+                .WithTorch()
+                .WithTorchVision()
+                .WithTorchExtraIndex("nightly/cu128");
+
+            if (installedPackage.PipOverrides != null)
+            {
+                pipArgs = pipArgs.WithUserOverrides(installedPackage.PipOverrides);
+            }
+            progress?.Report(
+                new ProgressReport(-1f, "Installing Torch for your shiny new GPU...", isIndeterminate: true)
+            );
+            await venvRunner.PipInstall(pipArgs, onConsoleOutput).ConfigureAwait(false);
+>>>>>>> 7ab289c8 (Merge pull request #996 from ionite34/more-blackwell-updates)
         }
         else
         {
@@ -168,6 +186,11 @@ public override async Task InstallPackage(
                 );
         }
 
+        if (isBlackwell && torchVersion is TorchIndex.Cuda)
+        {
+            pipArgs = new PipInstallArgs();
+        }
+
         pipArgs = pipArgs.WithParsedFromRequirementsTxt(requirementsContent, excludePattern: "torch");
 
         if (installedPackage.PipOverrides != null)