Browse Source

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

pull/109/head
JT 1 year ago
parent
commit
faa7b7c75c
  1. 10
      StabilityMatrix.Core/Helper/SharedFolders.cs
  2. 6
      StabilityMatrix.Core/Models/Packages/A3WebUI.cs
  3. 1
      StabilityMatrix.Core/Models/Packages/BasePackage.cs
  4. 3
      StabilityMatrix.Core/Models/Packages/ComfyUI.cs
  5. 5
      StabilityMatrix.Core/Models/Packages/DankDiffusion.cs
  6. 96
      StabilityMatrix.Core/Models/Packages/InvokeAI.cs
  7. 9
      StabilityMatrix.Core/Models/Packages/VladAutomatic.cs
  8. 6
      StabilityMatrix.Core/Models/Packages/VoltaML.cs

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

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

1
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.

3
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

5
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()
{

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

9
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<string> Update(InstalledPackage installedPackage,
IProgress<ProgressReport>? progress = null, bool includePrerelease = false)

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

Loading…
Cancel
Save