Merge branch 'main' of https://github.com/ionite34/StabilityMatrix

LykosAI · Aug 25, 2023 · 116dcc9 · 116dcc9
2 parents 1c77fa1 + 456a226
commit 116dcc9
Show file tree

Hide file tree

Showing 10 changed files with 218 additions and 88 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -5,6 +5,11 @@ All notable changes to Stability Matrix will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project adheres to [Semantic Versioning 2.0](https://semver.org/spec/v2.0.0.html).
 
+## v2.3.1
+### Fixed
+- Fixed Auto update not appearing in some regions due to date formatting issues
+- Local package import now migrates venvs and existing models
+
 ## v2.3.0
 ### Added
 - New installable Package - [Fooocus](https://github.com/lllyasviel/Fooocus)

diff --git a/StabilityMatrix.Avalonia/Languages/Resources.Designer.cs b/StabilityMatrix.Avalonia/Languages/Resources.Designer.cs
diff --git a/StabilityMatrix.Avalonia/Languages/Resources.resx b/StabilityMatrix.Avalonia/Languages/Resources.resx
@@ -66,4 +66,7 @@
     <data name="Label_Branches" xml:space="preserve">
         <value>Branches</value>
     </data>
+    <data name="Label_DragAndDropCheckpointsHereToImport" xml:space="preserve">
+        <value>Drag &amp; Drop checkpoints here to import</value>
+    </data>
 </root>
diff --git a/StabilityMatrix.Avalonia/ViewModels/CheckpointManager/CheckpointFolder.cs b/StabilityMatrix.Avalonia/ViewModels/CheckpointManager/CheckpointFolder.cs
@@ -1,7 +1,6 @@
 using System;
 using System.Collections.Generic;
 using System.Collections.Specialized;
-using System.Diagnostics;
 using System.IO;
 using System.Linq;
 using System.Threading.Tasks;
@@ -75,7 +74,7 @@ public partial class CheckpointFolder : ViewModelBase
             ? $"{Title} ({CheckpointFiles.Count + SubFolders.Sum(folder => folder.CheckpointFiles.Count)})"
             : Title;
 
-    public Base.ProgressViewModel Progress { get; } = new();
+    public ProgressViewModel Progress { get; } = new();
 
     public CheckpointFolder? ParentFolder { get; init; }
     public AdvancedObservableList<CheckpointFolder> SubFolders { get; init; } = new();

diff --git a/StabilityMatrix.Avalonia/ViewModels/Dialogs/PackageImportViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/Dialogs/PackageImportViewModel.cs
@@ -183,7 +183,7 @@ private void UpdateSelectedVersionToLatestMain()
         }
     }
 
-    public void AddPackageWithCurrentInputs()
+    public async Task AddPackageWithCurrentInputs()
     {
         if (SelectedBasePackage is null || PackagePath is null)
             return;
@@ -216,6 +216,15 @@ public void AddPackageWithCurrentInputs()
             LastUpdateCheck = DateTimeOffset.Now,
         };
 
+        // Recreate venv if it's a BaseGitPackage
+        if (SelectedBasePackage is BaseGitPackage gitPackage)
+        {
+            await gitPackage.SetupVenv(PackagePath, forceRecreate: true);
+        }
+
+        // Reconfigure shared links
+        await SelectedBasePackage.UpdateModelFolders(PackagePath);
+
         settingsManager.Transaction(s => s.InstalledPackages.Add(package));
     }
 }
diff --git a/StabilityMatrix.Avalonia/ViewModels/PackageManager/PackageCardViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/PackageManager/PackageCardViewModel.cs
@@ -1,4 +1,5 @@
 using System;
+using System.Collections.Generic;
 using System.Threading.Tasks;
 using Avalonia.Controls;
 using Avalonia.Controls.Notifications;
@@ -11,6 +12,7 @@
 using StabilityMatrix.Avalonia.Services;
 using StabilityMatrix.Avalonia.ViewModels.Base;
 using StabilityMatrix.Avalonia.ViewModels.Dialogs;
+using StabilityMatrix.Avalonia.Views.Dialogs;
 using StabilityMatrix.Core.Extensions;
 using StabilityMatrix.Core.Helper;
 using StabilityMatrix.Core.Helper.Factory;
