Remove from cache when done downloading

LykosAI · Aug 1, 2023 · 5a771df · 5a771df
1 parent b10a1ae
commit 5a771df
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 2 deletions.
diff --git a/StabilityMatrix.Avalonia/ViewModels/CheckpointBrowserViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/CheckpointBrowserViewModel.cs
@@ -217,9 +217,15 @@ private void UpdateModelCards(IEnumerable<CivitModel>? models, CivitMetadata? me
             var updateCards = models
                 .Select(model =>
                 {
-                    if (cache.Get(model.Id) != null)
+                    var cachedViewModel = cache.Get(model.Id);
+                    if (cachedViewModel != null)
                     {
-                        return cache.Get(model.Id);
+                        if (!cachedViewModel.IsImporting)
+                        {
+                            cache.Remove(model.Id);
+                        }
+
+                        return cachedViewModel;
                     }
 
                     var newCard = new CheckpointBrowserCardViewModel(model,

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