From 94115fd7007391dcf06c9be63794c8dbcf78676a Mon Sep 17 00:00:00 2001
From: JT <mohnjiles@gmail.com>
Date: Sun, 20 Aug 2023 16:01:08 -0700
Subject: [PATCH 1/8] Added Fooocus & select new data directory button

---
 CHANGELOG.md                                  |   4 +
 StabilityMatrix.Avalonia/App.axaml.cs         |   1 +
 .../StabilityMatrix.Avalonia.csproj           |   2 +-
 .../Dialogs/SelectModelVersionViewModel.cs    |   7 +-
 .../ViewModels/SettingsViewModel.cs           |  36 +++++
 .../Views/SettingsPage.axaml                  |  43 ++++--
 StabilityMatrix.Core/Helper/HardwareHelper.cs |  10 ++
 .../Models/Packages/Fooocus.cs                | 124 ++++++++++++++++++
 .../Models/Packages/VladAutomatic.cs          |  18 +--
 9 files changed, 214 insertions(+), 31 deletions(-)
 create mode 100644 StabilityMatrix.Core/Models/Packages/Fooocus.cs

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 764ee032..ebe4364c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -6,6 +6,10 @@ 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.2.1
+### Added
+- New installable Package - [Fooocus](https://github.com/lllyasviel/Fooocus)
+- Added "Select New Data Directory" button to Settings
+
 ### Fixed
 - Fixed SD.Next shared folders config not working with new config format, reverted to Junctions / Symlinks
 
diff --git a/StabilityMatrix.Avalonia/App.axaml.cs b/StabilityMatrix.Avalonia/App.axaml.cs
index 382f7274..487b73a0 100644
--- a/StabilityMatrix.Avalonia/App.axaml.cs
+++ b/StabilityMatrix.Avalonia/App.axaml.cs
@@ -308,6 +308,7 @@ public sealed class App : Application
         services.AddSingleton<BasePackage, ComfyUI>();
         services.AddSingleton<BasePackage, VoltaML>();
         services.AddSingleton<BasePackage, InvokeAI>();
+        services.AddSingleton<BasePackage, Fooocus>();
     }
 
     private static IServiceCollection ConfigureServices()
diff --git a/StabilityMatrix.Avalonia/StabilityMatrix.Avalonia.csproj b/StabilityMatrix.Avalonia/StabilityMatrix.Avalonia.csproj
index 2acc5271..f9d7a498 100644
--- a/StabilityMatrix.Avalonia/StabilityMatrix.Avalonia.csproj
+++ b/StabilityMatrix.Avalonia/StabilityMatrix.Avalonia.csproj
@@ -8,7 +8,7 @@
         <ApplicationManifest>app.manifest</ApplicationManifest>
         <AvaloniaUseCompiledBindingsByDefault>true</AvaloniaUseCompiledBindingsByDefault>
         <ApplicationIcon>./Assets/Icon.ico</ApplicationIcon>
-        <Version>2.1.1-dev.1</Version>
+        <Version>2.3.0-dev.1</Version>
         <InformationalVersion>$(Version)</InformationalVersion>
         <EnableWindowsTargeting>true</EnableWindowsTargeting>
     </PropertyGroup>
diff --git a/StabilityMatrix.Avalonia/ViewModels/Dialogs/SelectModelVersionViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/Dialogs/SelectModelVersionViewModel.cs
index d6435c85..082859db 100644
--- a/StabilityMatrix.Avalonia/ViewModels/Dialogs/SelectModelVersionViewModel.cs
+++ b/StabilityMatrix.Avalonia/ViewModels/Dialogs/SelectModelVersionViewModel.cs
@@ -42,8 +42,11 @@ public partial class SelectModelVersionViewModel : ContentDialogViewModelBase
         var firstImageUrl = value?.ModelVersion?.Images?.FirstOrDefault(
             img => nsfwEnabled || img.Nsfw == "None")?.Url;
 