@@ -219,23 +221,55 @@ public async Task Import()
     {
         if (!IsUnknownPackage || Design.IsDesignMode) return;
 
-        PackageImportViewModel viewModel = null!;
-        var dialog = vmFactory.GetDialog<PackageImportViewModel>(vm =>
+        var viewModel = vmFactory.Get<PackageImportViewModel>(vm =>
         {
-            vm.PackagePath = new DirectoryPath(Package?.FullPath ?? throw new InvalidOperationException());
-            viewModel = vm;
+            vm.PackagePath =
+                new DirectoryPath(Package?.FullPath ?? throw new InvalidOperationException());
         });
 
-        dialog.PrimaryButtonText = Resources.Action_Import;
-        dialog.CloseButtonText = Resources.Action_Cancel;
-        dialog.DefaultButton = ContentDialogButton.Primary;
+        var dialog = new TaskDialog
+        {
+            Content = new PackageImportDialog
+            {
+                DataContext = viewModel
+            },
+            ShowProgressBar = false,
+            Buttons = new List<TaskDialogButton>
+            {
+                new(Resources.Action_Import, TaskDialogStandardResult.Yes)
+                {
+                    IsDefault = true
+                },
+                new(Resources.Action_Cancel, TaskDialogStandardResult.Cancel)
+            }
+        };
 
-        var result = await dialog.ShowAsync();
-        if (result == ContentDialogResult.Primary)
+        dialog.Closing += async (sender, e) =>
         {
-            viewModel.AddPackageWithCurrentInputs();
-            EventManager.Instance.OnInstalledPackagesChanged();
-        }
+            // We only want to use the deferral on the 'Yes' Button
+            if ((TaskDialogStandardResult)e.Result == TaskDialogStandardResult.Yes)
+            {
+                var deferral = e.GetDeferral();
+
+                sender.ShowProgressBar = true;
+                sender.SetProgressBarState(0, TaskDialogProgressState.Indeterminate);
+
+                await using (new MinimumDelay(200, 300))
+                {
+                    var result = await notificationService.TryAsync(viewModel.AddPackageWithCurrentInputs());
+                    if (result.IsSuccessful)
+                    {
+                        EventManager.Instance.OnInstalledPackagesChanged();
+                    }
+                }
+
+                deferral.Complete();
+            }
+        };
+
+        dialog.XamlRoot = App.VisualRoot;
+
+        await dialog.ShowAsync(true);
     }
 
     public async Task OpenFolder()

diff --git a/StabilityMatrix.Avalonia/Views/CheckpointsPage.axaml b/StabilityMatrix.Avalonia/Views/CheckpointsPage.axaml
@@ -8,6 +8,7 @@
              xmlns:controls="clr-namespace:StabilityMatrix.Avalonia.Controls"
              xmlns:system="clr-namespace:System;assembly=System.Runtime"
              xmlns:checkpointManager="clr-namespace:StabilityMatrix.Avalonia.ViewModels.CheckpointManager"
+             xmlns:lang="clr-namespace:StabilityMatrix.Avalonia.Languages"
              d:DataContext="{x:Static mocks:DesignData.CheckpointsPageViewModel}"
              x:CompileBindings="True"
              x:DataType="vm:CheckpointsPageViewModel"
@@ -255,7 +256,6 @@
                           Background="Transparent"
                           DragDrop.AllowDrop="True">
                         <ItemsRepeater
-                            Classes="ItemsDragAndDrop"
                             ItemTemplate="{StaticResource CheckpointFileDataTemplate}"
                             ItemsSource="{Binding DisplayedCheckpointFiles}">
                             <ItemsRepeater.Layout>
@@ -266,8 +266,8 @@
                         <TextBlock VerticalAlignment="Center"
                                    TextAlignment="Center"
                                    HorizontalAlignment="Center"
-                                   Margin="8"
-                                   Text="Drag &amp; drop checkpoints here to import"
+                                   Margin="8,8,8,16"
+                                   Text="{x:Static lang:Resources.Label_DragAndDropCheckpointsHereToImport}"
                                    IsVisible="{Binding !CheckpointFiles.Count}"/>
                         <!-- Blurred background for drag and drop -->
                         <Border
@@ -284,7 +284,7 @@
                             Effect="{StaticResource TextDropShadowEffect}"
                             FontSize="24"
                             HorizontalAlignment="Center"
-                            Text="Drop a file here to import"
+                            Text="Drop file here to import"
                             VerticalAlignment="Center"
                             IsVisible="{Binding IsCurrentDragTarget}" />
                         <StackPanel

