Merge pull request LykosAI#184 from ionite34/cache-downloading-model-…

…cards Cache CheckpointBrowserCardViewModels when importing
LykosAI · Aug 1, 2023 · fbdbe49 · fbdbe49
2 parents 72f7f41 + 5a771df
commit fbdbe49
Show file tree

Hide file tree

Showing 4 changed files with 55 additions and 17 deletions.
diff --git a/StabilityMatrix.Avalonia/ViewModels/CheckpointBrowserCardViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/CheckpointBrowserCardViewModel.cs
@@ -38,6 +38,7 @@ public partial class CheckpointBrowserCardViewModel : ProgressViewModel
     private readonly ISettingsManager settingsManager;
     private readonly ServiceManager<ViewModelBase> dialogFactory;
     private readonly INotificationService notificationService;
+    private readonly Action<CheckpointBrowserCardViewModel>? onDownloadStart;
     public CivitModel CivitModel { get; init; }
     public override bool IsTextVisible => Value > 0;
 
@@ -51,12 +52,14 @@ public CheckpointBrowserCardViewModel(
         IDownloadService downloadService,
         ISettingsManager settingsManager,
         ServiceManager<ViewModelBase> dialogFactory,
-        INotificationService notificationService)
+        INotificationService notificationService,
+        Action<CheckpointBrowserCardViewModel>? onDownloadStart = null)
     {
         this.downloadService = downloadService;
         this.settingsManager = settingsManager;
         this.dialogFactory = dialogFactory;
         this.notificationService = notificationService;
+        this.onDownloadStart = onDownloadStart;
         CivitModel = civitModel;
 
         UpdateImage();
@@ -194,6 +197,8 @@ private async Task DoImport(CivitModel model, CivitModelVersion? selectedVersion
         IsImporting = true;
         Text = "Downloading...";
 
+        onDownloadStart?.Invoke(this);
+
         // Holds files to be deleted on errors
         var filesForCleanup = new HashSet<FilePath>();
 

diff --git a/StabilityMatrix.Avalonia/ViewModels/CheckpointBrowserViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/CheckpointBrowserViewModel.cs
@@ -21,6 +21,7 @@
 using StabilityMatrix.Core.Database;
 using StabilityMatrix.Core.Extensions;
 using StabilityMatrix.Core.Helper;
+using StabilityMatrix.Core.Helper.Cache;
 using StabilityMatrix.Core.Models;
 using StabilityMatrix.Core.Models.Api;
 using StabilityMatrix.Core.Models.Settings;
@@ -41,6 +42,7 @@ public partial class CheckpointBrowserViewModel : PageViewModelBase
     private readonly ILiteDbContext liteDbContext;
     private readonly INotificationService notificationService;
     private const int MaxModelsPerPage = 14;
+    private LRUCache<int /* model id */, CheckpointBrowserCardViewModel> cache = new(50);
 
     [ObservableProperty] private ObservableCollection<CheckpointBrowserCardViewModel>? modelCards;
     [ObservableProperty] private DataGridCollectionView? modelCardsView;
@@ -213,8 +215,29 @@ private void UpdateModelCards(IEnumerable<CivitModel>? models, CivitMetadata? me
         else
         {
             var updateCards = models
-                .Select(model => new CheckpointBrowserCardViewModel(model,
-                    downloadService, settingsManager, dialogFactory, notificationService)).ToList();
+                .Select(model =>
+                {
+                    var cachedViewModel = cache.Get(model.Id);
+                    if (cachedViewModel != null)
+                    {
+                        if (!cachedViewModel.IsImporting)
+                        {
+                            cache.Remove(model.Id);
+                        }
+
+                        return cachedViewModel;
+                    }
+
+                    var newCard = new CheckpointBrowserCardViewModel(model,
+                        downloadService, settingsManager, dialogFactory, notificationService,
+                        viewModel =>
+                        {
+                            if (cache.Get(viewModel.CivitModel.Id) != null) return;
+                            cache.Add(viewModel.CivitModel.Id, viewModel);
+                        });
+                    return newCard;
+                }).ToList();
+
             allModelCards = updateCards;
             ModelCards =
                 new ObservableCollection<CheckpointBrowserCardViewModel>(

diff --git a/StabilityMatrix.Avalonia/ViewModels/PackageManagerViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/PackageManagerViewModel.cs
@@ -289,17 +289,17 @@ private void DeleteDirectory(string targetDirectory)
     private async Task UpdateSelectedPackage()
     {
         if (SelectedPackage == null) return;
-
-        var package = packageFactory.FindPackageByName(SelectedPackage.PackageName);
+        var installedPackage = SelectedPackage;
+        var package = packageFactory.FindPackageByName(installedPackage.PackageName);
         if (package == null)
         {
             logger.LogError("Could not find package {SelectedPackagePackageName}", 
-                SelectedPackage.PackageName);
+                installedPackage.PackageName);
             return;
         }
 
-        ProgressText = $"Updating {SelectedPackage.DisplayName} to latest version...";
-        package.InstallLocation = SelectedPackage.FullPath!;
+        ProgressText = $"Updating {installedPackage.DisplayName} to latest version...";
+        package.InstallLocation = installedPackage.FullPath!;
         var progress = new Progress<ProgressReport>(progress =>
         {
             var percent = Convert.ToInt32(progress.Percentage);
@@ -311,14 +311,14 @@ private async Task UpdateSelectedPackage()
             EventManager.Instance.OnGlobalProgressChanged(percent);
         });
 
-        var updateResult = await package.Update(SelectedPackage, progress);
+        var updateResult = await package.Update(installedPackage, progress);
 
         if (string.IsNullOrWhiteSpace(updateResult))
         {
             var errorMsg =
-                $"There was an error updating {SelectedPackage.DisplayName}. Please try again later.";
+                $"There was an error updating {installedPackage.DisplayName}. Please try again later.";
 
-            if (SelectedPackage.PackageName == "automatic")
+            if (installedPackage.PackageName == "automatic")
             {
                 errorMsg = errorMsg.Replace("Please try again later.",
                     "Please stash any changes before updating, or manually update the package.");
@@ -329,10 +329,15 @@ private async Task UpdateSelectedPackage()
                 errorMsg, NotificationType.Error));
         }
 
-        settingsManager.UpdatePackageVersionNumber(SelectedPackage.Id, updateResult);
+        settingsManager.UpdatePackageVersionNumber(installedPackage.Id, updateResult);
         notificationService.Show("Update complete",
-            $"{SelectedPackage.DisplayName} has been updated to the latest version.",
+            $"{installedPackage.DisplayName} has been updated to the latest version.",
             NotificationType.Success);
+
+        installedPackage.UpdateAvailable = false;
+        UpdateAvailable = false;
+        InstallButtonText = "Launch";
+
         await OnLoadedAsync();
         await DelayedClearProgress(TimeSpan.FromSeconds(3));
     }
@@ -369,9 +374,5 @@ private async Task DelayedClearProgress(TimeSpan delay)
         ProgressText = string.Empty;
         ProgressValue = 0;
         IsIndeterminate = false;
-
-        SelectedPackage.UpdateAvailable = false;
-        UpdateAvailable = false;
-        InstallButtonText = "Launch";
     }
 }
diff --git a/StabilityMatrix.Core/Helper/Cache/LRUCache.cs b/StabilityMatrix.Core/Helper/Cache/LRUCache.cs
@@ -52,6 +52,15 @@ public void Add(TK key, TV val)
         cacheMap[key] = node;
     }
 
+    [MethodImpl(MethodImplOptions.Synchronized)]
+    public void Remove(TK key)
+    {
+        if (!cacheMap.TryGetValue(key, out var node)) return;
+
+        lruList.Remove(node);
+        cacheMap.Remove(key);
+    }
+
     private void RemoveFirst()
     {
         // Remove from LRUPriority