Comprehensive ROCm helper for centralized Windows package integration and installation support

StabilityMatrix.Core/Helper/Factory/PackageFactory.cs

@@ -4,6 +4,7 @@
 using StabilityMatrix.Core.Models.Packages;
 using StabilityMatrix.Core.Python;
 using StabilityMatrix.Core.Services;
+using StabilityMatrix.Core.Services.Rocm;
 
 namespace StabilityMatrix.Core.Helper.Factory;
 
@@ -18,6 +19,7 @@ public class PackageFactory : IPackageFactory
     private readonly IUvManager uvManager;
     private readonly IPyInstallationManager pyInstallationManager;
     private readonly IPipWheelService pipWheelService;
+    private readonly IRocmPackageHelper rocmPackageHelper;
 
     /// <summary>
     /// Mapping of package.Name to package
@@ -32,16 +34,20 @@ public PackageFactory(
         IPrerequisiteHelper prerequisiteHelper,
         IPyInstallationManager pyInstallationManager,
         IPyRunner pyRunner,
-        IPipWheelService pipWheelService
+        IUvManager uvManager,
+        IPipWheelService pipWheelService,
+        IRocmPackageHelper rocmPackageHelper
     )
     {
         this.githubApiCache = githubApiCache;
         this.settingsManager = settingsManager;
         this.downloadService = downloadService;
         this.prerequisiteHelper = prerequisiteHelper;
         this.pyRunner = pyRunner;
+        this.uvManager = uvManager;
         this.pyInstallationManager = pyInstallationManager;
         this.pipWheelService = pipWheelService;
+        this.rocmPackageHelper = rocmPackageHelper;
         this.basePackages = basePackages.ToDictionary(x => x.Name);
     }
 
@@ -55,7 +61,8 @@ public BasePackage GetNewBasePackage(InstalledPackage installedPackage)
                 downloadService,
                 prerequisiteHelper,
                 pyInstallationManager,
