Shared folder strategies

StabilityMatrix.Tests/Helper/PackageFactoryTests.cs

@@ -16,7 +16,7 @@ public void Setup()
         fakeBasePackages = new List<BasePackage>
         {
             // TODO: inject mocks
-            new DankDiffusion(null, null, null, null)
+            new DankDiffusion(null, null, null, null, null)
         };
         packageFactory = new PackageFactory(fakeBasePackages);
     }

StabilityMatrix/App.xaml.cs

@@ -1,4 +1,4 @@
-using System;
+using System;
 using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
 using System.IO;
@@ -231,6 +231,8 @@ private void App_OnStartup(object sender, StartupEventArgs e)
             serviceCollection.AddSingleton<BasePackage, A3WebUI>();
             serviceCollection.AddSingleton<BasePackage, VladAutomatic>();
             serviceCollection.AddSingleton<BasePackage, ComfyUI>();
+            serviceCollection.AddSingleton<VladAutomaticSharedFolderStrategy>();
+            serviceCollection.AddTransient<LinkedFolderSharedFolderStrategy>();
             serviceCollection.AddSingleton<Wpf.Ui.Contracts.ISnackbarService, Wpf.Ui.Services.SnackbarService>();
             serviceCollection.AddSingleton<IPrerequisiteHelper, PrerequisiteHelper>();
             serviceCollection.AddSingleton<ISnackbarService, SnackbarService>();

StabilityMatrix/Models/Packages/A3WebUI.cs

@@ -24,10 +24,14 @@ public class A3WebUI : BaseGitPackage
     public string RelativeArgsDefinitionScriptPath => "modules.cmd_args";
 
 
-    public A3WebUI(IGithubApiCache githubApi, ISettingsManager settingsManager, IDownloadService downloadService,
-        IPrerequisiteHelper prerequisiteHelper) :
+    public A3WebUI(IGithubApiCache githubApi, 
+        ISettingsManager settingsManager, 
+        IDownloadService downloadService,
+        IPrerequisiteHelper prerequisiteHelper, 
+        LinkedFolderSharedFolderStrategy sharedFolderStrategy) :
         base(githubApi, settingsManager, downloadService, prerequisiteHelper)
     {
+        SharedFolderStrategy = sharedFolderStrategy;
     }
 
     // From https://github.com/AUTOMATIC1111/stable-diffusion-webui/tree/master/models
@@ -48,6 +52,8 @@ public A3WebUI(IGithubApiCache githubApi, ISettingsManager settingsManager, IDow
         [SharedFolderType.ControlNet] = "models/ControlNet"
     };
 
+    public override ISharedFolderStrategy SharedFolderStrategy { get; protected set; }
+
     [SuppressMessage("ReSharper", "ArrangeObjectCreationWhenTypeNotEvident")]
     public override List<LaunchOptionDefinition> LaunchOptions => new()
     {

StabilityMatrix/Models/Packages/BaseGitPackage.cs

@@ -28,7 +28,11 @@ public abstract class BaseGitPackage : BasePackage
     protected readonly IDownloadService DownloadService;
     protected readonly IPrerequisiteHelper PrerequisiteHelper;
     protected PyVenvRunner? VenvRunner;
-
+
+    public BaseGitPackage()
+    {
+
+    }
     /// <summary>
     /// URL of the hosted web page on launch
     /// </summary>

StabilityMatrix/Models/Packages/BasePackage.cs

@@ -1,4 +1,4 @@
-using System;
+using System;
 using System.Collections.Generic;
 using System.Linq;
 using System.Threading.Tasks;
@@ -31,13 +31,15 @@ public abstract class BasePackage
 
     public abstract List<LaunchOptionDefinition> LaunchOptions { get; }
     public virtual string? ExtraLaunchArguments { get; set; } = null;
-    
+
     /// <summary>
     /// The shared folders that this package supports.
     /// Mapping of <see cref="SharedFolderType"/> to the relative path from the package root.
     /// </summary>
     public virtual Dictionary<SharedFolderType, string>? SharedFolders { get; }
-
+
+    public abstract ISharedFolderStrategy SharedFolderStrategy { get; protected set; }
+
     public abstract Task<string> GetLatestVersion();
     public abstract Task<IEnumerable<PackageVersion>> GetAllVersions(bool isReleaseMode = true);
     public abstract Task<IReadOnlyList<GitHubCommit>?> GetAllCommits(string branch, int page = 1, int perPage = 10);

StabilityMatrix/Models/Packages/ComfyUI.cs

@@ -26,9 +26,10 @@ public class ComfyUI : BaseGitPackage
     public override bool ShouldIgnoreReleases => true;
 
     public ComfyUI(IGithubApiCache githubApi, ISettingsManager settingsManager, IDownloadService downloadService,
-        IPrerequisiteHelper prerequisiteHelper) :
+        IPrerequisiteHelper prerequisiteHelper, LinkedFolderSharedFolderStrategy sharedFolderStrategy) :
         base(githubApi, settingsManager, downloadService, prerequisiteHelper)
     {
+        SharedFolderStrategy = sharedFolderStrategy;
     }
 
     // https://github.com/comfyanonymous/ComfyUI/blob/master/folder_paths.py#L11
