diff --git a/StabilityMatrix.Avalonia/DesignData/DesignData.cs b/StabilityMatrix.Avalonia/DesignData/DesignData.cs index 914ce421..89b132a4 100644 --- a/StabilityMatrix.Avalonia/DesignData/DesignData.cs +++ b/StabilityMatrix.Avalonia/DesignData/DesignData.cs @@ -5,6 +5,7 @@ using System.Collections.ObjectModel; using System.Diagnostics.CodeAnalysis; using System.IO; using System.Net.Http; +using DynamicData.Binding; using Microsoft.Extensions.DependencyInjection; using StabilityMatrix.Avalonia.Models; using StabilityMatrix.Avalonia.Services; @@ -166,65 +167,72 @@ public static class DesignData InstallerViewModel.ReleaseNotes = "## Release Notes\nThis is a test release note."; // Checkpoints page - CheckpointsPageViewModel.CheckpointFolders = new ObservableCollection - { - new(settingsManager, downloadService, modelFinder, notificationService) + CheckpointsPageViewModel.CheckpointFolders = + new ObservableCollectionExtended { - Title = "StableDiffusion", - DirectoryPath = "Models/StableDiffusion", - CheckpointFiles = new AdvancedObservableList + new(settingsManager, downloadService, modelFinder, notificationService) { - new() + Title = "StableDiffusion", + DirectoryPath = "Models/StableDiffusion", + CheckpointFiles = new AdvancedObservableList { - FilePath = "~/Models/StableDiffusion/electricity-light.safetensors", - Title = "Auroral Background", - PreviewImagePath = - "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/" - + "78fd2a0a-42b6-42b0-9815-81cb11bb3d05/00009-2423234823.jpeg", - ConnectedModel = new ConnectedModelInfo + new() { - VersionName = "Lightning Auroral", - BaseModel = "SD 1.5", - ModelName = "Auroral Background", - ModelType = CivitModelType.Model, - FileMetadata = new CivitFileMetadata + FilePath = "~/Models/StableDiffusion/electricity-light.safetensors", + Title = "Auroral Background", + PreviewImagePath = + "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/" + + "78fd2a0a-42b6-42b0-9815-81cb11bb3d05/00009-2423234823.jpeg", + ConnectedModel = new ConnectedModelInfo { - Format = CivitModelFormat.SafeTensor, - Fp = CivitModelFpType.fp16, - Size = CivitModelSize.pruned, + VersionName = "Lightning Auroral", + BaseModel = "SD 1.5", + ModelName = "Auroral Background", + ModelType = CivitModelType.Model, + FileMetadata = new CivitFileMetadata + { + Format = CivitModelFormat.SafeTensor, + Fp = CivitModelFpType.fp16, + Size = CivitModelSize.pruned, + } } - } + }, + new() + { + FilePath = "~/Models/Lora/model.safetensors", + Title = "Some model" + }, }, - new() { FilePath = "~/Models/Lora/model.safetensors", Title = "Some model" }, }, - }, - new(settingsManager, downloadService, modelFinder, notificationService) - { - Title = "Lora", - DirectoryPath = "Packages/Lora", - SubFolders = new AdvancedObservableList() + new(settingsManager, downloadService, modelFinder, notificationService) { - new(settingsManager, downloadService, modelFinder, notificationService) + Title = "Lora", + DirectoryPath = "Packages/Lora", + SubFolders = new AdvancedObservableList() { - Title = "StableDiffusion", - DirectoryPath = "Packages/Lora/Subfolder", + new(settingsManager, downloadService, modelFinder, notificationService) + { + Title = "StableDiffusion", + DirectoryPath = "Packages/Lora/Subfolder", + }, + new(settingsManager, downloadService, modelFinder, notificationService) + { + Title = "Lora", + DirectoryPath = "Packages/StableDiffusion/Subfolder", + } }, - new(settingsManager, downloadService, modelFinder, notificationService) + CheckpointFiles = new AdvancedObservableList { - Title = "Lora", - DirectoryPath = "Packages/StableDiffusion/Subfolder", + new() { FilePath = "~/Models/Lora/lora_v2.pt", Title = "Best Lora v2", } } - }, - CheckpointFiles = new AdvancedObservableList - { - new() { FilePath = "~/Models/Lora/lora_v2.pt", Title = "Best Lora v2", } } - } - }; + }; foreach (var folder in CheckpointsPageViewModel.CheckpointFolders) { - folder.DisplayedCheckpointFiles = folder.CheckpointFiles; + folder.DisplayedCheckpointFiles = new AdvancedObservableList( + folder.CheckpointFiles + ); } CheckpointBrowserViewModel.ModelCards = diff --git a/StabilityMatrix.Avalonia/ViewModels/CheckpointManager/CheckpointFile.cs b/StabilityMatrix.Avalonia/ViewModels/CheckpointManager/CheckpointFile.cs index 0245ef12..239eff85 100644 --- a/StabilityMatrix.Avalonia/ViewModels/CheckpointManager/CheckpointFile.cs +++ b/StabilityMatrix.Avalonia/ViewModels/CheckpointManager/CheckpointFile.cs @@ -56,7 +56,7 @@ public partial class CheckpointFile : ViewModelBase public ObservableCollection Badges { get; set; } = new(); - private static readonly string[] SupportedCheckpointExtensions = + public static readonly string[] SupportedCheckpointExtensions = { ".safetensors", ".pt", diff --git a/StabilityMatrix.Avalonia/ViewModels/CheckpointsPageViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/CheckpointsPageViewModel.cs index 954a4c4b..bb04fc7a 100644 --- a/StabilityMatrix.Avalonia/ViewModels/CheckpointsPageViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/CheckpointsPageViewModel.cs @@ -8,6 +8,7 @@ using Avalonia.Controls; using Avalonia.Threading; using CommunityToolkit.Mvvm.ComponentModel; using CommunityToolkit.Mvvm.Input; +using DynamicData; using FluentAvalonia.UI.Controls; using NLog; using StabilityMatrix.Avalonia.Services; @@ -32,6 +33,8 @@ public partial class CheckpointsPageViewModel : PageViewModelBase private readonly ModelFinder modelFinder; private readonly IDownloadService downloadService; private readonly INotificationService notificationService; + private readonly SourceCache checkpointFoldersCache = + new(x => x.Title); public override string Title => "Checkpoints"; @@ -65,11 +68,10 @@ public partial class CheckpointsPageViewModel : PageViewModelBase } [ObservableProperty] - private ObservableCollection checkpointFolders = new(); + private ObservableCollection checkpointFolders = new(); [ObservableProperty] - private ObservableCollection displayedCheckpointFolders = - new(); + private ObservableCollection displayedCheckpointFolders = new(); public CheckpointsPageViewModel( ISharedFolders sharedFolders, @@ -97,16 +99,13 @@ public partial class CheckpointsPageViewModel : PageViewModelBase // Refresh search filter OnSearchFilterChanged(string.Empty); - Logger.Info( - $"Loaded {DisplayedCheckpointFolders.Count} checkpoint folders in {sw.ElapsedMilliseconds}ms" - ); - if (Design.IsDesignMode) return; IsLoading = CheckpointFolders.Count == 0; IsIndexing = CheckpointFolders.Count > 0; - await IndexFolders(); + GetStuff(); + //await IndexFolders(); IsLoading = false; IsIndexing = false; @@ -119,12 +118,15 @@ public partial class CheckpointsPageViewModel : PageViewModelBase var sw = Stopwatch.StartNew(); if (string.IsNullOrWhiteSpace(SearchFilter)) { - DisplayedCheckpointFolders = new ObservableCollection( - CheckpointFolders.Select(x => - { - x.SearchFilter = SearchFilter; - return x; - }) + checkpointFoldersCache.Edit( + s => + s.Load( + CheckpointFolders.Select(x => + { + x.SearchFilter = SearchFilter; + return x; + }) + ) ); sw.Stop(); Logger.Info($"OnSearchFilterChanged in {sw.ElapsedMilliseconds}ms"); @@ -138,10 +140,11 @@ public partial class CheckpointsPageViewModel : PageViewModelBase { folder.SearchFilter = SearchFilter; } + + checkpointFoldersCache.Edit(s => s.Load(filteredFolders)); + sw.Stop(); Logger.Info($"ContainsSearchFilter in {sw.ElapsedMilliseconds}ms"); - - DisplayedCheckpointFolders = new ObservableCollection(filteredFolders); } partial void OnShowConnectedModelImagesChanged(bool value) @@ -233,6 +236,99 @@ public partial class CheckpointsPageViewModel : PageViewModelBase } } + private void GetStuff() + { + CheckpointFolders.Clear(); + var allFiles = Directory.EnumerateFiles( + settingsManager.ModelsDirectory, + "*.*", + SearchOption.AllDirectories + ); + foreach (var file in allFiles) + { + var extension = Path.GetExtension(file); + if (!CheckpointFile.SupportedCheckpointExtensions.Contains(extension)) + continue; + + var folder = + Path.GetDirectoryName(file) + ?.Replace( + $"{settingsManager.ModelsDirectory}{Path.DirectorySeparatorChar}", + string.Empty + ) ?? string.Empty; + + if (string.IsNullOrWhiteSpace(folder)) + continue; + + var isRootFolder = !folder.Contains(Path.DirectorySeparatorChar); + var rootFolderName = isRootFolder + ? folder + : folder.Split(Path.DirectorySeparatorChar).First(); + + var rootCheckpointFolder = CheckpointFolders.FirstOrDefault( + x => x.Title == rootFolderName + ); + if (rootCheckpointFolder == null) + { + rootCheckpointFolder = new CheckpointFolder( + settingsManager, + downloadService, + modelFinder, + notificationService + ) + { + Title = rootFolderName, + DirectoryPath = Path.Combine(settingsManager.ModelsDirectory, rootFolderName), + IsExpanded = isRootFolder, // Top level folders expanded by default + }; + CheckpointFolders.Add(rootCheckpointFolder); + } + + if (isRootFolder) + { + rootCheckpointFolder.CheckpointFiles.Add( + new CheckpointFile { Title = Path.GetFileName(file), FilePath = file } + ); + continue; + } + + // recursively add subfolders + var subFolderNames = folder.Split(Path.DirectorySeparatorChar).Skip(1); + + foreach (var subFolderName in subFolderNames) + { + var subFolder = rootCheckpointFolder.SubFolders.FirstOrDefault( + x => x.Title == subFolderName + ); + if (subFolder == null) + { + subFolder = new CheckpointFolder( + settingsManager, + downloadService, + modelFinder, + notificationService + ) + { + Title = subFolderName, + DirectoryPath = Path.Combine( + rootCheckpointFolder.DirectoryPath, + subFolderName + ), + ParentFolder = rootCheckpointFolder + }; + rootCheckpointFolder.SubFolders.Add(subFolder); + } + rootCheckpointFolder = subFolder; + } + + rootCheckpointFolder.CheckpointFiles.Add( + new CheckpointFile { Title = Path.GetFileName(file), FilePath = file } + ); + } + + DisplayedCheckpointFolders = CheckpointFolders; + } + [RelayCommand] private async Task OpenModelsFolder() { diff --git a/StabilityMatrix.Avalonia/Views/CheckpointsPage.axaml b/StabilityMatrix.Avalonia/Views/CheckpointsPage.axaml index e1f8e17f..23f05e7b 100644 --- a/StabilityMatrix.Avalonia/Views/CheckpointsPage.axaml +++ b/StabilityMatrix.Avalonia/Views/CheckpointsPage.axaml @@ -177,22 +177,18 @@ TextTrimming="CharacterEllipsis" ToolTip.Tip="{Binding FileName}" VerticalAlignment="Bottom" /> - - - - - - - + + + + +