Add AMD ROCm support to Forge Neo (Windows and Linux)

StabilityMatrix.Core/Helper/Factory/PackageFactory.cs

@@ -260,7 +260,8 @@ public BasePackage GetNewBasePackage(InstalledPackage installedPackage)
                 downloadService,
                 prerequisiteHelper,
                 pyInstallationManager,
-                pipWheelService
+                pipWheelService,
+                rocmPackageHelper
             ),
             "framepack" => new FramePack(
                 githubApiCache,

StabilityMatrix.Core/Models/Packages/ForgeClassic.cs

@@ -219,6 +219,18 @@ public override async Task InstallPackage(
             )
             .ConfigureAwait(false);
 
+        // Allow derived packages (e.g. Forge Neo on ROCm) to prepare the environment before the
+        // upstream launch.py install runs, such as pre-installing an alternative torch build.
+        await PrepareRocmInstallAsync(
+                venvRunner,
+                installedPackage,
+                options,
+                onConsoleOutput,
+                progress,
+                cancellationToken
+            )
+            .ConfigureAwait(false);
+
         progress?.Report(new ProgressReport(-1f, "Running install script...", isIndeterminate: true));
 
         // Build args for their launch.py - use --uv for fast installs, --exit to quit after setup
@@ -293,6 +305,23 @@ public override async Task InstallPackage(
         progress?.Report(new ProgressReport(1f, "Install complete", isIndeterminate: false));
     }
 
+    /// <summary>
+    /// Hook for derived packages to prepare the environment before the upstream Forge install script
+    /// runs. The default implementation does nothing. Used by Forge Neo to pre-install an alternative
+    /// ROCm torch build so the upstream launch.py does not install its default (CUDA) torch build.
+    /// </summary>
+    protected virtual Task PrepareRocmInstallAsync(
+        IPyVenvRunner venvRunner,
+        InstalledPackage installedPackage,
+        InstallPackageOptions options,
+        Action<ProcessOutput>? onConsoleOutput,
+        IProgress<ProgressReport>? progress,
+        CancellationToken cancellationToken
+    )
+    {
+        return Task.CompletedTask;
+    }
+
     private async Task<int> RunInstallScriptWithPromptHandling(
         IPyVenvRunner venvRunner,
         IReadOnlyCollection<string> launchArgs,

StabilityMatrix.Core/Models/Packages/ForgeNeo.cs

@@ -1,9 +1,15 @@
-using Injectio.Attributes;
+using System.Collections.Immutable;
+using Injectio.Attributes;
 using StabilityMatrix.Core.Helper;
 using StabilityMatrix.Core.Helper.Cache;
+using StabilityMatrix.Core.Helper.HardwareInfo;
 using StabilityMatrix.Core.Models;
+using StabilityMatrix.Core.Models.Progress;
+using StabilityMatrix.Core.Models.Rocm;
+using StabilityMatrix.Core.Processes;
 using StabilityMatrix.Core.Python;
 using StabilityMatrix.Core.Services;
+using StabilityMatrix.Core.Services.Rocm;
 
 namespace StabilityMatrix.Core.Models.Packages;
 
@@ -14,7 +20,8 @@ public class ForgeNeo(
     IDownloadService downloadService,
     IPrerequisiteHelper prerequisiteHelper,
     IPyInstallationManager pyInstallationManager,
-    IPipWheelService pipWheelService
+    IPipWheelService pipWheelService,
+    IRocmPackageHelper rocmPackageHelper
 )
     : ForgeClassic(
         githubApi,
@@ -25,10 +32,60 @@ IPipWheelService pipWheelService
         pipWheelService
     )
 {
+    internal const string LinuxRocmIndexUrl = "https://download.pytorch.org/whl/rocm7.14";
+
+    // Linux uses the upstream PyTorch ROCm index rather than AMD repo wheels. Kept inline because it
+    // collapses into the helper once its OS-agnostic path lands.
+    private static readonly RocmPackageProfile LinuxRocmProfile = new()
+    {
+        ExtraEnvironmentFactory = _ => new Dictionary<string, string>
+        {
+            ["TORCH_INDEX_URL"] = LinuxRocmIndexUrl,
+            ["TORCH_COMMAND"] = $"pip install --index-url {LinuxRocmIndexUrl} torch torchvision",
+        },
+    };
+
+    /// <summary>
+    /// Windows installs use AMD repo wheels; Linux uses the upstream PyTorch ROCm index.
+    /// </summary>
+    private RocmPackageProfile ActiveRocmProfile =>
+        Compat.IsWindows ? ForgeNeoWindowsRocmProfile.Default : LinuxRocmProfile;
+
     public override string Name => "forge-neo";
     public override string DisplayName { get; set; } = "Stable Diffusion WebUI Forge - Neo";
     public override string MainBranch => "neo";
     public override PackageType PackageType => PackageType.SdInference;
+
+    // CUDA for NVIDIA; ROCm for AMD.
+    public override IEnumerable<TorchIndex> AvailableTorchIndices => [TorchIndex.Cuda, TorchIndex.Rocm];
+
+    public override bool IsCompatible => HardwareHelper.HasNvidiaGpu() || HasRocmSupport();
+
+    /// <summary>
+    /// Uses the shared ROCm helper for ROCm eligibility checks so Forge Neo does not maintain its own
+    /// support matrix. The helper resolves compatibility for both Windows and Linux.
+    /// </summary>
+    private bool HasRocmSupport()
+    {
+        return rocmPackageHelper.GetCompatibility().IsCompatible;
+    }
+
+    public override TorchIndex GetRecommendedTorchVersion()
+    {
+        var preferNvidia = SettingsManager.Settings.PreferredGpu?.IsNvidia ?? HardwareHelper.HasNvidiaGpu();
+        if (preferNvidia)
+        {
+            return TorchIndex.Cuda;
+        }
+
+        if (HasRocmSupport())
+        {
+            return TorchIndex.Rocm;
+        }
+
+        return base.GetRecommendedTorchVersion();
+    }
+
     public override List<LaunchOptionDefinition> LaunchOptions
     {
         get
@@ -45,10 +102,122 @@ public override List<LaunchOptionDefinition> LaunchOptions
                     Options = ["--bnb"],
                 }
             );
+
+            ForgeNeoWindowsRocmProfile.Default.ApplyWindowsRocmLaunchDefaults(options, rocmPackageHelper);
+
             return options;
         }
     }
 
