From faa7b7c75ca95a95c6c62cfd3853ae1b0c6dd68b Mon Sep 17 00:00:00 2001 From: JT Date: Mon, 14 Aug 2023 20:45:27 -0700 Subject: [PATCH] Added shared model folder handling for InvokeAI & made the packages handle their own RemoveModelFolderLinks --- StabilityMatrix.Core/Helper/SharedFolders.cs | 10 +- .../Models/Packages/A3WebUI.cs | 6 ++ .../Models/Packages/BasePackage.cs | 1 + .../Models/Packages/ComfyUI.cs | 3 + .../Models/Packages/DankDiffusion.cs | 5 + .../Models/Packages/InvokeAI.cs | 96 +++++++++++++++++-- .../Models/Packages/VladAutomatic.cs | 9 +- .../Models/Packages/VoltaML.cs | 6 ++ 8 files changed, 120 insertions(+), 16 deletions(-) diff --git a/StabilityMatrix.Core/Helper/SharedFolders.cs b/StabilityMatrix.Core/Helper/SharedFolders.cs index 52e4fe7a..fa2278cd 100644 --- a/StabilityMatrix.Core/Helper/SharedFolders.cs +++ b/StabilityMatrix.Core/Helper/SharedFolders.cs @@ -1,4 +1,5 @@ -using NLog; +using AsyncAwaitBestPractices; +using NLog; using StabilityMatrix.Core.Extensions; using StabilityMatrix.Core.Helper.Factory; using StabilityMatrix.Core.Models; @@ -147,7 +148,7 @@ public class SharedFolders : ISharedFolders } } - private static void RemoveLinksForPackage(BasePackage package, DirectoryPath installPath) + public static void RemoveLinksForPackage(BasePackage package, DirectoryPath installPath) { var sharedFolders = package.SharedFolders; if (sharedFolders == null) @@ -171,7 +172,6 @@ public class SharedFolders : ISharedFolders public void RemoveLinksForAllPackages() { - var libraryDir = settingsManager.LibraryDir; var packages = settingsManager.Settings.InstalledPackages; foreach (var package in packages) { @@ -182,11 +182,11 @@ public class SharedFolders : ISharedFolders try { - RemoveLinksForPackage(basePackage, Path.Combine(libraryDir, package.LibraryPath)); + basePackage.RemoveModelFolderLinks(package.FullPath).GetAwaiter().GetResult(); } catch (Exception e) { - Logger.Warn("Failed to remove junction links for package {Package} " + + Logger.Warn("Failed to remove links for package {Package} " + "({DisplayName}): {Message}", package.PackageName, package.DisplayName, e.Message); } } diff --git a/StabilityMatrix.Core/Models/Packages/A3WebUI.cs b/StabilityMatrix.Core/Models/Packages/A3WebUI.cs index 2397a5ef..1c976fa4 100644 --- a/StabilityMatrix.Core/Models/Packages/A3WebUI.cs +++ b/StabilityMatrix.Core/Models/Packages/A3WebUI.cs @@ -207,4 +207,10 @@ public class A3WebUI : BaseGitPackage await StabilityMatrix.Core.Helper.SharedFolders.UpdateLinksForPackage(this, SettingsManager.ModelsDirectory, installDirectory).ConfigureAwait(false); } + + public override Task RemoveModelFolderLinks(DirectoryPath installDirectory) + { + StabilityMatrix.Core.Helper.SharedFolders.RemoveLinksForPackage(this, installDirectory); + return Task.CompletedTask; + } } diff --git a/StabilityMatrix.Core/Models/Packages/BasePackage.cs b/StabilityMatrix.Core/Models/Packages/BasePackage.cs index 837f2ae1..285e5b10 100644 --- a/StabilityMatrix.Core/Models/Packages/BasePackage.cs +++ b/StabilityMatrix.Core/Models/Packages/BasePackage.cs @@ -39,6 +39,7 @@ public abstract class BasePackage public abstract Task RunPackage(string installedPackagePath, string command, string arguments); public abstract Task SetupModelFolders(DirectoryPath installDirectory); public abstract Task UpdateModelFolders(DirectoryPath installDirectory); + public abstract Task RemoveModelFolderLinks(DirectoryPath installDirectory); /// /// Shuts down the subprocess, canceling any pending streams. diff --git a/StabilityMatrix.Core/Models/Packages/ComfyUI.cs b/StabilityMatrix.Core/Models/Packages/ComfyUI.cs index 141511ac..c7863ca5 100644 --- a/StabilityMatrix.Core/Models/Packages/ComfyUI.cs +++ b/StabilityMatrix.Core/Models/Packages/ComfyUI.cs @@ -225,6 +225,9 @@ public class ComfyUI : BaseGitPackage public override Task UpdateModelFolders(DirectoryPath installDirectory) => SetupModelFolders(installDirectory); + public override Task RemoveModelFolderLinks(DirectoryPath installDirectory) => + Task.CompletedTask; + public class ComfyModelPathsYaml { public class SmData diff --git a/StabilityMatrix.Core/Models/Packages/DankDiffusion.cs b/StabilityMatrix.Core/Models/Packages/DankDiffusion.cs index 4777f09f..5f55b4d7 100644 --- a/StabilityMatrix.Core/Models/Packages/DankDiffusion.cs +++ b/StabilityMatrix.Core/Models/Packages/DankDiffusion.cs @@ -44,6 +44,11 @@ public class DankDiffusion : BaseGitPackage throw new NotImplementedException(); } + public override Task RemoveModelFolderLinks(DirectoryPath installDirectory) + { + throw new NotImplementedException(); + } + public override List LaunchOptions { get; } public override Task GetLatestVersion() { diff --git a/StabilityMatrix.Core/Models/Packages/InvokeAI.cs b/StabilityMatrix.Core/Models/Packages/InvokeAI.cs index 178246e9..dcad4af6 100644 --- a/StabilityMatrix.Core/Models/Packages/InvokeAI.cs +++ b/StabilityMatrix.Core/Models/Packages/InvokeAI.cs @@ -1,9 +1,11 @@ using System.Data; +using System.Text.Json; using System.Text.RegularExpressions; using NLog; using StabilityMatrix.Core.Extensions; using StabilityMatrix.Core.Helper; using StabilityMatrix.Core.Helper.Cache; +using StabilityMatrix.Core.Models.Api; using StabilityMatrix.Core.Models.FileInterfaces; using StabilityMatrix.Core.Models.Progress; using StabilityMatrix.Core.Processes; @@ -28,9 +30,6 @@ public class InvokeAI : BaseGitPackage public override string Blurb => "Professional Creative Tools for Stable Diffusion"; public override string LaunchCommand => "invokeai-web"; - public override string Disclaimer => - "Note: InvokeAI support is currently experimental, and checkpoints in the shared models folder will not be available with InvokeAI."; - public override IReadOnlyList ExtraLaunchCommands => new[] { "invokeai-configure", @@ -293,10 +292,12 @@ public class InvokeAI : BaseGitPackage { await SetupVenv(installedPackagePath).ConfigureAwait(false); - if (command.Equals("invokeai-configure")) + arguments = command switch { - arguments = "--yes --skip-sd-weights"; - } + "invokeai-configure" => "--yes --skip-sd-weights", + "invokeai-model-install" => "--yes", + _ => arguments + }; VenvRunner.EnvironmentVariables = GetEnvVars(installedPackagePath); @@ -355,11 +356,92 @@ public class InvokeAI : BaseGitPackage public override Task SetupModelFolders(DirectoryPath installDirectory) { + var modelsFolder = SettingsManager.ModelsDirectory; + if (!Directory.Exists(modelsFolder)) + return Task.CompletedTask; + + var connectedModelJsons = + Directory.GetFiles(modelsFolder, "*.cm-info.json", SearchOption.AllDirectories); + + var connectedModelDict = new Dictionary(); + + foreach (var jsonFile in connectedModelJsons) + { + var json = File.ReadAllText(jsonFile); + var connectedModelInfo = JsonSerializer.Deserialize(json); + var extension = connectedModelInfo?.FileMetadata.Format switch + { + CivitModelFormat.SafeTensor => ".safetensors", + CivitModelFormat.PickleTensor => ".pt", + _ => string.Empty + }; + + if (string.IsNullOrWhiteSpace(extension) || connectedModelInfo == null) + continue; + + var modelFilePath = jsonFile.Replace(".cm-info.json", extension); + if (File.Exists(modelFilePath)) + { + connectedModelDict[modelFilePath] = connectedModelInfo; + } + } + + foreach (var modelFilePath in connectedModelDict.Keys) + { + var model = connectedModelDict[modelFilePath]; + var modelType = model.ModelType switch + { + CivitModelType.Checkpoint => "main", + CivitModelType.LORA => "lora", + CivitModelType.TextualInversion => "embedding", + CivitModelType.Controlnet => "controlnet", + _ => string.Empty + }; + + if (string.IsNullOrWhiteSpace(modelType)) + continue; + + var sourcePath = modelFilePath; + var destinationPath = Path.Combine(installDirectory, RelativeRootPath, "autoimport", + modelType, Path.GetFileName(modelFilePath)); + + try + { + File.CreateSymbolicLink(destinationPath, sourcePath); + } + catch (IOException e) + { + // File already exists + Logger.Warn(e, + $"Could not create symlink for {sourcePath} to {destinationPath} - file already exists"); + } + } + return Task.CompletedTask; } - public override Task UpdateModelFolders(DirectoryPath installDirectory) + public override Task UpdateModelFolders(DirectoryPath installDirectory) => + SetupModelFolders(installDirectory); + + public override Task RemoveModelFolderLinks(DirectoryPath installDirectory) { + var autoImportDir = Path.Combine(installDirectory, RelativeRootPath, "autoimport"); + var allSymlinks = + Directory.GetFiles(autoImportDir, "*.*", SearchOption.AllDirectories) + .Select(path => new FileInfo(path)).Where(file => file.LinkTarget != null); + + foreach (var link in allSymlinks) + { + try + { + link.Delete(); + } + catch (IOException e) + { + Logger.Warn(e, $"Could not delete symlink {link.FullName}"); + } + } + return Task.CompletedTask; } diff --git a/StabilityMatrix.Core/Models/Packages/VladAutomatic.cs b/StabilityMatrix.Core/Models/Packages/VladAutomatic.cs index cf0292f3..9d63c549 100644 --- a/StabilityMatrix.Core/Models/Packages/VladAutomatic.cs +++ b/StabilityMatrix.Core/Models/Packages/VladAutomatic.cs @@ -294,10 +294,11 @@ public class VladAutomatic : BaseGitPackage return Task.CompletedTask; } - public override Task UpdateModelFolders(DirectoryPath installDirectory) - { - return SetupModelFolders(installDirectory); - } + public override Task UpdateModelFolders(DirectoryPath installDirectory) => + SetupModelFolders(installDirectory); + + public override Task RemoveModelFolderLinks(DirectoryPath installDirectory) => + Task.CompletedTask; public override async Task Update(InstalledPackage installedPackage, IProgress? progress = null, bool includePrerelease = false) diff --git a/StabilityMatrix.Core/Models/Packages/VoltaML.cs b/StabilityMatrix.Core/Models/Packages/VoltaML.cs index f10c16dc..c566bd01 100644 --- a/StabilityMatrix.Core/Models/Packages/VoltaML.cs +++ b/StabilityMatrix.Core/Models/Packages/VoltaML.cs @@ -191,4 +191,10 @@ public class VoltaML : BaseGitPackage await StabilityMatrix.Core.Helper.SharedFolders.UpdateLinksForPackage(this, SettingsManager.ModelsDirectory, installDirectory).ConfigureAwait(false); } + + public override Task RemoveModelFolderLinks(DirectoryPath installDirectory) + { + StabilityMatrix.Core.Helper.SharedFolders.RemoveLinksForPackage(this, installDirectory); + return Task.CompletedTask; + } }