-        Dispatcher.UIThread.InvokeAsync(async
-            () => await UpdateImage(firstImageUrl));
+        Dispatcher.UIThread.InvokeAsync(async () =>
+        {
+            SelectedFile = value?.CivitFileViewModels.FirstOrDefault();
+            await UpdateImage(firstImageUrl);
+        });
     }
 
     partial void OnSelectedFileChanged(CivitFileViewModel? value)
diff --git a/StabilityMatrix.Avalonia/ViewModels/SettingsViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/SettingsViewModel.cs
index 54026801..7c013f2e 100644
--- a/StabilityMatrix.Avalonia/ViewModels/SettingsViewModel.cs
+++ b/StabilityMatrix.Avalonia/ViewModels/SettingsViewModel.cs
@@ -329,6 +329,42 @@ public partial class SettingsViewModel : PageViewModelBase
             "Stability Matrix has been added to the Start Menu for all users.", NotificationType.Success);
     }
 
+    public async Task PickNewDataDirectory()
+    {
+        var viewModel = dialogFactory.Get<SelectDataDirectoryViewModel>();
+        var dialog = new BetterContentDialog
+        {
+            IsPrimaryButtonEnabled = false,
+            IsSecondaryButtonEnabled = false,
+            IsFooterVisible = false,
+            Content = new SelectDataDirectoryDialog
+            {
+                DataContext = viewModel
+            }
+        };
+
+        var result = await dialog.ShowAsync();
+        if (result == ContentDialogResult.Primary)
+        {
+            // 1. For portable mode, call settings.SetPortableMode()
+            if (viewModel.IsPortableMode)
+            {
+                settingsManager.SetPortableMode();
+            }
+            // 2. For custom path, call settings.SetLibraryPath(path)
+            else
+            {
+                settingsManager.SetLibraryPath(viewModel.DataDirectory);
+            }
+            
+            // Try to find library again, should be found now
+            if (!settingsManager.TryFindLibrary())
+            {
+                throw new Exception("Could not find library after setting path");
+            }            
+        }
+    }
+
     #endregion
     
     #region Debug Section
diff --git a/StabilityMatrix.Avalonia/Views/SettingsPage.axaml b/StabilityMatrix.Avalonia/Views/SettingsPage.axaml
index 8e3d6d0d..01443d17 100644
--- a/StabilityMatrix.Avalonia/Views/SettingsPage.axaml
+++ b/StabilityMatrix.Avalonia/Views/SettingsPage.axaml
@@ -6,6 +6,7 @@
              xmlns:ui="using:FluentAvalonia.UI.Controls"
              xmlns:mocks="clr-namespace:StabilityMatrix.Avalonia.DesignData"
              xmlns:controls="clr-namespace:StabilityMatrix.Avalonia.Controls"
+             xmlns:avalonia="clr-namespace:Projektanker.Icons.Avalonia;assembly=Projektanker.Icons.Avalonia"
              mc:Ignorable="d" d:DesignWidth="800" d:DesignHeight="700"
              x:DataType="vm:SettingsViewModel"
              x:CompileBindings="True"
@@ -124,7 +125,7 @@
                     </Grid>
                     
                     <!-- System Options -->
-                    <Grid Grid.Row="4" Margin="0,8,0,0" RowDefinitions="auto,*">
+                    <Grid Grid.Row="4" Margin="0,8,0,0" RowDefinitions="auto, auto, auto">
                         <TextBlock 
                             FontWeight="Medium"
                             Text="System" 
@@ -166,6 +167,24 @@
                                 </StackPanel>
                             </ui:SettingsExpander.Footer>
                         </ui:SettingsExpander>
