Checkpoints Page Improvements

- Improved Checkpoints page indexing performance
- Fixed issue where Checkpoints page may not show all checkpoints after clearing search filter
- Fixed issue where Checkpoints page may show incorrect checkpoints for the given filter after changing pages

CHANGELOG.md

@@ -13,6 +13,9 @@ and this project adheres to [Semantic Versioning 2.0](https://semver.org/spec/v2
 ### Fixed
 - Fixed SD.Next shared folders config not working with new config format, reverted to Junctions / Symlinks
 - Fixed issue where model version wouldn't be selected in the "All Versions" section of the Model Browser
+- Improved Checkpoints page indexing performance
+- Fixed issue where Checkpoints page may not show all checkpoints after clearing search filter
+- Fixed issue where Checkpoints page may show incorrect checkpoints for the given filter after changing pages
 
 ## v2.2.0
 ### Added

StabilityMatrix.Avalonia/ViewModels/CheckpointManager/CheckpointFile.cs

@@ -196,43 +196,33 @@ public partial class CheckpointFile : ViewModelBase
     /// </summary>
     public static IEnumerable<CheckpointFile> FromDirectoryIndex(string directory, SearchOption searchOption = SearchOption.TopDirectoryOnly)
     {
-        // Get all files with supported extensions
-        var allExtensions = SupportedCheckpointExtensions
-            .Concat(SupportedImageExtensions)
-            .Concat(SupportedMetadataExtensions);
-
-        var files = allExtensions.AsParallel()
-            .SelectMany(pattern => Directory.EnumerateFiles(directory, $"*{pattern}", searchOption)).ToDictionary<string, string>(Path.GetFileName);
-
-        foreach (var file in files.Keys.Where(k => SupportedCheckpointExtensions.Contains(Path.GetExtension(k))))
+        foreach (var file in Directory.EnumerateFiles(directory, "*.*", searchOption))
         {
-            var checkpointFile = new CheckpointFile()
+            if (!SupportedCheckpointExtensions.Any(ext => file.Contains(ext)))
+                continue;
+            
+            var checkpointFile = new CheckpointFile
             {
                 Title = Path.GetFileNameWithoutExtension(file),
                 FilePath = Path.Combine(directory, file),
             };
             
-            // Check for connected model info
-            var fileNameWithoutExtension = Path.GetFileNameWithoutExtension(file);
-            var cmInfoPath = $"{fileNameWithoutExtension}.cm-info.json";
-            if (files.TryGetValue(cmInfoPath, out var jsonPath))
+            var jsonPath = Path.Combine(directory, $"{Path.GetFileNameWithoutExtension(file)}.cm-info.json");
+            if (File.Exists(jsonPath)) 
             {
-                try
-                {
-                    var jsonData = File.ReadAllText(jsonPath);
-                    checkpointFile.ConnectedModel = ConnectedModelInfo.FromJson(jsonData);
-                }
-                catch (IOException e)
-                {
-                    Debug.WriteLine($"Failed to parse {cmInfoPath}: {e}");
-                }
+                var json = File.ReadAllText(jsonPath);
+                var connectedModelInfo = ConnectedModelInfo.FromJson(json);
+                checkpointFile.ConnectedModel = connectedModelInfo;
             }
 
-            // Check for preview image
-            var previewImage = SupportedImageExtensions.Select(ext => $"{fileNameWithoutExtension}.preview{ext}").FirstOrDefault(files.ContainsKey);
-            if (previewImage != null)
+            var possibleImagePaths = SupportedImageExtensions.Select(ext =>
+                    Path.Combine(directory,
+                        $"{Path.GetFileNameWithoutExtension(file)}.preview{ext}"))
+                .Where(File.Exists).ToList();
+
+            if (possibleImagePaths.Any())
             {
-                checkpointFile.PreviewImagePath = files[previewImage];
+                checkpointFile.PreviewImagePath = possibleImagePaths.First();
             }
 
             yield return checkpointFile;

StabilityMatrix.Avalonia/ViewModels/CheckpointsPageViewModel.cs

@@ -1,5 +1,6 @@
 using System;
 using System.Collections.ObjectModel;
+using System.Diagnostics;
 using System.IO;
 using System.Linq;
 using System.Threading.Tasks;
@@ -10,6 +11,7 @@ using CommunityToolkit.Mvvm.Input;
 using FluentAvalonia.UI.Controls;
 using NLog;
 using StabilityMatrix.Avalonia.ViewModels.Base;
+using StabilityMatrix.Avalonia.ViewModels.CheckpointManager;
 using StabilityMatrix.Avalonia.Views;
 using StabilityMatrix.Core.Attributes;
 using StabilityMatrix.Core.Helper;
@@ -95,7 +97,12 @@ public partial class CheckpointsPageViewModel : PageViewModelBase
     {
         if (string.IsNullOrWhiteSpace(SearchFilter))
         {
-            DisplayedCheckpointFolders = CheckpointFolders;
+            DisplayedCheckpointFolders = new ObservableCollection<CheckpointFolder>(
+                CheckpointFolders.Select(x =>
+                {
+                    x.SearchFilter = SearchFilter;
+                    return x;
+                }));
             return;
         }
         
@@ -106,7 +113,7 @@ public partial class CheckpointsPageViewModel : PageViewModelBase
             folder.SearchFilter = SearchFilter;
         }
 
-        DisplayedCheckpointFolders = new ObservableCollection<CheckpointManager.CheckpointFolder>(filteredFolders);
+        DisplayedCheckpointFolders = new ObservableCollection<CheckpointFolder>(filteredFolders);
     }
     
     private bool ContainsSearchFilter(CheckpointManager.CheckpointFolder folder)
@@ -136,8 +143,10 @@ public partial class CheckpointsPageViewModel : PageViewModelBase
 
         var folders = Directory.GetDirectories(modelsDirectory);
 
+        var sw = new Stopwatch();
+        sw.Start();
         // Index all folders
-        var indexTasks = folders.Select(f => Task.Run(async () =>
+        var indexTasks = folders.Select(async f =>
         {
             var checkpointFolder =
                 new CheckpointManager.CheckpointFolder(settingsManager, downloadService, modelFinder)
@@ -148,21 +157,29 @@ public partial class CheckpointsPageViewModel : PageViewModelBase
                 };
             await checkpointFolder.IndexAsync();
             return checkpointFolder;
-        })).ToList();
+        }).ToList();
 
         await Task.WhenAll(indexTasks);
 