@@ -46,7 +47,9 @@ public ComfyUI(IGithubApiCache githubApi, ISettingsManager settingsManager, IDow
         [SharedFolderType.ESRGAN] = "models/upscale_models",
         [SharedFolderType.Hypernetwork] = "models/hypernetworks",
     };
-
+
+    public override ISharedFolderStrategy SharedFolderStrategy { get; protected set; }
+
     public override List<LaunchOptionDefinition> LaunchOptions => new()
     {
         new()

StabilityMatrix/Models/Packages/DankDiffusion.cs

@@ -10,9 +10,10 @@ namespace StabilityMatrix.Models.Packages;
 public class DankDiffusion : BaseGitPackage
 {
     public DankDiffusion(IGithubApiCache githubApi, ISettingsManager settingsManager, IDownloadService downloadService,
-        IPrerequisiteHelper prerequisiteHelper) :
+        IPrerequisiteHelper prerequisiteHelper, LinkedFolderSharedFolderStrategy sharedFolderStrategy) :
         base(githubApi, settingsManager, downloadService, prerequisiteHelper)
     {
+        SharedFolderStrategy = sharedFolderStrategy;
     }
 
     public override string Name => "dank-diffusion";
@@ -27,6 +28,8 @@ public override Task RunPackage(string installedPackagePath, string arguments)
     }
 
     public override List<LaunchOptionDefinition> LaunchOptions { get; }
+    public override ISharedFolderStrategy SharedFolderStrategy { get; protected set; }
+
     public override Task<string> GetLatestVersion()
     {
         throw new System.NotImplementedException();

StabilityMatrix/Models/Packages/ISharedFolderStrategy.cs

@@ -0,0 +1,8 @@
+using System.Threading.Tasks;
+
+namespace StabilityMatrix.Models.Packages;
+
+public interface ISharedFolderStrategy
+{
+    Task ExecuteAsync(BasePackage package);
+}

StabilityMatrix/Models/Packages/LinkedFolderSharedFolderStrategy.cs

@@ -0,0 +1,24 @@
+using System;
+using System.Threading.Tasks;
+using Microsoft.Extensions.DependencyInjection;
+
+namespace StabilityMatrix.Models.Packages;
+
+public class LinkedFolderSharedFolderStrategy : ISharedFolderStrategy
+{
+    private readonly IServiceProvider serviceProvider;
+
+    public LinkedFolderSharedFolderStrategy(IServiceProvider serviceProvider)
+    {
+        this.serviceProvider = serviceProvider;
+    }
+
+    public Task ExecuteAsync(BasePackage package)
+    {
+        // TODO: Move SharedFolders logic here
+        // NOTE: We're using this awkward solution because a circular dependency is generated in the graph otherwise
+        var sharedFolders = serviceProvider.GetRequiredService<ISharedFolders>();
+        sharedFolders.SetupLinksForPackage(package, package.InstallLocation);
+        return Task.CompletedTask;
+    }
+}

StabilityMatrix/Models/Packages/VladAutomatic.cs

@@ -27,9 +27,10 @@ public class VladAutomatic : BaseGitPackage
     public override bool ShouldIgnoreReleases => true;
 
     public VladAutomatic(IGithubApiCache githubApi, ISettingsManager settingsManager, IDownloadService downloadService,
-        IPrerequisiteHelper prerequisiteHelper) :
+        IPrerequisiteHelper prerequisiteHelper, VladAutomaticSharedFolderStrategy sharedFolderStrategy) :
         base(githubApi, settingsManager, downloadService, prerequisiteHelper)
     {
+        SharedFolderStrategy = sharedFolderStrategy;
     }
 
     // https://github.com/vladmandic/automatic/blob/master/modules/shared.py#L324
@@ -53,6 +54,8 @@ public VladAutomatic(IGithubApiCache githubApi, ISettingsManager settingsManager
         [SharedFolderType.LyCORIS] = "models/LyCORIS",
     };
 