+                        <ui:SettingsExpander Grid.Row="2"
+                                             Header="Select New Data Directory"
+                                             Description="Does not move existing data"
+                                             IconSource="MoveToFolder"
+                                             Margin="8,0">
+                            <ui:SettingsExpander.Footer>
+                                <Button Command="{Binding PickNewDataDirectory}">
+                                    <Grid ColumnDefinitions="Auto, Auto">
+                                        <avalonia:Icon Grid.Row="0" Value="fa-solid fa-folder-open"
+                                                     Margin="0,0,8,0"
+                                                     VerticalAlignment="Center" />
+                                        <TextBlock Grid.Column="1"
+                                                   VerticalAlignment="Center"
+                                                   Text="Select Directory"/>
+                                    </Grid>
+                                </Button>
+                            </ui:SettingsExpander.Footer>
+                        </ui:SettingsExpander>
                     </Grid>
                     
                     <!-- Debug Options -->
@@ -181,24 +200,24 @@
                                 IconSource="Code"
                                 Command="{Binding LoadDebugInfo}"
                                 Header="Debug Options"
-                                Margin="8, 0,8,4">
+                                Margin="8, 0,8,0">
                                 
                                 <ui:SettingsExpanderItem Description="Paths" IconSource="Folder"
-                                                         Margin="4">
+                                                         Margin="4, 0">
                                     <SelectableTextBlock Text="{Binding DebugPaths}"
                                                          Foreground="{DynamicResource TextControlPlaceholderForeground}"
                                                          TextWrapping="WrapWithOverflow" />
                                 </ui:SettingsExpanderItem>
 
                                 <ui:SettingsExpanderItem Description="Compat Info" IconSource="StarFilled"
-                                                         Margin="4,0,4,4">
+                                                         Margin="4,0">
                                     <SelectableTextBlock Text="{Binding DebugCompatInfo}"
                                                          Foreground="{DynamicResource TextControlPlaceholderForeground}"
                                                          TextWrapping="WrapWithOverflow" />
                                 </ui:SettingsExpanderItem>
                                 
                                 <ui:SettingsExpanderItem Description="GPU Info" IconSource="FullScreenMaximize"
-                                                         Margin="4,0,4,4">
+                                                         Margin="4,0">
                                     <SelectableTextBlock Text="{Binding DebugGpuInfo}"
                                                          Foreground="{DynamicResource TextControlPlaceholderForeground}"
                                                          TextWrapping="WrapWithOverflow" />
@@ -206,10 +225,9 @@
 
                                 <ui:SettingsExpanderItem Content="Animation Scale" IconSource="Clock"
                                                          Description="Lower values = faster animations. 0x means animations are instant."
-                                                         Margin="4,0,4,4">
+                                                         Margin="4,0">
                                     <ui:SettingsExpanderItem.Footer>
-                                        <ComboBox Margin="0, 8"
-                                                  ItemsSource="{Binding AnimationScaleOptions}"
+                                        <ComboBox ItemsSource="{Binding AnimationScaleOptions}"
                                                   SelectedItem="{Binding SelectedAnimationScale}">
                                             <ComboBox.ItemTemplate>
                                                 <DataTemplate>
@@ -223,30 +241,27 @@
                                 </ui:SettingsExpanderItem>
                                 
                                 <ui:SettingsExpanderItem Content="Notification" IconSource="CommentAdd"
-                                                         Margin="4,0,4,4">
+                                                         Margin="4,0">
                                     <ui:SettingsExpanderItem.Footer>
                                         <Button
-                                            Margin="0, 8"
                                             Command="{Binding DebugNotificationCommand}"
                                             Content="New Notification"/>
                                     </ui:SettingsExpanderItem.Footer>
                                 </ui:SettingsExpanderItem>
                                 
                                 <ui:SettingsExpanderItem Content="Content Dialog" IconSource="NewWindow"
-                                                         Margin="4,0,4,4">
+                                                         Margin="4,0">
                                     <ui:SettingsExpanderItem.Footer>
                                         <Button
