From 77a4aa441fdf18d4620014d832baea44a12a7baf Mon Sep 17 00:00:00 2001
From: Ionite <dev@ionite.io>
Date: Sun, 10 Sep 2023 20:45:26 -0400
Subject: [PATCH] Invalidate cached images on deletion

---
 .../FallbackRamCachedWebImageLoader.cs        | 34 ++++++++++++++-----
 .../Inference/ImageFolderCardViewModel.cs     | 26 ++++----------
 2 files changed, 32 insertions(+), 28 deletions(-)

diff --git a/StabilityMatrix.Avalonia/FallbackRamCachedWebImageLoader.cs b/StabilityMatrix.Avalonia/FallbackRamCachedWebImageLoader.cs
index 5b79748e..9dbe6ef9 100644
--- a/StabilityMatrix.Avalonia/FallbackRamCachedWebImageLoader.cs
+++ b/StabilityMatrix.Avalonia/FallbackRamCachedWebImageLoader.cs
@@ -1,10 +1,12 @@
 using System;
+using System.Collections.Concurrent;
 using System.Diagnostics.CodeAnalysis;
 using System.IO;
 using System.Threading.Tasks;
 using AsyncAwaitBestPractices;
 using AsyncImageLoader.Loaders;
 using Avalonia.Media.Imaging;
+using StabilityMatrix.Core.Extensions;
 
 namespace StabilityMatrix.Avalonia;
 
@@ -14,15 +16,19 @@ public readonly record struct ImageLoadFailedEventArgs(string Url, Exception Exc
 public class FallbackRamCachedWebImageLoader : RamCachedWebImageLoader
 {
     private readonly WeakEventManager<ImageLoadFailedEventArgs> loadFailedEventManager = new();
-    
+
     public event EventHandler<ImageLoadFailedEventArgs> LoadFailed
     {
         add => loadFailedEventManager.AddEventHandler(value);
         remove => loadFailedEventManager.RemoveEventHandler(value);
     }
-    
-    protected void OnLoadFailed(string url, Exception exception) => loadFailedEventManager.RaiseEvent(
-        this, new ImageLoadFailedEventArgs(url, exception), nameof(LoadFailed));
+
+    protected void OnLoadFailed(string url, Exception exception) =>
+        loadFailedEventManager.RaiseEvent(
+            this,
+            new ImageLoadFailedEventArgs(url, exception),
+            nameof(LoadFailed)
+        );
 
     /// <summary>
     /// Attempts to load bitmap
@@ -44,17 +50,19 @@ public class FallbackRamCachedWebImageLoader : RamCachedWebImageLoader
                 return null;
             }
         }
-        
-        var internalOrCachedBitmap = 
+
+        var internalOrCachedBitmap =
             await LoadFromInternalAsync(url).ConfigureAwait(false)
             ?? await LoadFromGlobalCache(url).ConfigureAwait(false);
-        
-        if (internalOrCachedBitmap != null) return internalOrCachedBitmap;
+
+        if (internalOrCachedBitmap != null)
+            return internalOrCachedBitmap;
 
         try
         {
             var externalBytes = await LoadDataFromExternalAsync(url).ConfigureAwait(false);
-            if (externalBytes == null) return null;
+            if (externalBytes == null)
+                return null;
 
             using var memoryStream = new MemoryStream(externalBytes);
             var bitmap = new Bitmap(memoryStream);
@@ -67,4 +75,12 @@ public class FallbackRamCachedWebImageLoader : RamCachedWebImageLoader
         }
     }
 
+    public void RemoveFromCache(string url)
+    {
+        // ConcurrentDictionary<string, Task<Bitmap?>> _memoryCache
+        var cache =
+            this.GetPrivateField<ConcurrentDictionary<string, Task<Bitmap?>>>("_memoryCache")
+            ?? throw new NullReferenceException("Memory cache not found");
+        cache.TryRemove(url, out _);
+    }
 }
diff --git a/StabilityMatrix.Avalonia/ViewModels/Inference/ImageFolderCardViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/Inference/ImageFolderCardViewModel.cs
index c2db1d97..3ef00b3b 100644
--- a/StabilityMatrix.Avalonia/ViewModels/Inference/ImageFolderCardViewModel.cs
+++ b/StabilityMatrix.Avalonia/ViewModels/Inference/ImageFolderCardViewModel.cs
@@ -1,6 +1,7 @@
 using System;
 using System.IO;
 using System.Threading.Tasks;
+using AsyncImageLoader;
 using Avalonia.Controls.Notifications;
 using Avalonia.Controls.Primitives;
 using Avalonia.Input;
@@ -64,25 +65,6 @@ public partial class ImageFolderCardViewModel : ViewModelBase
         this.settingsManager = settingsManager;
         this.notificationService = notificationService;
 
-        // var minDatetime = DateTimeOffset.FromUnixTimeMilliseconds(0);
-
-        /*localImagesSource
-            .Connect()
-            .DeferUntilLoaded()
-            .Transform(
-                imageFile =>
-                    new ImageFolderCardItemViewModel
-                    {
-                        LocalImageFile = imageFile,
-                        ImagePath = Design.IsDesignMode
-                            ? imageFile.RelativePath
-                            : imageFile.GetFullPath(settingsManager.ImagesDirectory)
-                    }
-            )
-            .SortBy(x => x.LocalImageFile?.LastModifiedAt ?? minDatetime, SortDirection.Descending)
-            .Bind(Items)
-            .Subscribe();*/
-
         localImagesSource
             .Connect()
             .DeferUntilLoaded()
@@ -163,6 +145,12 @@ public partial class ImageFolderCardViewModel : ViewModelBase
 
         // Remove from index
         await imageIndexService.RemoveImage(item);
+
+        // Invalidate cache
+        if (ImageLoader.AsyncImageLoader is FallbackRamCachedWebImageLoader loader)
+        {
+            loader.RemoveFromCache(imagePath);
+        }
     }
 
     /// <summary>