diff --git a/StabilityMatrix.Core/Helper/Factory/PackageFactory.cs b/StabilityMatrix.Core/Helper/Factory/PackageFactory.cs index ca1b68828..2d3d674c5 100644 --- a/StabilityMatrix.Core/Helper/Factory/PackageFactory.cs +++ b/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, diff --git a/StabilityMatrix.Core/Models/Packages/ForgeClassic.cs b/StabilityMatrix.Core/Models/Packages/ForgeClassic.cs index f9365ab75..f106e85e1 100644 --- a/StabilityMatrix.Core/Models/Packages/ForgeClassic.cs +++ b/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)); } + /// + /// 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. + /// + protected virtual Task PrepareRocmInstallAsync( + IPyVenvRunner venvRunner, + InstalledPackage installedPackage, + InstallPackageOptions options, + Action? onConsoleOutput, + IProgress? progress, + CancellationToken cancellationToken + ) + { + return Task.CompletedTask; + } + private async Task RunInstallScriptWithPromptHandling( IPyVenvRunner venvRunner, IReadOnlyCollection launchArgs, diff --git a/StabilityMatrix.Core/Models/Packages/ForgeNeo.cs b/StabilityMatrix.Core/Models/Packages/ForgeNeo.cs index ee8f4f9f2..0dc0c29eb 100644 --- a/StabilityMatrix.Core/Models/Packages/ForgeNeo.cs +++ b/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 + { + ["TORCH_INDEX_URL"] = LinuxRocmIndexUrl, + ["TORCH_COMMAND"] = $"pip install --index-url {LinuxRocmIndexUrl} torch torchvision", + }, + }; + + /// + /// Windows installs use AMD repo wheels; Linux uses the upstream PyTorch ROCm index. + /// + 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 AvailableTorchIndices => [TorchIndex.Cuda, TorchIndex.Rocm]; + + public override bool IsCompatible => HardwareHelper.HasNvidiaGpu() || HasRocmSupport(); + + /// + /// 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. + /// + 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 LaunchOptions { get @@ -45,10 +102,122 @@ public override List LaunchOptions Options = ["--bnb"], } ); + + ForgeNeoWindowsRocmProfile.Default.ApplyWindowsRocmLaunchDefaults(options, rocmPackageHelper); + return options; } } + protected override async Task PrepareRocmInstallAsync( + IPyVenvRunner venvRunner, + InstalledPackage installedPackage, + InstallPackageOptions options, + Action? onConsoleOutput, + IProgress? 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 GetEnvVars( + ImmutableDictionary 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 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."; } diff --git a/StabilityMatrix.Core/Models/Rocm/ForgeNeoWindowsRocmProfile.cs b/StabilityMatrix.Core/Models/Rocm/ForgeNeoWindowsRocmProfile.cs new file mode 100644 index 000000000..69d5e56b3 --- /dev/null +++ b/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; + +/// +/// 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. +/// +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 BuildEnvironment(RocmRuntimeContext runtimeContext) + { + var indexUrl = WindowsRocmSupport.GetMultiArchPythonPackageIndexUrl(runtimeContext.RuntimeGfxArch); + var deviceExtra = WindowsRocmSupport.TryGetMultiArchDeviceExtra(runtimeContext.RuntimeGfxArch); + + if (string.IsNullOrWhiteSpace(deviceExtra)) + { + return new Dictionary(); + } + + // 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 + { + ["TORCH_INDEX_URL"] = indexUrl, + ["TORCH_COMMAND"] = torchCommand, + }; + } + + public override void ApplyWindowsRocmLaunchDefaults( + List 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 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); + } +} diff --git a/docs/advanced/hardware-support.md b/docs/advanced/hardware-support.md index 87b3621f3..8af0d18ca 100644 --- a/docs/advanced/hardware-support.md +++ b/docs/advanced/hardware-support.md @@ -58,7 +58,7 @@ Architectures in the `gfx110x`, `gfx115x`, and `gfx120x` ranges are treated as " - Torch is installed from AMD's ROCm multi-arch index (`repo.amd.com/rocm/whl-multi-arch/`) as device-specific wheels (`torch[device-gfxNNNN]`). Vega parts (`gfx900` / `gfx906`) pull from the 'TheRock' nightly multi-arch feed instead, since these architecture builds currently are only available there instead of the stable production distribution stream. - On modern architectures it applies a set of ROCm performance and attention environment variables at launch (MIOpen find-mode tuning, AOTriton experimental flash attention, `COMFYUI_ENABLE_MIOPEN`, and an allocator tuning string). AOTriton is excluded on the `gfx1152` / `gfx1153` APU architectures, which it does not yet support. Legacy architectures instead force a math SDP fallback. The full variable list and exactly which ones are auto-applied are documented in [Environment Variables](environment-variables.md#amd-and-rocm-variables). - ComfyUI offers optional extra commands for supported AMD GPUs, including **Install Triton and SageAttention (ROCm)** (Sage Attention 1.x), **Install Flash Attention (ROCm)** (legacy architectures), an **Install ROCm Development SDK** step, and an **Install bitsandbytes (ROCm)** step for Python 3.12 environments. -- **Packages:** ComfyUI, Stable Diffusion WebUI Reforge, InvokeAI, SwarmUI, and Wan2GP are supported by this install path. +- **Packages:** ComfyUI, Stable Diffusion WebUI Reforge, Stable Diffusion WebUI Forge - Neo, InvokeAI, SwarmUI, and Wan2GP are supported by this install path. > [!NOTE] While not managed by Stability Matrix for ROCm installs, SD.Next has a Windows-native ROCm install when the "ROCm" Pytorch index is selected in Advanced Installation Options during initial package install and `--use-rocm` is set in launch options. This install path is internally handled by SD.Next itself and currently only supports RDNA2 dedicated GPUs, RDNA3, RDNA3.5, and RDNA4 GPUs. **Caveats:** @@ -100,7 +100,10 @@ On Linux, AMD GPUs use native ROCm directly, which is the mature AMD path. - **Caveats:** - Native ROCm on Linux depends on a system-level ROCm installation and a compatible kernel/driver stack, which Stability Matrix does not install for you. - The Windows-only ROCm performance environment overrides described above are not auto-applied on Linux, so if you want them you can set them yourself via the [Environment Variables](environment-variables.md) editor. -- **Packages:** ComfyUI, Stable Diffusion WebUI, SD.Next, Stable Diffusion WebUI Forge, InvokeAI, SwarmUI, SDFX, OneTrainer, and Wan2GP list ROCm support. +- **Packages:** ComfyUI, Stable Diffusion WebUI, SD.Next, Stable Diffusion WebUI Forge, Stable Diffusion WebUI Forge - Neo, InvokeAI, SwarmUI, SDFX, OneTrainer, and Wan2GP list ROCm support. + + > [!NOTE] + > Forge Neo's upstream project officially supports NVIDIA CUDA only, so its ROCm support here is added by Stability Matrix and is not supported by the upstream developers. ## Intel (IPEX)