-                                            Margin="0, 8"
                                             Command="{Binding DebugContentDialogCommand}"
                                             Content="Show Dialog"/>
                                     </ui:SettingsExpanderItem.Footer>
                                 </ui:SettingsExpanderItem>
                                 
                                 <ui:SettingsExpanderItem Content="Exceptions" IconSource="Flag"
-                                                         Margin="4,0,4,4">
+                                                         Margin="4,0">
                                     <ui:SettingsExpanderItem.Footer>
                                         <Button
-                                            Margin="0, 8"
                                             Command="{Binding DebugThrowExceptionCommand}"
                                             Content="Unhandled Exception"/>
                                     </ui:SettingsExpanderItem.Footer>
diff --git a/StabilityMatrix.Core/Helper/HardwareHelper.cs b/StabilityMatrix.Core/Helper/HardwareHelper.cs
index 66d678ca..eaef9df8 100644
--- a/StabilityMatrix.Core/Helper/HardwareHelper.cs
+++ b/StabilityMatrix.Core/Helper/HardwareHelper.cs
@@ -123,6 +123,16 @@ public static partial class HardwareHelper
     {
         return IterGpuInfo().Any(gpu => gpu.IsAmd);
     }
+    
+    // Set ROCm for default if AMD and Linux
+    public static bool PreferRocm() => !HardwareHelper.HasNvidiaGpu()
+                                       && HardwareHelper.HasAmdGpu()
+                                       && Compat.IsLinux;
+    
+    // Set DirectML for default if AMD and Windows
+    public static bool PreferDirectML() => !HardwareHelper.HasNvidiaGpu()
+                                           && HardwareHelper.HasAmdGpu()
+                                           && Compat.IsWindows;
 }
 
 public enum Level
