Add ComfyUI

StabilityMatrix/App.xaml.cs

@@ -81,6 +81,7 @@ namespace StabilityMatrix
             
             serviceCollection.AddSingleton<BasePackage, A3WebUI>();
             serviceCollection.AddSingleton<BasePackage, VladAutomatic>();
+            serviceCollection.AddSingleton<BasePackage, ComfyUI>();
             serviceCollection.AddSingleton<ISnackbarService, SnackbarService>();
             serviceCollection.AddSingleton<ISettingsManager, SettingsManager>();
             serviceCollection.AddSingleton<IPrerequisiteHelper, PrerequisiteHelper>();

StabilityMatrix/Models/Packages/BaseGitPackage.cs

@@ -165,7 +165,7 @@ public abstract class BaseGitPackage : BasePackage
         return version;
     }
     
-    private void UnzipPackage(bool isUpdate = false)
+    public void UnzipPackage(bool isUpdate = false)
     {
         if (isUpdate)
         {

StabilityMatrix/Models/Packages/ComfyUI.cs

@@ -0,0 +1,125 @@
+using System;
+using System.Collections.Generic;
+using System.Diagnostics;
+using System.IO;
+using System.Linq;
+using System.Text.RegularExpressions;
+using System.Threading.Tasks;
+using NLog;
+using StabilityMatrix.Helper;
+using StabilityMatrix.Helper.Cache;
+using StabilityMatrix.Python;
+
+namespace StabilityMatrix.Models.Packages;
+
+public class ComfyUI : BaseGitPackage
+{
+    private static readonly Logger Logger = LogManager.GetCurrentClassLogger();
+    public override string Name => "ComfyUI";
+    public override string DisplayName { get; set; } = "ComfyUI";
+    public override string Author => "comfyanonymous";
+    public override string LaunchCommand => "main.py";
+
+    public ComfyUI(IGithubApiCache githubApi, ISettingsManager settingsManager) : base(githubApi, settingsManager)
+    {
+    }
+
+    public override List<LaunchOptionDefinition> LaunchOptions => new()
+    {
+        new()
+        {
+            Name = "VRAM",
+            Options = new() { "--lowvram", "--medvram" }
+        },
+        new()
+        {
+            Name = "Xformers",
+            Options = new() { "--xformers" }
+        },
+        new()
+        {
+            Name = "API",
+            Options = new() { "--api" }
+        },
+        LaunchOptionDefinition.Extras
+    };
+
+    public override Task<string> GetLatestVersion() => Task.FromResult("master");
+
+    public override async Task<IEnumerable<PackageVersion>> GetAllVersions(bool isReleaseMode = true)
+    {
+        var allBranches = await GetAllBranches();
+        return allBranches.Select(b => new PackageVersion
+        {
+            TagName = $"{b.Name}", 
+            ReleaseNotesMarkdown = string.Empty
+        });
+    }
+
+    public override async Task InstallPackage(bool isUpdate = false)
+    {
+        UnzipPackage(isUpdate);
+        
+        if (isUpdate)
+        {
+            OnUpdateComplete("Update complete");
+            return;
+        }
+        // On first install, need to setup dependencies
+        OnInstallProgressChanged(-1); // Indeterminate progress bar
+        // Setup venv
+        Logger.Debug("Setting up venv");
+        await SetupVenv(InstallLocation);
+        var venvRunner = new PyVenvRunner(Path.Combine(InstallLocation, "venv"));
+        
+        void OnConsoleOutput(string? s)
+        {
+            Debug.WriteLine($"venv stdout: {s}");
+        }
+        
+        // Install torch
+        Logger.Debug("Starting torch install...");
+        await venvRunner.PipInstall(venvRunner.GetTorchInstallCommand(), InstallLocation, OnConsoleOutput);
+        // Install requirements
+        Logger.Debug("Starting requirements install...");
+        await venvRunner.PipInstall("-r requirements.txt", InstallLocation, OnConsoleOutput);
+        
+        Logger.Debug("Finished installing requirements!");
+        OnInstallComplete("Installation complete");
+    }
+    
+    public override async Task RunPackage(string installedPackagePath, string arguments)
+    {
+        await SetupVenv(installedPackagePath);
+
+        void HandleConsoleOutput(string? s)
+        {
+            if (s == null) return;
+
+            if (s.Contains("To see the GUI go to", StringComparison.OrdinalIgnoreCase))
+            {
+                var regex = new Regex(
+                    @"(https?:\/\/)([^:\s]+):(\d+)");
+                var match = regex.Match(s);
+                if (match.Success)
+                {
+                    WebUrl = match.Value;
+                }
+                OnStartupComplete(WebUrl);
+            }
+
+            Debug.WriteLine($"process stdout: {s}");
+            OnConsoleOutput($"{s}\n");
+        }
+
+        void HandleExit(int i)
+        {
+            Debug.WriteLine($"Venv process exited with code {i}");
+            OnExit(i);
+        }
+
+        var args = $"\"{Path.Combine(installedPackagePath, LaunchCommand)}\" {arguments}";
+
+        VenvRunner?.RunDetached(args.TrimEnd(), HandleConsoleOutput, HandleExit, workingDirectory: installedPackagePath);
+    }
+}

StabilityMatrix/ViewModels/InstallerViewModel.cs

@@ -221,20 +221,22 @@ public partial class InstallerViewModel : ObservableObject
         SelectedPackage.InstallLocation = $"{InstallPath}\\{InstallName}";
         SelectedPackage.DisplayName = InstallName;
 
-        var version = await DownloadPackage(SelectedVersion.TagName);
-        await InstallPackage();
-
-        ProgressText = "Installing dependencies...";
-        await pyRunner.Initialize();
-        if (!PyRunner.PipInstalled)
+        if (!PyRunner.PipInstalled || !PyRunner.VenvInstalled)
         {
-            await pyRunner.SetupPip();
+            ProgressText = "Installing dependencies...";
+            await pyRunner.Initialize();
+            if (!PyRunner.PipInstalled)
+            {
+                await pyRunner.SetupPip();
+            }
+            if (!PyRunner.VenvInstalled)
+            {
+                await pyRunner.InstallPackage("virtualenv");
+            }
         }
 
-        if (!PyRunner.VenvInstalled)
-        {
-            await pyRunner.InstallPackage("virtualenv");
-        }
+        var version = await DownloadPackage(SelectedVersion.TagName);
+        await InstallPackage();
 
         ProgressText = "Done";