+    protected override async Task PrepareRocmInstallAsync(
+        IPyVenvRunner venvRunner,
+        InstalledPackage installedPackage,
+        InstallPackageOptions options,
+        Action<ProcessOutput>? onConsoleOutput,
+        IProgress<ProgressReport>? progress,
+        CancellationToken cancellationToken
+    )
+    {
+        var torchIndex = options.PythonOptions.TorchIndex ?? GetRecommendedTorchVersion();
+        if (torchIndex != TorchIndex.Rocm)
+        {
+            return;
+        }
+
+        if (!HasRocmSupport())
+        {
+            throw new InvalidOperationException(
+                rocmPackageHelper.GetCompatibility().FailureReason
+                    ?? "The current machine is not compatible with the selected ROCm torch index."
+            );
+        }
+
+        var profile = ActiveRocmProfile;
+
+        progress?.Report(new ProgressReport(-1f, "Preparing ROCm environment...", isIndeterminate: true));
+
+        // Apply the ROCm launch environment (including the TORCH_COMMAND/TORCH_INDEX_URL override) so both
+        // the torch install and the subsequent upstream launch.py see it.
+        venvRunner.UpdateEnvironmentVariables(env =>
+            env.SetItems(rocmPackageHelper.BuildLaunchEnvironment(profile))
+        );
+
+        if (Compat.IsWindows)
+        {
+            // Install the AMD multi-arch torch build before launch.py runs, so its is_installed("torch")
+            // gate skips the default NVIDIA CUDA torch install.
+            await rocmPackageHelper
+                .InstallWindowsNativeTorchAsync(
+                    venvRunner,
+                    installedPackage,
+                    profile,
+                    progress,
+                    onConsoleOutput,
+                    cancellationToken
+                )
+                .ConfigureAwait(false);
+            return;
+        }
+
+        // Linux: install from the upstream PyTorch ROCm index. An exclusive --index-url (rather than
+        // --extra-index-url) keeps pip from quietly resolving the CPU build from PyPI.
+        var torchArgs = new PipInstallArgs()
+            .AddKeyedArgs("--index-url", ["--index-url", LinuxRocmIndexUrl])
+            .AddArg("torch")
+            .AddArg("torchvision");
+
+        await venvRunner.PipInstall(torchArgs, onConsoleOutput).ConfigureAwait(false);
+
+        await VerifyLinuxRocmTorchInstallAsync(venvRunner, cancellationToken).ConfigureAwait(false);
+    }
+
+    private static async Task VerifyLinuxRocmTorchInstallAsync(
+        IPyVenvRunner venvRunner,
+        CancellationToken cancellationToken
+    )
+    {
+        cancellationToken.ThrowIfCancellationRequested();
+
+        var torchInfo = await venvRunner.PipShow("torch").ConfigureAwait(false);
+        if (torchInfo is null)
+        {
+            throw new InvalidOperationException("torch was not installed after Linux ROCm setup.");
+        }
+
+        if (
+            string.IsNullOrWhiteSpace(torchInfo.Version)
+            || !torchInfo.Version.Contains("rocm", StringComparison.OrdinalIgnoreCase)
+        )
+        {
+            throw new InvalidOperationException(
+                $"Installed torch is not a ROCm build (detected version: {torchInfo.Version})."
+            );
+        }
+    }
+
+    protected override ImmutableDictionary<string, string> GetEnvVars(
+        ImmutableDictionary<string, string> env,
+        InstalledPackage installedPackage
+    )
+    {
+        env = base.GetEnvVars(env, installedPackage);
+
+        var selectedTorchIndex = installedPackage.PreferredTorchIndex ?? GetRecommendedTorchVersion();
+        if (selectedTorchIndex != TorchIndex.Rocm || !HasRocmSupport())
+        {
+            return env;
+        }
+
+        return env.SetItems(rocmPackageHelper.BuildLaunchEnvironment(ActiveRocmProfile));
+    }
+
+    protected override IReadOnlyList<string> GetLaunchNoticeLines(InstalledPackage installedPackage)
+    {
+        var selectedTorchIndex = installedPackage.PreferredTorchIndex ?? GetRecommendedTorchVersion();
+        // The helper returns an empty list on non-Windows platforms.
+        return rocmPackageHelper.GetWindowsLaunchNoticeLines(selectedTorchIndex);
+    }
+
     public override string Blurb =>
         "Neo mainly serves as an continuation for the \"latest\" version of Forge. Additionally, this fork is focused on optimization and usability, with the main goal of being the lightest WebUI without any bloatwares.";
 }