diff --git a/StabilityMatrix.Core/Models/Packages/Fooocus.cs b/StabilityMatrix.Core/Models/Packages/Fooocus.cs
new file mode 100644
index 00000000..69444a73
--- /dev/null
+++ b/StabilityMatrix.Core/Models/Packages/Fooocus.cs
@@ -0,0 +1,124 @@
+using System.Diagnostics;
+using System.Text.RegularExpressions;
+using StabilityMatrix.Core.Helper;
+using StabilityMatrix.Core.Helper.Cache;
+using StabilityMatrix.Core.Models.Progress;
+using StabilityMatrix.Core.Processes;
+using StabilityMatrix.Core.Services;
+
+namespace StabilityMatrix.Core.Models.Packages;
+
+public class Fooocus : BaseGitPackage
+{
+    public Fooocus(IGithubApiCache githubApi, ISettingsManager settingsManager,
+        IDownloadService downloadService, IPrerequisiteHelper prerequisiteHelper) : base(githubApi,
+        settingsManager, downloadService, prerequisiteHelper)
+    {
+    }
+
+    public override string Name => "Fooocus";
+    public override string DisplayName { get; set; } = "Fooocus";
+    public override string Author => "lllyasviel";
+
+    public override string Blurb =>
+        "Fooocus is a rethinking of Stable Diffusion and Midjourney’s designs";
+
+    public override string LicenseType => "GPL-3.0";
+    public override string LicenseUrl => "https://github.com/lllyasviel/Fooocus/blob/main/LICENSE";
+    public override string LaunchCommand => "launch.py";
+
+    public override Uri PreviewImageUri =>
+        new("https://user-images.githubusercontent.com/19834515/261830306-f79c5981-cf80-4ee3-b06b-3fef3f8bfbc7.png");
+
+    public override List<LaunchOptionDefinition> LaunchOptions => new()
+    {
+        LaunchOptionDefinition.Extras
+    };
+
+    public override Dictionary<SharedFolderType, IReadOnlyList<string>> SharedFolders => new()
+    {
+        [SharedFolderType.StableDiffusion] = new[] {"models/checkpoints"},
+        [SharedFolderType.Diffusers] = new[] {"models/diffusers"},
+        [SharedFolderType.Lora] = new[] {"models/loras"},
+        [SharedFolderType.CLIP] = new[] {"models/clip"},
+        [SharedFolderType.TextualInversion] = new[] {"models/embeddings"},
+        [SharedFolderType.VAE] = new[] {"models/vae"},
+        [SharedFolderType.ApproxVAE] = new[] {"models/vae_approx"},
+        [SharedFolderType.ControlNet] = new[] {"models/controlnet"},
+        [SharedFolderType.GLIGEN] = new[] {"models/gligen"},
+        [SharedFolderType.ESRGAN] = new[] {"models/upscale_models"},
+        [SharedFolderType.Hypernetwork] = new[] {"models/hypernetworks"}
+    };
+    
+    public override async Task<string> GetLatestVersion()
+    {
+        var release = await GetLatestRelease().ConfigureAwait(false);
+        return release.TagName!;
+    }
+
+    public override async Task InstallPackage(IProgress<ProgressReport>? progress = null)
+    {
+        await base.InstallPackage(progress).ConfigureAwait(false);
+        var venvRunner = await SetupVenv(InstallLocation).ConfigureAwait(false);
+
+        progress?.Report(new ProgressReport(-1f, "Installing torch...", isIndeterminate: true));
+        
+        var torchVersion = "cpu";
+        var gpus = HardwareHelper.IterGpuInfo().ToList();
+        
+        // If has Nvidia Gpu, install CUDA version
+        if (gpus.Any(g => g.IsNvidia))
+        {
+            torchVersion = "cu118";
+        }
+        // For AMD, Install ROCm version
+        else if (gpus.Any(g => g.IsAmd))
+        {
+            torchVersion= "rocm5.4.2";
+        }
+
+        await venvRunner
+            .PipInstall(
+                $"torch==2.0.1 torchvision==0.15.2 --extra-index-url https://download.pytorch.org/whl/{torchVersion}",
+                OnConsoleOutput).ConfigureAwait(false);
+
+        progress?.Report(new ProgressReport(-1f, "Installing requirements...",
+            isIndeterminate: true));
+        await venvRunner.PipInstall("-r requirements_versions.txt", OnConsoleOutput)
+            .ConfigureAwait(false);
+    }
+
+    public override async Task RunPackage(string installedPackagePath, string command, string arguments)
+    {
+        await SetupVenv(installedPackagePath).ConfigureAwait(false);
+
+        void HandleConsoleOutput(ProcessOutput s)
+        {
+            OnConsoleOutput(s);
+            
+            if (s.Text.Contains("To see the GUI go to", StringComparison.OrdinalIgnoreCase))
+            {
+                var regex = new Regex(@"(https?:\/\/)([^:\s]+):(\d+)");
+                var match = regex.Match(s.Text);
+                if (match.Success)
+                {
+                    WebUrl = match.Value;
+                }
+                OnStartupComplete(WebUrl);
+            }
+        }
+
+        void HandleExit(int i)
+        {
+            Debug.WriteLine($"Venv process exited with code {i}");
+            OnExit(i);
+        }
+
+        var args = $"\"{Path.Combine(installedPackagePath, command)}\" {arguments}";
+
+        VenvRunner?.RunDetached(
+            args.TrimEnd(), 
+            HandleConsoleOutput, 
+            HandleExit);
+    }
+}
diff --git a/StabilityMatrix.Core/Models/Packages/VladAutomatic.cs b/StabilityMatrix.Core/Models/Packages/VladAutomatic.cs
index bcf940a2..679840ee 100644
--- a/StabilityMatrix.Core/Models/Packages/VladAutomatic.cs
+++ b/StabilityMatrix.Core/Models/Packages/VladAutomatic.cs
@@ -98,7 +98,7 @@ public class VladAutomatic : BaseGitPackage
         {
             Name = "Use DirectML if no compatible GPU is detected",
             Type = LaunchOptionType.Bool,
-            InitialValue = PreferDirectML(),
+            InitialValue = HardwareHelper.PreferDirectML(),
             Options = new() { "--use-directml" }
         },
         new()