-                pipWheelService
+                pipWheelService,
+                rocmPackageHelper
             ),
             "Fooocus" => new Fooocus(
                 githubApiCache,

StabilityMatrix.Core/Helper/HardwareInfo/GpuInfo.cs

@@ -1,4 +1,6 @@
-namespace StabilityMatrix.Core.Helper.HardwareInfo;
+using StabilityMatrix.Core.Models.Rocm;
+
+namespace StabilityMatrix.Core.Helper.HardwareInfo;
 
 public record GpuInfo
 {
@@ -62,11 +64,7 @@ public bool IsLegacyNvidiaGpu()
 
     public bool IsWindowsRocmSupportedGpu()
     {
-        var gfx = GetAmdGfxArch();
-        if (gfx is null)
-            return false;
-
-        return gfx.StartsWith("gfx110") || gfx.StartsWith("gfx120") || gfx.Equals("gfx1151");
+        return WindowsRocmSupport.IsSupportedGpu(this);
     }
 
     public bool IsAmd => Name?.Contains("amd", StringComparison.OrdinalIgnoreCase) ?? false;
@@ -84,7 +82,7 @@ public bool IsWindowsRocmSupportedGpu()
         return name switch
         {
             // RDNA4
-            _ when Has("R9700") || Has("9070") => "gfx1201",
+            _ when Has("R9700") || Has("R9600") || Has("9070") => "gfx1201",
             _ when Has("9060") => "gfx1200",
 
             // RDNA3.5 APUs
@@ -112,6 +110,9 @@ _ when Has("660M") || Has("680M") => "gfx1035",
             _ when Has("6300") || Has("6400") || Has("6450") || Has("6500") || Has("6550") || Has("6500M") =>
                 "gfx1034",
 
+            // RDNA2 Steam Deck APU
+            _ when Has("Van Gogh") || Has("Sephiroth") => "gfx1033",
+
             // RDNA2 Navi23
             _ when Has("6600") || Has("6650") || Has("6700S") || Has("6800S") || Has("6600M") => "gfx1032",
 
@@ -121,6 +122,16 @@ _ when Has("6700") || Has("6750") || Has("6800M") || Has("6850M") => "gfx1031",
             // RDNA2 Navi21 (big die)
             _ when Has("6800") || Has("6900") || Has("6950") => "gfx1030",
 
+            // RDNA1 Navi10 XT (incl. Pro card)
+            _ when Has("5600") || Has("5700") || Has("v520") => "gfx1010",
+
+            // RDNA1 Navi10 XTX
+            _ when Has("5500") => "gfx1012",
+
+            // Vega/GCN5 Dedicated GPUs
+            _ when Has("pro vii") || HasNoSpace("provii") => "gfx90X",
+            _ when Has("rx vega") || Has("vega 64") || Has("vega 56") || Has("vega frontier") => "gfx900",
+            _ when Has("radeon vii") || HasNoSpace("radeonvii") => "gfx906",
             _ => null,
         };
 

StabilityMatrix.Core/Helper/HardwareInfo/HardwareHelper.cs

@@ -7,6 +7,7 @@
 using Microsoft.Win32;
 using NLog;
 using StabilityMatrix.Core.Extensions;
+using StabilityMatrix.Core.Models.Rocm;
 
 namespace StabilityMatrix.Core.Helper.HardwareInfo;
 
@@ -316,12 +317,11 @@ public static bool HasAmdGpu()
         return IterGpuInfo().Any(gpu => gpu.IsAmd);
     }
 
-    public static bool HasWindowsRocmSupportedGpu() =>
-        IterGpuInfo().Any(gpu => gpu is { IsAmd: true, Name: not null } && gpu.IsWindowsRocmSupportedGpu());
+    public static bool HasWindowsRocmSupportedGpu() => IterGpuInfo().Any(WindowsRocmSupport.IsSupportedGpu);
 
     public static GpuInfo? GetWindowsRocmSupportedGpu()
     {
-        return IterGpuInfo().FirstOrDefault(gpu => gpu.IsWindowsRocmSupportedGpu());
+        return IterGpuInfo().FirstOrDefault(WindowsRocmSupport.IsSupportedGpu);
     }
 
     public static bool HasIntelGpu() => IterGpuInfo().Any(gpu => gpu.IsIntel);

StabilityMatrix.Core/Models/Packages/ComfyUI.cs

@@ -13,9 +13,11 @@
 using StabilityMatrix.Core.Models.Packages.Config;
 using StabilityMatrix.Core.Models.Packages.Extensions;
 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;
 
@@ -26,7 +28,8 @@ public class ComfyUI(
     IDownloadService downloadService,
     IPrerequisiteHelper prerequisiteHelper,
     IPyInstallationManager pyInstallationManager,
-    IPipWheelService pipWheelService
+    IPipWheelService pipWheelService,
+    IRocmPackageHelper? rocmPackageHelper = null
 )
     : BaseGitPackage(
         githubApi,
@@ -38,6 +41,7 @@ IPipWheelService pipWheelService
     )
 {
     private static readonly Logger Logger = LogManager.GetCurrentClassLogger();
+
     public override string Name => "ComfyUI";
     public override string DisplayName { get; set; } = "ComfyUI";
     public override string Author => "comfyanonymous";
@@ -247,7 +251,7 @@ IPipWheelService pipWheelService
                 Name = "Enable DirectML",
                 Type = LaunchOptionType.Bool,
                 InitialValue =
-                    !HardwareHelper.HasWindowsRocmSupportedGpu()
+                    !HasWindowsRocmSupport()
                     && HardwareHelper.PreferDirectMLOrZluda()
                     && this is not ComfyZluda,
                 Options = ["--directml"],
@@ -264,7 +268,9 @@ IPipWheelService pipWheelService
             {
                 Name = "Cross Attention Method",
                 Type = LaunchOptionType.Bool,
-                InitialValue = "--use-pytorch-cross-attention",
+                InitialValue = DefaultToQuadCrossAttention()
+                    ? "--use-quad-cross-attention" // For Legacy AMD GPUs.
+                    : "--use-pytorch-cross-attention",
                 Options =
                 [
                     "--use-split-cross-attention",
@@ -362,69 +368,36 @@ public override async Task InstallPackage(
             .ConfigureAwait(false);
 
         var torchIndex = options.PythonOptions.TorchIndex ?? GetRecommendedTorchVersion();
-        var gfxArch =
-            SettingsManager.Settings.PreferredGpu?.GetAmdGfxArch()
-            ?? HardwareHelper.GetWindowsRocmSupportedGpu()?.GetAmdGfxArch();
-
-        // Special case for Windows ROCm Nightly builds
-        if (
-            Compat.IsWindows
-            && !string.IsNullOrWhiteSpace(gfxArch)
-            && torchIndex is TorchIndex.Rocm
-            && options.PythonOptions.PythonVersion >= PyVersion.Parse("3.11.0")
-        )
+        var isLegacyNvidia =
+            torchIndex == TorchIndex.Cuda
+            && (
+                SettingsManager.Settings.PreferredGpu?.IsLegacyNvidiaGpu()
+                ?? HardwareHelper.HasLegacyNvidiaGpu()
+            );
+
+        if (Compat.IsWindows && torchIndex == TorchIndex.Rocm && HasWindowsRocmSupport())
         {
-            var config = new PipInstallConfig
+            if (rocmPackageHelper is null)
             {
-                RequirementsFilePaths = ["requirements.txt"],
-                ExtraPipArgs = ["numpy<2"],
-                SkipTorchInstall = true,
-                PostInstallPipArgs = ["typing-extensions>=4.15.0"],
-            };
-            await StandardPipInstallProcessAsync(
+                throw new InvalidOperationException(
+                    "Windows ROCm installation requires the shared ROCm helper to resolve gfx-specific index URLs."
+                );
+            }
+
+            await rocmPackageHelper
+                .InstallWindowsNativePackageAsync(
                     venvRunner,
-                    options,
+                    installLocation,
                     installedPackage,
-                    config,
-                    onConsoleOutput,
+                    WindowsRocmProfile,
                     progress,
+                    onConsoleOutput,
                     cancellationToken
                 )
                 .ConfigureAwait(false);
-
-            progress?.Report(
-                new ProgressReport(-1f, "Installing ROCm nightly torch...", isIndeterminate: true)
-            );
-            var indexUrl = gfxArch switch
-            {
-                "gfx1150" => "https://rocm.nightlies.amd.com/v2-staging/gfx1150", // Strix/Gorgon Point
-                "gfx1151" => "https://rocm.nightlies.amd.com/v2/gfx1151", // Strix Halo
-                _ when gfxArch.StartsWith("gfx110") => "https://rocm.nightlies.amd.com/v2/gfx110X-all",
-                _ when gfxArch.StartsWith("gfx120") => "https://rocm.nightlies.amd.com/v2/gfx120X-all",
-                _ => throw new ArgumentOutOfRangeException(
-                    nameof(gfxArch),
-                    $"Unsupported GFX Arch: {gfxArch}"
-                ),
-            };
-
-            var torchPipArgs = new PipInstallArgs()
-                .AddArgs("--pre", "--upgrade")
-                .WithTorch()
-                .WithTorchVision()
-                .WithTorchAudio()
-                .AddArgs("--index-url", indexUrl);
-
-            await venvRunner.PipInstall(torchPipArgs, onConsoleOutput).ConfigureAwait(false);
         }
-        else // Standard installation path for all other cases
+        else
         {
-            var isLegacyNvidia =
-                torchIndex == TorchIndex.Cuda
-                && (
-                    SettingsManager.Settings.PreferredGpu?.IsLegacyNvidiaGpu()
-                    ?? HardwareHelper.HasLegacyNvidiaGpu()
-                );
-
             var config = new PipInstallConfig
             {
                 RequirementsFilePaths = ["requirements.txt"],
@@ -479,7 +452,11 @@ await StandardPipInstallProcessAsync(
                             SettingsManager.Settings.PreferredGpu?.IsBlackwellGpu()
                             ?? HardwareHelper.HasBlackwellGpu(),
                         WorkingDirectory = installLocation,
-                        EnvironmentVariables = GetEnvVars(venvRunner.EnvironmentVariables),
+                        EnvironmentVariables = GetEnvVars(
+                            venvRunner.EnvironmentVariables,
+                            installLocation,
+                            installedPackage
+                        ),
                     };
 
                     await step.ExecuteAsync(progress).ConfigureAwait(false);
@@ -529,7 +506,7 @@ public override async Task RunPackage(
         await SetupVenv(installLocation, pythonVersion: PyVersion.Parse(installedPackage.PythonVersion))
             .ConfigureAwait(false);
 
-        VenvRunner.UpdateEnvironmentVariables(GetEnvVars);
+        VenvRunner.UpdateEnvironmentVariables(env => GetEnvVars(env, installLocation, installedPackage));
 
         // Check for old NVIDIA driver version with cu130 installations
         var isNvidia = SettingsManager.Settings.PreferredGpu?.IsNvidia ?? HardwareHelper.HasNvidiaGpu();
@@ -613,13 +590,7 @@ public override TorchIndex GetRecommendedTorchVersion()
     {
         var preferRocm =
             (Compat.IsLinux && (SettingsManager.Settings.PreferredGpu?.IsAmd ?? HardwareHelper.PreferRocm()))
-            || (
-                Compat.IsWindows
-                && (
-                    SettingsManager.Settings.PreferredGpu?.IsWindowsRocmSupportedGpu()
-                    ?? HardwareHelper.HasWindowsRocmSupportedGpu()
-                )
-            );
+            || HasWindowsRocmSupport();
 
         if (AvailableTorchIndices.Contains(TorchIndex.Rocm) && preferRocm)
         {
@@ -629,6 +600,54 @@ public override TorchIndex GetRecommendedTorchVersion()
         return base.GetRecommendedTorchVersion();
     }
 
+    /// Windows ROCm install profile for ComfyUI.
+    private static readonly RocmPackageProfile WindowsRocmProfile = new()
+    {
+        RequiresRocmSdk = true,
+        ExtraInstallPipArgs = ["numpy<2"],
+        PostInstallPipArgs = ["typing-extensions>=4.15.0"],
+        UpgradePackages = true,
+        ExtraEnvironmentFactory = BuildComfyWindowsRocmEnvironment,
+    };
+
+    private static IReadOnlyDictionary<string, string> BuildComfyWindowsRocmEnvironment(
+        RocmRuntimeContext runtimeContext
+    )
+    {
+        return WindowsRocmSupport.IsModernArchitecture(runtimeContext.RuntimeGfxArch)
+            ? new Dictionary<string, string> { ["COMFYUI_ENABLE_MIOPEN"] = "1" }
+            : new Dictionary<string, string>();
+    }
+
+    /// Uses the shared ROCm helper for Windows ROCm eligibility checks so ComfyUI does not maintain its own support matrix.
+    private bool HasWindowsRocmSupport()
+    {
+        if (!Compat.IsWindows)
+            return false;
+
+        if (rocmPackageHelper is null)
+            return false;
+
+        var compatibility = rocmPackageHelper.GetCompatibility(WindowsRocmProfile);
+
+        return compatibility.IsCompatible;
+    }
+
+    /// Defaults legacy Windows ROCm GPUs to quad cross-attention because PyTorch cross-attention is considerably slower
+    /// and not as supported on older AMD architectures.
+    private bool DefaultToQuadCrossAttention()
+    {
+        if (!Compat.IsWindows || !HasWindowsRocmSupport())
+            return false;
+
+        var gpu = SettingsManager.Settings.PreferredGpu;
+        var gfxArch = WindowsRocmSupport.IsSupportedGpu(gpu)
+            ? gpu?.GetAmdGfxArch()
+            : HardwareHelper.GetWindowsRocmSupportedGpu()?.GetAmdGfxArch();
+
+        return WindowsRocmSupport.PreferLegacyAttentionFallback(gfxArch);
+    }
+
     public override IPackageExtensionManager ExtensionManager =>
         new ComfyExtensionManager(this, settingsManager);
 
@@ -979,21 +998,29 @@ await PipWheelService
             .ConfigureAwait(false);
     }
 
-    private ImmutableDictionary<string, string> GetEnvVars(ImmutableDictionary<string, string> env)
+    private ImmutableDictionary<string, string> GetEnvVars(
+        ImmutableDictionary<string, string> env,
+        string installLocation,
+        InstalledPackage installedPackage
+    )
     {
         // if we're not on windows or we don't have a windows rocm gpu, return original env
-        var hasRocmGpu =
-            SettingsManager.Settings.PreferredGpu?.IsWindowsRocmSupportedGpu()
-            ?? HardwareHelper.HasWindowsRocmSupportedGpu();
+        var hasRocmGpu = HasWindowsRocmSupport();
 
         if (!Compat.IsWindows || !hasRocmGpu)
             return env;
 
-        // set some experimental speed improving env vars for Windows ROCm
-        return env.SetItem("PYTORCH_TUNABLEOP_ENABLED", "1")
-            .SetItem("MIOPEN_FIND_MODE", "2")
-            .SetItem("TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL", "1")
-            .SetItem("PYTORCH_ALLOC_CONF", "max_split_size_mb:6144,garbage_collection_threshold:0.8") // greatly helps prevent GPU OOM and instability/driver timeouts/OS hard locks and decreases dependency on Tiled VAE at standard res's
-            .SetItem("COMFYUI_ENABLE_MIOPEN", "1"); // re-enables "cudnn" in ComfyUI as it's needed for MiOpen to function properly
+        if (rocmPackageHelper is not null)
+        {
+            var rocmEnvironment = rocmPackageHelper.BuildLaunchEnvironment(
+                installLocation,
+                installedPackage,
+                WindowsRocmProfile
+            );
+
+            return env.SetItems(rocmEnvironment);
+        }
+
+        return env;
     }
 }