StabilityMatrix.Core/Models/Rocm/ForgeNeoWindowsRocmProfile.cs

@@ -0,0 +1,121 @@
+using StabilityMatrix.Core.Helper;
+using StabilityMatrix.Core.Models.Packages;
+using StabilityMatrix.Core.Services.Rocm;
+
+namespace StabilityMatrix.Core.Models.Rocm;
+
+/// <summary>
+/// Windows ROCm profile for Forge Neo.
+/// Forge Neo has no native AMD support, so the helper owns the ROCm torch install and the launch
+/// environment, while Forge Neo's own launch.py still installs the rest of its dependencies.
+/// </summary>
+public class ForgeNeoWindowsRocmProfile : RocmPackageProfile
+{
+    private const string DisableSmartMemoryOptionName = "Disable Smart Memory";
+    private const string UsePytorchCrossAttentionOptionName = "Use PyTorch Cross Attention";
+
+    // Recommended launch defaults from the "AMD Forge Neo with ROCm" setup guide.
+    private static readonly string[] EnableLaunchOptionNames = ["Pin Shared Memory", "CUDA Stream"];
+    private static readonly string[] DisableLaunchOptionNames = ["CUDA Malloc"];
+
+    public ForgeNeoWindowsRocmProfile()
+    {
+        InstallConfig = new PipInstallConfig
+        {
+            // Forge Neo installs its own requirements via launch.py; the helper only needs to add the
+            // ROCm-oriented accelerators that upstream would otherwise pull as CUDA-only builds.
+            PostTorchInstallPipArgs = ["bitsandbytes", "triton-windows"],
+            UpgradePackages = true,
+        };
+
+        ExtraEnvironmentFactory = BuildEnvironment;
+    }
+
+    public static ForgeNeoWindowsRocmProfile Default { get; } = new ForgeNeoWindowsRocmProfile();
+
+    private static IReadOnlyDictionary<string, string> BuildEnvironment(RocmRuntimeContext runtimeContext)
+    {
+        var indexUrl = WindowsRocmSupport.GetMultiArchPythonPackageIndexUrl(runtimeContext.RuntimeGfxArch);
+        var deviceExtra = WindowsRocmSupport.TryGetMultiArchDeviceExtra(runtimeContext.RuntimeGfxArch);
+
+        if (string.IsNullOrWhiteSpace(deviceExtra))
+        {
+            return new Dictionary<string, string>();
+        }
+
+        // Persisted so that if Forge Neo re-runs its own torch install step it targets AMD's index
+        // instead of pulling the NVIDIA CUDA build from PyTorch.
+        var torchCommand =
+            $"pip install --index-url {indexUrl} "
+            + $"\"torch[{deviceExtra}]\" \"torchvision[{deviceExtra}]\" torchaudio \"rocm[devel]\"";
+
+        return new Dictionary<string, string>
+        {
+            ["TORCH_INDEX_URL"] = indexUrl,
+            ["TORCH_COMMAND"] = torchCommand,
+        };
+    }
+
+    public override void ApplyWindowsRocmLaunchDefaults(
+        List<LaunchOptionDefinition> launchOptions,
+        IRocmPackageHelper rocmPackageHelper
+    )
+    {
+        if (!(Compat.IsWindows && rocmPackageHelper.GetCompatibility().IsCompatible))
+        {
+            return;
+        }
+
+        for (var i = 0; i < launchOptions.Count; i++)
+        {
+            var option = launchOptions[i];
+
+            if (EnableLaunchOptionNames.Contains(option.Name))
+            {
+                launchOptions[i] = option with { InitialValue = true };
+            }
+            else if (DisableLaunchOptionNames.Contains(option.Name))
+            {
+                launchOptions[i] = option with { InitialValue = null };
+            }
+        }
+
+        InsertBoolOptionIfMissing(
+            launchOptions,
+            DisableSmartMemoryOptionName,
+            "--disable-smart-memory",
+            "Aggressively offload to RAM instead of keeping models in VRAM when possible"
+        );
+        InsertBoolOptionIfMissing(
+            launchOptions,
+            UsePytorchCrossAttentionOptionName,
+            "--use-pytorch-cross-attention",
+            "Use the PyTorch cross attention (override sageattention/flash_attn/xformers)"
+        );
+    }
+
+    private static void InsertBoolOptionIfMissing(
+        List<LaunchOptionDefinition> launchOptions,
+        string name,
+        string option,
+        string description
+    )
+    {
+        if (launchOptions.Any(x => x.Name == name))
+        {
+            return;
+        }
+
+        var definition = new LaunchOptionDefinition
+        {
+            Name = name,
+            Type = LaunchOptionType.Bool,
+            Description = description,
+            Options = [option],
+            InitialValue = true,
+        };
+
+        var extrasIndex = launchOptions.FindIndex(x => x.Name == LaunchOptionDefinition.Extras.Name);
+        launchOptions.Insert(extrasIndex < 0 ? launchOptions.Count : extrasIndex, definition);
+    }
+}