@@ -112,7 +112,7 @@ public class VladAutomatic : BaseGitPackage
         {
             Name = "Force use of AMD ROCm backend",
             Type = LaunchOptionType.Bool,
-            InitialValue = PreferRocm(),
+            InitialValue = HardwareHelper.PreferRocm(),
             Options = new() { "--use-rocm" }
         },
         new()
@@ -137,16 +137,6 @@ public class VladAutomatic : BaseGitPackage
     };
 
     public override string ExtraLaunchArguments => "";
-
-    // Set ROCm for default if AMD and Linux
-    private static bool PreferRocm() => !HardwareHelper.HasNvidiaGpu()
-                                     && HardwareHelper.HasAmdGpu()
-                                     && Compat.IsLinux;
-    
-    // Set DirectML for default if AMD and Windows
-    private static bool PreferDirectML() => !HardwareHelper.HasNvidiaGpu()
-                                            && HardwareHelper.HasAmdGpu()
-                                            && Compat.IsWindows;
     
     public override Task<string> GetLatestVersion() => Task.FromResult("master");
 
@@ -177,13 +167,13 @@ public class VladAutomatic : BaseGitPackage
             await venvRunner.CustomInstall("launch.py --use-cuda --debug --test", OnConsoleOutput)
                 .ConfigureAwait(false);
         }
-        else if (PreferRocm())
+        else if (HardwareHelper.PreferRocm())
         {
             // ROCm
             await venvRunner.CustomInstall("launch.py --use-rocm --debug --test", OnConsoleOutput)
                 .ConfigureAwait(false);
         }
-        else if (PreferDirectML())
+        else if (HardwareHelper.PreferDirectML())
         {
             // DirectML
             await venvRunner.CustomInstall("launch.py --use-directml --debug --test", OnConsoleOutput)

From 2de0ec729d65f5a60f8f6880f18db7512a189698 Mon Sep 17 00:00:00 2001
From: JT <mohnjiles@gmail.com>
Date: Sun, 20 Aug 2023 16:06:11 -0700
Subject: [PATCH 2/8] use PreferRocm instead

---
 StabilityMatrix.Core/Models/Packages/Fooocus.cs | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/StabilityMatrix.Core/Models/Packages/Fooocus.cs b/StabilityMatrix.Core/Models/Packages/Fooocus.cs
index 69444a73..965599f6 100644
--- a/StabilityMatrix.Core/Models/Packages/Fooocus.cs
+++ b/StabilityMatrix.Core/Models/Packages/Fooocus.cs
@@ -66,15 +66,13 @@ public class Fooocus : BaseGitPackage
         var torchVersion = "cpu";
         var gpus = HardwareHelper.IterGpuInfo().ToList();
         
-        // If has Nvidia Gpu, install CUDA version
         if (gpus.Any(g => g.IsNvidia))
         {
             torchVersion = "cu118";
         }
-        // For AMD, Install ROCm version
-        else if (gpus.Any(g => g.IsAmd))
+        else if (HardwareHelper.PreferRocm())
         {
-            torchVersion= "rocm5.4.2";
+            torchVersion = "rocm5.4.2";
         }
 
         await venvRunner

From e50b5a3d25bb096104cff21a8772fbacd9022d8d Mon Sep 17 00:00:00 2001
From: JT <mohnjiles@gmail.com>
Date: Sun, 20 Aug 2023 16:18:10 -0700
Subject: [PATCH 3/8] update chagenlog

---
 CHANGELOG.md | 1 +
 1 file changed, 1 insertion(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ebe4364c..9ef8d865 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -12,6 +12,7 @@ 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
 
 ## v2.2.0
 ### Added

From bc1bb6a75775331ded692e7371b24695e5b5cb0d Mon Sep 17 00:00:00 2001
From: JT <mohnjiles@gmail.com>
Date: Sun, 20 Aug 2023 18:10:36 -0700
Subject: [PATCH 4/8] show dialog & restart app when changing data directory

---
 .../ViewModels/SettingsViewModel.cs             | 17 +++++++++++++----
 1 file changed, 13 insertions(+), 4 deletions(-)

diff --git a/StabilityMatrix.Avalonia/ViewModels/SettingsViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/SettingsViewModel.cs
index 7c013f2e..000438bd 100644
--- a/StabilityMatrix.Avalonia/ViewModels/SettingsViewModel.cs
+++ b/StabilityMatrix.Avalonia/ViewModels/SettingsViewModel.cs
@@ -2,6 +2,7 @@
 using System.Collections.Generic;
 using System.Collections.ObjectModel;
 using System.ComponentModel;
+using System.Diagnostics;
 using System.IO;
 using System.Linq;
 using System.Reflection;
@@ -357,11 +358,19 @@ public partial class SettingsViewModel : PageViewModelBase
                 settingsManager.SetLibraryPath(viewModel.DataDirectory);
             }
             
-            // Try to find library again, should be found now
-            if (!settingsManager.TryFindLibrary())
+            // Restart
+            var restartDialog = new BetterContentDialog
             {
-                throw new Exception("Could not find library after setting path");
-            }            
+                Title = "Restart required",
+                Content = "Stability Matrix must be restarted for the changes to take effect.",
+                PrimaryButtonText = "Restart",
+                DefaultButton = ContentDialogButton.Primary,
+                IsSecondaryButtonEnabled = false,
+            };
+            await restartDialog.ShowAsync();
+            
+            Process.Start(Compat.AppCurrentPath);
+            Environment.Exit(0);
         }
     }
 

