From 93a16cbd3e9a903b554af4d0ba940a01518e552d Mon Sep 17 00:00:00 2001
From: Ionite <dev@ionite.io>
Date: Thu, 9 May 2024 16:51:42 -0400
Subject: [PATCH] Remove preview torch from mps install

Stable supports MPS now, pre flag getting dev builds instead

(cherry picked from commit 6e079ae84e0523c360115fba8e9efaa993981544)
---
 StabilityMatrix.Core/Models/Packages/A3WebUI.cs           | 2 +-
 StabilityMatrix.Core/Models/Packages/ComfyUI.cs           | 3 +--
 StabilityMatrix.Core/Models/Packages/SDWebForge.cs        | 2 +-
 StabilityMatrix.Core/Models/Packages/StableDiffusionUx.cs | 2 +-
 4 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/StabilityMatrix.Core/Models/Packages/A3WebUI.cs b/StabilityMatrix.Core/Models/Packages/A3WebUI.cs
index 9e9b5b80..16dbf4f7 100644
--- a/StabilityMatrix.Core/Models/Packages/A3WebUI.cs
+++ b/StabilityMatrix.Core/Models/Packages/A3WebUI.cs
@@ -223,7 +223,7 @@ public class A3WebUI(
                     TorchVersion.Cpu => "cpu",
                     TorchVersion.Cuda => "cu121",
                     TorchVersion.Rocm => "rocm5.6",
-                    TorchVersion.Mps => "nightly/cpu",
+                    TorchVersion.Mps => "cpu",
                     _ => throw new ArgumentOutOfRangeException(nameof(torchVersion), torchVersion, null)
                 }
             )
diff --git a/StabilityMatrix.Core/Models/Packages/ComfyUI.cs b/StabilityMatrix.Core/Models/Packages/ComfyUI.cs
index 61c10f96..f172b29d 100644
--- a/StabilityMatrix.Core/Models/Packages/ComfyUI.cs
+++ b/StabilityMatrix.Core/Models/Packages/ComfyUI.cs
@@ -209,8 +209,7 @@ public class ComfyUI(
         pipArgs = torchVersion switch
         {
             TorchVersion.DirectMl => pipArgs.WithTorchDirectML(),
-            TorchVersion.Mps
-                => pipArgs.AddArg("--pre").WithTorch().WithTorchVision().WithTorchExtraIndex("nightly/cpu"),
+            TorchVersion.Mps => pipArgs.WithTorch().WithTorchVision().WithTorchExtraIndex("cpu"),
             _
                 => pipArgs
                     .AddArg("--upgrade")
diff --git a/StabilityMatrix.Core/Models/Packages/SDWebForge.cs b/StabilityMatrix.Core/Models/Packages/SDWebForge.cs
index 43fa5e22..7bda64fc 100644
--- a/StabilityMatrix.Core/Models/Packages/SDWebForge.cs
+++ b/StabilityMatrix.Core/Models/Packages/SDWebForge.cs
@@ -162,7 +162,7 @@ public class SDWebForge(
                         TorchVersion.Cpu => "cpu",
                         TorchVersion.Cuda => "cu121",
                         TorchVersion.Rocm => "rocm5.6",
-                        TorchVersion.Mps => "nightly/cpu",
+                        TorchVersion.Mps => "cpu",
                         _ => throw new ArgumentOutOfRangeException(nameof(torchVersion), torchVersion, null)
                     }
                 );
diff --git a/StabilityMatrix.Core/Models/Packages/StableDiffusionUx.cs b/StabilityMatrix.Core/Models/Packages/StableDiffusionUx.cs
index 5b5cacfa..b881efe4 100644
--- a/StabilityMatrix.Core/Models/Packages/StableDiffusionUx.cs
+++ b/StabilityMatrix.Core/Models/Packages/StableDiffusionUx.cs
@@ -211,7 +211,7 @@ public class StableDiffusionUx(
                         new PipInstallArgs()
                             .WithTorch("==2.1.2")
                             .WithTorchVision()
-                            .WithTorchExtraIndex("nightly/cpu"),
+                            .WithTorchExtraIndex("cpu"),
                         onConsoleOutput
                     )
                     .ConfigureAwait(false);