+        sw.Stop();
+        Logger.Debug("Indexed {Count} folders in {Elapsed} ms", indexTasks.Count, sw.ElapsedMilliseconds);
+
         // Set new observable collection, ordered by alphabetical order
         CheckpointFolders =
-            new ObservableCollection<CheckpointManager.CheckpointFolder>(indexTasks
+            new ObservableCollection<CheckpointFolder>(indexTasks
                 .Select(t => t.Result)
                 .OrderBy(f => f.Title));
 
         if (!string.IsNullOrWhiteSpace(SearchFilter))
         {
-            DisplayedCheckpointFolders = new ObservableCollection<CheckpointManager.CheckpointFolder>(
-                CheckpointFolders
-                    .Where(x => x.CheckpointFiles.Any(y => y.FileName.Contains(SearchFilter))));
+            var filtered = CheckpointFolders
+                .Where(x => x.CheckpointFiles.Any(y => y.FileName.Contains(SearchFilter))).Select(
+                    f =>
+                    {
+                        f.SearchFilter = SearchFilter;
+                        return f;
+                    });
+            DisplayedCheckpointFolders = new ObservableCollection<CheckpointFolder>(filtered);
         }
         else
         {

StabilityMatrix.Avalonia/Views/Dialogs/SelectDataDirectoryDialog.axaml

@@ -50,7 +50,7 @@
             </Grid>
             
             <TextBlock
-                Text="This is where the model checkpoints, LORAs, web UIs, settings, etc. will be installed. If you were satisfied with the previous versions, you don't need to change anything here."
+                Text="This is where the model checkpoints, LORAs, web UIs, settings, etc. will be installed."
                 TextWrapping="Wrap"
                 Foreground="LightGray"
                 FontSize="12"
@@ -59,7 +59,7 @@
             <CheckBox
                 Content="Portable Mode"
                 IsChecked="{Binding IsPortableMode, Mode=TwoWay}"
-                Margin="0,16,0,0" />
+                Margin="0,32,0,0" />
 
             <ui:InfoBar
                 IsClosable="False"