Added shared model folder handling for InvokeAI & made the packages handle their own RemoveModelFolderLinks

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);
             }
         }

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;
+    }
 }

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);
     
     /// <summary>
     /// Shuts down the subprocess, canceling any pending streams.

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

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<LaunchOptionDefinition> LaunchOptions { get; }
     public override Task<string> GetLatestVersion()
     {

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<string> 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<string, ConnectedModelInfo>();
+
+        foreach (var jsonFile in connectedModelJsons)
+        {
+            var json = File.ReadAllText(jsonFile);
+            var connectedModelInfo = JsonSerializer.Deserialize<ConnectedModelInfo>(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;
     }
 

StabilityMatrix.Core/Models/Packages/VladAutomatic.cs

@@ -294,10 +294,11 @@ public class VladAutomatic : BaseGitPackage
         return Task.CompletedTask;
     }
 
-    public override Task UpdateModelFolders(DirectoryPath installDirectory)
+    public override Task UpdateModelFolders(DirectoryPath installDirectory) =>
-    {
+        SetupModelFolders(installDirectory);
-        return SetupModelFolders(installDirectory);
+
-    }
+    public override Task RemoveModelFolderLinks(DirectoryPath installDirectory) =>
+        Task.CompletedTask;
 
     public override async Task<string> Update(InstalledPackage installedPackage,
         IProgress<ProgressReport>? progress = null, bool includePrerelease = false)

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;
+    }
 }