Merge pull request LykosAI#1644 from NeuralFault/remove-normalvram

ComfyUI: Remove `--normalvram`

Author: mohnjiles

StabilityMatrix.Core/Models/Packages/ComfyUI.cs

@@ -221,10 +221,9 @@ IPipWheelService pipWheelService
                 InitialValue = HardwareHelper.IterGpuInfo().Select(gpu => gpu.MemoryLevel).Max() switch
                 {
                     MemoryLevel.Low => "--lowvram",
-                    MemoryLevel.Medium => "--normalvram",
                     _ => null,
                 },
-                Options = ["--highvram", "--normalvram", "--lowvram", "--novram"],
+                Options = ["--highvram", "--lowvram", "--novram"],
             },
             new()
             {
@@ -529,6 +528,8 @@ public override async Task RunPackage(
         await SetupVenv(installLocation, pythonVersion: PyVersion.Parse(installedPackage.PythonVersion))
             .ConfigureAwait(false);
 
+        var launchArguments = NormalizeLaunchArguments(installedPackage, options.Arguments);
+
         VenvRunner.UpdateEnvironmentVariables(GetEnvVars);
 
         // Check for old NVIDIA driver version with cu130 installations
@@ -585,7 +586,7 @@ older torch index (e.g. cu128)
         }
 
         VenvRunner.RunDetached(
-            [Path.Combine(installLocation, options.Command ?? LaunchCommand), .. options.Arguments],
+            [Path.Combine(installLocation, options.Command ?? LaunchCommand), .. launchArguments],
             HandleConsoleOutput,
             OnExit
         );
@@ -609,6 +610,34 @@ void HandleConsoleOutput(ProcessOutput s)
         }
     }
 
+    protected ProcessArgs NormalizeLaunchArguments(
+        InstalledPackage installedPackage,
+        ProcessArgs fallbackArguments
+    )
+    {
+        if (installedPackage.LaunchArgs is not { Count: > 0 })
+        {
+            return fallbackArguments;
+        }
+
+        var removedCount = installedPackage.LaunchArgs.RemoveAll(option =>
+            string.Equals(option.Name, "--normalvram", StringComparison.OrdinalIgnoreCase)
+        );
+
+        if (removedCount == 0)
+        {
+            return fallbackArguments;
+        }
+
+        Logger.Info("Removed {RemovedCount} obsolete ComfyUI launch args before launch", removedCount);
+
+        SettingsManager.SaveLaunchArgs(installedPackage.Id, installedPackage.LaunchArgs);
+
+        return ProcessArgs.FromQuoted(
+            installedPackage.LaunchArgs.Select(option => option.ToArgString()).OfType<string>()
+        );
+    }
+
     public override TorchIndex GetRecommendedTorchVersion()
     {
         var preferRocm =

StabilityMatrix.Core/Models/Packages/ComfyZluda.cs

@@ -205,8 +205,10 @@ public override async Task RunPackage(
         await SetupVenv(installLocation, pythonVersion: PyVersion.Parse(installedPackage.PythonVersion))
             .ConfigureAwait(false);
 
+        var launchArguments = NormalizeLaunchArguments(installedPackage, options.Arguments);
+
         var zludaPath = Path.Combine(installLocation, LaunchCommand);
-        ProcessArgs args = ["--", VenvRunner.PythonPath.ToString(), "main.py", .. options.Arguments];
+        ProcessArgs args = ["--", VenvRunner.PythonPath.ToString(), "main.py", .. launchArguments];
         zludaProcess = ProcessRunner.StartAnsiProcess(
             zludaPath,
             args,

StabilityMatrix.Tests/Models/Packages/ComfyLaunchArgMigrationTests.cs

@@ -0,0 +1,86 @@
+using NSubstitute;
+using StabilityMatrix.Core.Helper;
+using StabilityMatrix.Core.Helper.Cache;
+using StabilityMatrix.Core.Models;
+using StabilityMatrix.Core.Models.Packages;
+using StabilityMatrix.Core.Processes;
+using StabilityMatrix.Core.Python;
+using StabilityMatrix.Core.Services;
+
+namespace StabilityMatrix.Tests.Models.Packages;
+
+[TestClass]
+public class ComfyLaunchArgMigrationTests
+{
+    [TestMethod]
+    public void NormalizeLaunchArguments_StripsObsoleteNormalVramAndPersistsUpdatedArgs()
+    {
+        var settingsManager = Substitute.For<ISettingsManager>();
+        var package = new TestComfyUI(
+            Substitute.For<IGithubApiCache>(),
+            settingsManager,
+            Substitute.For<IDownloadService>(),
+            Substitute.For<IPrerequisiteHelper>(),
+            Substitute.For<IPyInstallationManager>(),
+            Substitute.For<IPipWheelService>()
+        );
+
+        var installedPackage = new InstalledPackage
+        {
+            Id = Guid.NewGuid(),
+            PackageName = "ComfyUI",
+            LaunchArgs =
+            [
+                new LaunchOption
+                {
+                    Name = "--normalvram",
+                    Type = LaunchOptionType.Bool,
+                    OptionValue = true,
+                },
+            ],
+        };
+
+        var fallbackArguments = ProcessArgs.FromQuoted(
+            installedPackage
+                .LaunchArgs.Select(option => option.ToArgString())
+                .Where(argument => argument is not null)
+                .Select(argument => argument!)
+        );
+
+        var normalizedArguments = package.Normalize(installedPackage, fallbackArguments);
+
+        Assert.IsFalse(installedPackage.LaunchArgs.Any(option => option.Name == "--normalvram"));
+        Assert.IsFalse(normalizedArguments.Contains("--normalvram"));
+        settingsManager
+            .Received(1)
+            .SaveLaunchArgs(
+                installedPackage.Id,
+                Arg.Is<IEnumerable<LaunchOption>>(options =>
+                    options.All(option => option.Name != "--normalvram")
+                )
+            );
+    }
+
+    private sealed class TestComfyUI : ComfyUI
+    {
+        public TestComfyUI(
+            IGithubApiCache githubApi,
+            ISettingsManager settingsManager,
+            IDownloadService downloadService,
+            IPrerequisiteHelper prerequisiteHelper,
+            IPyInstallationManager pyInstallationManager,
+            IPipWheelService pipWheelService
+        )
+            : base(
+                githubApi,
+                settingsManager,
+                downloadService,
+                prerequisiteHelper,
+                pyInstallationManager,
+                pipWheelService
+            ) { }
+
+        public ProcessArgs Normalize(InstalledPackage installedPackage, ProcessArgs fallbackArguments) =>
+            NormalizeLaunchArguments(installedPackage, fallbackArguments);
+    }
+}