+    public override ISharedFolderStrategy SharedFolderStrategy { get; protected set; }
+
     [SuppressMessage("ReSharper", "ArrangeObjectCreationWhenTypeNotEvident")]
     public override List<LaunchOptionDefinition> LaunchOptions => new()
     {

StabilityMatrix/Models/Packages/VladAutomaticSharedFolderStrategy.cs

@@ -0,0 +1,67 @@
+using System.Collections.Generic;
+using System.IO;
+using System.Linq;
+using System.Threading.Tasks;
+using Newtonsoft.Json;
+using Newtonsoft.Json.Linq;
+using StabilityMatrix.Extensions;
+using StabilityMatrix.Helper;
+using StabilityMatrix.Models.FileInterfaces;
+
+namespace StabilityMatrix.Models.Packages;
+
+public class VladAutomaticSharedFolderStrategy : ISharedFolderStrategy
+{
+    private readonly ISettingsManager settingsManager;
+
+    public VladAutomaticSharedFolderStrategy(ISettingsManager settingsManager)
+    {
+        this.settingsManager = settingsManager;
+    }
+
+    public async Task ExecuteAsync(BasePackage package)
+    {
+        var installedPackage = settingsManager
+            .Settings
+            .InstalledPackages
+            .Single(p => p.PackageName == package.Name);
+        var configFilePath = Path.Combine(settingsManager.LibraryDir, installedPackage.LibraryPath!, "config.json");
+
+        // Load the configuration file
+        var json = await File.ReadAllTextAsync(configFilePath);
+        var job = JsonConvert.DeserializeObject<JObject>(json)!;
+
+        // Update the configuration values
+        var modelsDirectory = new DirectoryPath(settingsManager.ModelsDirectory);
+        foreach (var (sharedFolderType, configKey) in map)
+        {
+            var value = Path.Combine(modelsDirectory.FullPath, sharedFolderType.GetStringValue());
+            job[configKey] = value;
+        }
+
+        // Write the configuration file
+        await File.WriteAllTextAsync(configFilePath, JsonConvert.SerializeObject(job, Formatting.Indented));
+    }
+
+    private Dictionary<SharedFolderType, string> map = new()
+    {
+        { SharedFolderType.StableDiffusion, "ckpt_dir" },
+        { SharedFolderType.Diffusers, "diffusers_dir" },
+        { SharedFolderType.VAE, "vae_dir" },
+        { SharedFolderType.Lora, "lora_dir" },
+        { SharedFolderType.LyCORIS, "lyco_dir" },
+        // { SharedFolderType.Styles, "styles_dir"},
+        { SharedFolderType.TextualInversion, "embeddings_dir" },
+        { SharedFolderType.Hypernetwork, "hypernetwork_dir" },
+        { SharedFolderType.Codeformer, "codeformer_models_path" },
+        { SharedFolderType.GFPGAN, "gfpgan_models_path" },
+        { SharedFolderType.ESRGAN, "esrgan_models_path" },
+        { SharedFolderType.BSRGAN , "bsrgan_models_path"},
+        { SharedFolderType.RealESRGAN, "realesrgan_models_path" },
+        { SharedFolderType.ScuNET, "scunet_models_path" },
+        { SharedFolderType.SwinIR, "swinir_models_path" },
+        { SharedFolderType.LDSR, "ldsr_models_path" },
+        { SharedFolderType.CLIP, "clip_models_path" }
+    };
+} 
+

StabilityMatrix/StabilityMatrix.csproj

@@ -30,6 +30,7 @@
       <PackageReference Include="Microsoft.Toolkit.Uwp.Notifications" Version="7.1.3" />
       <PackageReference Include="Microsoft.Web.WebView2" Version="1.0.1823.32" />
       <PackageReference Include="Microsoft.Xaml.Behaviors.Wpf" Version="1.1.39" />
+      <PackageReference Include="Newtonsoft.Json" Version="13.0.1" />
       <PackageReference Include="NLog" Version="5.2.0" />
       <PackageReference Include="NLog.Extensions.Logging" Version="5.3.0" />
       <PackageReference Include="NSec.Cryptography" Version="22.4.0" />

StabilityMatrix/ViewModels/InstallerViewModel.cs

@@ -318,8 +318,9 @@ private async Task ActuallyInstall()
         await InstallPackage();
 
         ProgressText = "Setting up shared folder links...";
-        sharedFolders.SetupLinksForPackage(SelectedPackage, SelectedPackage.InstallLocation);
-
+
+        await SelectedPackage.SharedFolderStrategy.ExecuteAsync(SelectedPackage);
+
         ProgressText = "Done";
         IsIndeterminate = false;
         ProgressValue = 100;

StabilityMatrix/ViewModels/LaunchViewModel.cs

@@ -157,7 +157,10 @@ private void ToastNotificationManagerCompatOnOnActivated(
         basePackage.StartupComplete += RunningPackageOnStartupComplete;
 
         // Update shared folder links (in case library paths changed)
-        sharedFolders.UpdateLinksForPackage(basePackage, packagePath);
+        if (basePackage.SharedFolderStrategy is LinkedFolderSharedFolderStrategy)
+            sharedFolders.UpdateLinksForPackage(basePackage, packagePath);
+        else
+            await basePackage.SharedFolderStrategy.ExecuteAsync(basePackage);
 
         // Load user launch args from settings and convert to string
         var userArgs = settingsManager.GetLaunchArgs(activeInstall.Id);