From a76b693aa1f17d6a0e4543655874275d47e00641 Mon Sep 17 00:00:00 2001
From: JT <mohnjiles@gmail.com>
Date: Sun, 20 Aug 2023 18:13:45 -0700
Subject: [PATCH 5/8] Shutdown instead of Exit

---
 StabilityMatrix.Avalonia/ViewModels/SettingsViewModel.cs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/StabilityMatrix.Avalonia/ViewModels/SettingsViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/SettingsViewModel.cs
index 000438bd..2e7faab7 100644
--- a/StabilityMatrix.Avalonia/ViewModels/SettingsViewModel.cs
+++ b/StabilityMatrix.Avalonia/ViewModels/SettingsViewModel.cs
@@ -370,7 +370,7 @@ public partial class SettingsViewModel : PageViewModelBase
             await restartDialog.ShowAsync();
             
             Process.Start(Compat.AppCurrentPath);
-            Environment.Exit(0);
+            App.Shutdown();
         }
     }
 

From 8c55676428fbe33f7df6fd097ba4697d3ddc330b Mon Sep 17 00:00:00 2001
From: JT <mohnjiles@gmail.com>
Date: Sun, 20 Aug 2023 19:34:53 -0700
Subject: [PATCH 6/8] 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                                  |  3 ++
 .../CheckpointManager/CheckpointFile.cs       | 44 +++++++------------
 .../ViewModels/CheckpointsPageViewModel.cs    | 33 ++++++++++----
 .../Dialogs/SelectDataDirectoryDialog.axaml   |  4 +-
 4 files changed, 47 insertions(+), 37 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9ef8d865..9d7a10ae 100644