diff --git a/StabilityMatrix.Core/Database/LiteDbContext.cs b/StabilityMatrix.Core/Database/LiteDbContext.cs
@@ -1,4 +1,6 @@
-using LiteDB.Async;
+using LiteDB;
+using LiteDB.Async;
+using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Options;
 using StabilityMatrix.Core.Extensions;
 using StabilityMatrix.Core.Models.Api;
@@ -10,27 +12,12 @@ namespace StabilityMatrix.Core.Database;
 
 public class LiteDbContext : ILiteDbContext
 {
+    private readonly ILogger<LiteDbContext> logger;
     private readonly ISettingsManager settingsManager;
     private readonly DebugOptions debugOptions;
 
     private LiteDatabaseAsync? database;
-    public LiteDatabaseAsync Database
-    {
-        get
-        {
-            if (database != null) return database;
-
-            var connectionString = debugOptions.TempDatabase ? ":temp:"
-                : $"Filename={Path.Combine(settingsManager.LibraryDir, "StabilityMatrix.db")};Mode=Shared";
-            database = new LiteDatabaseAsync(connectionString);
-
-            // Register reference fields
-            LiteDBExtensions.Register<CivitModel, CivitModelVersion>(m => m.ModelVersions, "CivitModelVersions");
-            LiteDBExtensions.Register<CivitModelQueryCacheEntry, CivitModel>(e => e.Items, "CivitModels");
-
-            return database;
-        }
-    }
+    public LiteDatabaseAsync Database => database ??= CreateDatabase();
 
     // Notification events
     public event EventHandler? CivitModelsChanged;
@@ -41,11 +28,49 @@ public LiteDatabaseAsync Database
     public ILiteCollectionAsync<CivitModelQueryCacheEntry> CivitModelQueryCache => Database.GetCollection<CivitModelQueryCacheEntry>("CivitModelQueryCache");
     public ILiteCollectionAsync<GithubCacheEntry> GithubCache => Database.GetCollection<GithubCacheEntry>("GithubCache");
 
-    public LiteDbContext(ISettingsManager settingsManager, IOptions<DebugOptions> debugOptions)
+    public LiteDbContext(
+        ILogger<LiteDbContext> logger,
+        ISettingsManager settingsManager, 
+        IOptions<DebugOptions> debugOptions)
     {
+        this.logger = logger;
         this.settingsManager = settingsManager;
         this.debugOptions = debugOptions.Value;
     }
+
+    private LiteDatabaseAsync CreateDatabase()
+    {
+        LiteDatabaseAsync? db = null;
+
+        if (debugOptions.TempDatabase)
+        {
+            db = new LiteDatabaseAsync(":temp:");
+        }
+
+        // Attempt to create connection, might be in use
+        try
+        {
+            var dbPath = Path.Combine(settingsManager.LibraryDir, "StabilityMatrix.db");
+            db = new LiteDatabaseAsync(new ConnectionString()
+            {
+                Filename = dbPath,
+                Connection = ConnectionType.Shared,
+            });
+        }
+        catch (IOException e)
+        {
+            logger.LogWarning("Database in use or not accessible ({Message}), using temporary database", e.Message);
+        }
+
+        // Fallback to temporary database
+        db ??= new LiteDatabaseAsync(":temp:");
+
+        // Register reference fields
+        LiteDBExtensions.Register<CivitModel, CivitModelVersion>(m => m.ModelVersions, "CivitModelVersions");
+        LiteDBExtensions.Register<CivitModelQueryCacheEntry, CivitModel>(e => e.Items, "CivitModels");
+
+        return db;
+    }
 
     public async Task<(CivitModel?, CivitModelVersion?)> FindCivitModelFromFileHashAsync(string hashBlake3)
     {
@@ -54,22 +79,26 @@ public LiteDbContext(ISettingsManager settingsManager, IOptions<DebugOptions> de
                 .Select(f => f.Hashes)
                 .Select(hashes => hashes.BLAKE3)
                 .Any(hash => hash == hashBlake3))
-            .FirstOrDefaultAsync();
+            .FirstOrDefaultAsync()
+            .ConfigureAwait(false);
+
         if (version is null) return (null, null);
+
         var model = await CivitModels.Query()
             .Include(m => m.ModelVersions)
             .Where(m => m.ModelVersions!
                 .Select(v => v.Id)
                 .Any(id => id == version.Id))
-            .FirstOrDefaultAsync();
+            .FirstOrDefaultAsync().ConfigureAwait(false);
+
         return (model, version);
     }
 
     public async Task<bool> UpsertCivitModelAsync(CivitModel civitModel)
     {
         // Insert model versions first then model
-        var versionsUpdated = await CivitModelVersions.UpsertAsync(civitModel.ModelVersions);
-        var updated = await CivitModels.UpsertAsync(civitModel);
+        var versionsUpdated = await CivitModelVersions.UpsertAsync(civitModel.ModelVersions).ConfigureAwait(false);
+        var updated = await CivitModels.UpsertAsync(civitModel).ConfigureAwait(false);
         // Notify listeners on any change
         var anyUpdated = versionsUpdated > 0 || updated;
         if (anyUpdated)
@@ -84,8 +113,8 @@ public async Task<bool> UpsertCivitModelAsync(IEnumerable<CivitModel> civitModel
         var civitModelsArray = civitModels.ToArray();
         // Get all model versions then insert models
         var versions = civitModelsArray.SelectMany(model => model.ModelVersions ?? new());
-        var versionsUpdated = await CivitModelVersions.UpsertAsync(versions);
-        var updated = await CivitModels.UpsertAsync(civitModelsArray);
+        var versionsUpdated = await CivitModelVersions.UpsertAsync(versions).ConfigureAwait(false);
+        var updated = await CivitModels.UpsertAsync(civitModelsArray).ConfigureAwait(false);
         // Notify listeners on any change
         var anyUpdated = versionsUpdated > 0 || updated > 0;
         if (updated > 0 || versionsUpdated > 0)
@@ -98,7 +127,7 @@ public async Task<bool> UpsertCivitModelAsync(IEnumerable<CivitModel> civitModel
     // Add to cache
     public async Task<bool> UpsertCivitModelQueryCacheEntryAsync(CivitModelQueryCacheEntry entry)
     {
-        var changed = await CivitModelQueryCache.UpsertAsync(entry);
+        var changed = await CivitModelQueryCache.UpsertAsync(entry).ConfigureAwait(false);
         if (changed)
         {
             CivitModelsChanged?.Invoke(this, EventArgs.Empty);
@@ -107,8 +136,14 @@ public async Task<bool> UpsertCivitModelQueryCacheEntryAsync(CivitModelQueryCach
         return changed;
     }
 
-    public Task<GithubCacheEntry?> GetGithubCacheEntry(string cacheKey) =>
-        GithubCache.FindByIdAsync(cacheKey);
+    public async Task<GithubCacheEntry?> GetGithubCacheEntry(string cacheKey)
+    {
+        if (await GithubCache.FindByIdAsync(cacheKey).ConfigureAwait(false) is { } result)
+        {
+            return result;
+        }
+        return null;
+    }
 
     public Task<bool> UpsertGithubCacheEntry(GithubCacheEntry cacheEntry) =>
         GithubCache.UpsertAsync(cacheEntry);

diff --git a/StabilityMatrix.Core/Models/Packages/BaseGitPackage.cs b/StabilityMatrix.Core/Models/Packages/BaseGitPackage.cs
@@ -105,13 +105,13 @@ public override async Task<IEnumerable<PackageVersion>> GetAllVersions(bool isRe
     /// <summary>
     /// Setup the virtual environment for the package.
     /// </summary>
-    /// <param name="installedPackagePath"></param>
-    /// <param name="venvName"></param>
-    /// <returns></returns>
     [MemberNotNull(nameof(VenvRunner))]
-    protected async Task<PyVenvRunner> SetupVenv(string installedPackagePath, string venvName = "venv")
+    public async Task<PyVenvRunner> SetupVenv(
+        string installedPackagePath, 
+        string venvName = "venv",
+        bool forceRecreate = false)
     {
-        var venvPath = Path.Combine(installedPackagePath, "venv");
+        var venvPath = Path.Combine(installedPackagePath, venvName);
         if (VenvRunner != null)
         {
             await VenvRunner.DisposeAsync().ConfigureAwait(false);
@@ -123,9 +123,9 @@ protected async Task<PyVenvRunner> SetupVenv(string installedPackagePath, string
             EnvironmentVariables = SettingsManager.Settings.EnvironmentVariables,
         };
 
-        if (!VenvRunner.Exists())
+        if (!VenvRunner.Exists() || forceRecreate)
         {
-            await VenvRunner.Setup().ConfigureAwait(false);
+            await VenvRunner.Setup(forceRecreate).ConfigureAwait(false);
         }
         return VenvRunner;
     }