Merge pull request #1157 from ionite34/comfy-rocm-env-vars

Use recommended env vars for windows+rocm comfyUI

(cherry picked from commit 6b0f954e7977472edd56e92242a8500a7132efe1)

Author: mohnjiles

CHANGELOG.md

@@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning 2.0](https://semver.org/spec/v2
 - Changed ComfyUI torch index from `cu129` back to `cu128` for better compatibility with custom nodes
 - Updated SageAttention installer to install v2.2.0-windows.post3
 - Updated Nunchaku installer to install v1.0.1
+- Updated Windows ROCm ComfyUI installs to use recommended environment variables by default
 ### Fixed
 - Fixed [#1372](https://github.com/LykosAI/StabilityMatrix/issues/1372), [#1399](https://github.com/LykosAI/StabilityMatrix/issues/1399) - LiteAsyncException upon starting Stability Matrix v2.15.0
 - Fixed [#1391](https://github.com/LykosAI/StabilityMatrix/issues/1391) - "Failed to parse" error when upgrading pip packages with extra index url

StabilityMatrix.Core/Models/Packages/ComfyUI.cs

@@ -1,4 +1,5 @@
-using System.Text.Json;
+using System.Collections.Immutable;
+using System.Text.Json;
 using System.Text.RegularExpressions;
 using Injectio.Attributes;
 using NLog;
@@ -443,6 +444,8 @@ public override async Task RunPackage(
         await SetupVenv(installLocation, pythonVersion: PyVersion.Parse(installedPackage.PythonVersion))
             .ConfigureAwait(false);
 
+        VenvRunner.UpdateEnvironmentVariables(GetEnvVars);
+
         VenvRunner.RunDetached(
             [Path.Combine(installLocation, options.Command ?? LaunchCommand), .. options.Arguments],
             HandleConsoleOutput,
@@ -808,4 +811,20 @@ private async Task InstallNunchaku(InstalledPackage? installedPackage)
         EventManager.Instance.OnPackageInstallProgressAdded(runner);
         await runner.ExecuteSteps([installNunchaku]).ConfigureAwait(false);
     }
+
+    private ImmutableDictionary<string, string> GetEnvVars(ImmutableDictionary<string, string> env)
+    {
+        // 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();
+
+        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");
+    }
 }