--- a/CHANGELOG.md
+++ b/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
diff --git a/StabilityMatrix.Avalonia/ViewModels/CheckpointManager/CheckpointFile.cs b/StabilityMatrix.Avalonia/ViewModels/CheckpointManager/CheckpointFile.cs
index e540599c..5f25d91a 100644
--- a/StabilityMatrix.Avalonia/ViewModels/CheckpointManager/CheckpointFile.cs
+++ b/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;
diff --git a/StabilityMatrix.Avalonia/ViewModels/CheckpointsPageViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/CheckpointsPageViewModel.cs
index a5a59945..fe0e10d0 100644
--- a/StabilityMatrix.Avalonia/ViewModels/CheckpointsPageViewModel.cs
+++ b/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
         {
diff --git a/StabilityMatrix.Avalonia/Views/Dialogs/SelectDataDirectoryDialog.axaml b/StabilityMatrix.Avalonia/Views/Dialogs/SelectDataDirectoryDialog.axaml
index bdbccdc5..2fe90df8 100644
--- a/StabilityMatrix.Avalonia/Views/Dialogs/SelectDataDirectoryDialog.axaml
+++ b/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"

From d9e701ce92777c5dff7a0b1d0fe18c6afe507cb5 Mon Sep 17 00:00:00 2001
From: JT <mohnjiles@gmail.com>
Date: Sun, 20 Aug 2023 19:39:24 -0700
Subject: [PATCH 7/8] more linq

---
 .../ViewModels/CheckpointManager/CheckpointFile.cs  | 13 ++++---------
 1 file changed, 4 insertions(+), 9 deletions(-)

diff --git a/StabilityMatrix.Avalonia/ViewModels/CheckpointManager/CheckpointFile.cs b/StabilityMatrix.Avalonia/ViewModels/CheckpointManager/CheckpointFile.cs
index 5f25d91a..a8f83521 100644
--- a/StabilityMatrix.Avalonia/ViewModels/CheckpointManager/CheckpointFile.cs
+++ b/StabilityMatrix.Avalonia/ViewModels/CheckpointManager/CheckpointFile.cs
@@ -215,15 +215,10 @@ public partial class CheckpointFile : ViewModelBase
                 checkpointFile.ConnectedModel = connectedModelInfo;
             }
 
-            var possibleImagePaths = SupportedImageExtensions.Select(ext =>
-                    Path.Combine(directory,
-                        $"{Path.GetFileNameWithoutExtension(file)}.preview{ext}"))
-                .Where(File.Exists).ToList();
-
-            if (possibleImagePaths.Any())
-            {
-                checkpointFile.PreviewImagePath = possibleImagePaths.First();
-            }
+            checkpointFile.PreviewImagePath = SupportedImageExtensions
+                .Select(ext => Path.Combine(directory,
+                    $"{Path.GetFileNameWithoutExtension(file)}.preview{ext}")).Where(File.Exists)
+                .FirstOrDefault();
 
             yield return checkpointFile;
         }

From 1bfaa45fadbe8327dece1db7089c14502ac158b7 Mon Sep 17 00:00:00 2001
From: JT <mohnjiles@gmail.com>
Date: Sun, 20 Aug 2023 19:40:52 -0700
Subject: [PATCH 8/8] Remove stopwatch

---
 .../ViewModels/CheckpointsPageViewModel.cs                   | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/StabilityMatrix.Avalonia/ViewModels/CheckpointsPageViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/CheckpointsPageViewModel.cs
index fe0e10d0..19d1e84c 100644
--- a/StabilityMatrix.Avalonia/ViewModels/CheckpointsPageViewModel.cs
+++ b/StabilityMatrix.Avalonia/ViewModels/CheckpointsPageViewModel.cs
@@ -143,8 +143,6 @@ public partial class CheckpointsPageViewModel : PageViewModelBase
 
         var folders = Directory.GetDirectories(modelsDirectory);
 
-        var sw = new Stopwatch();
-        sw.Start();
         // Index all folders
         var indexTasks = folders.Select(async f =>
         {
@@ -161,9 +159,6 @@ public partial class CheckpointsPageViewModel : PageViewModelBase
 
         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<CheckpointFolder>(indexTasks