WIP new page for connected checkpoint updates/management

StabilityMatrix.Avalonia/App.axaml.cs

@@ -204,6 +204,7 @@ public sealed class App : Application
             .AddSingleton<SettingsViewModel>()
             .AddSingleton<CheckpointBrowserViewModel>()
             .AddSingleton<CheckpointsPageViewModel>()
+            .AddSingleton<NewCheckpointsPageViewModel>()
             .AddSingleton<LaunchPageViewModel>()
             .AddSingleton<ProgressManagerViewModel>();
         
@@ -217,6 +218,7 @@ public sealed class App : Application
                     provider.GetRequiredService<LaunchPageViewModel>(),
                     provider.GetRequiredService<PackageManagerViewModel>(),
                     provider.GetRequiredService<CheckpointsPageViewModel>(),
+                    provider.GetRequiredService<NewCheckpointsPageViewModel>(),
                     provider.GetRequiredService<CheckpointBrowserViewModel>(),
                 },
                 FooterPages =
@@ -285,6 +287,7 @@ public sealed class App : Application
         services.AddSingleton<SettingsPage>();
         services.AddSingleton<CheckpointBrowserPage>();
         services.AddSingleton<ProgressManagerPage>();
+        services.AddSingleton<NewCheckpointsPage>();
         
         // Dialogs
         services.AddTransient<SelectDataDirectoryDialog>();

StabilityMatrix.Avalonia/DesignData/DesignData.cs

@@ -223,6 +223,35 @@ public static class DesignData
                 })
             };
 
+        NewCheckpointsPageViewModel.AllCheckpoints = new ObservableCollection<CheckpointFile>
+        {
+            new()
+            {
+                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
+                {
+                    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"
+            }
+        };
+
         ProgressManagerViewModel.ProgressItems = new ObservableCollection<ProgressItemViewModel>
         {
             new(new ProgressItem(Guid.NewGuid(), "Test File.exe",
@@ -273,6 +302,9 @@ public static class DesignData
     public static CheckpointsPageViewModel CheckpointsPageViewModel =>
         Services.GetRequiredService<CheckpointsPageViewModel>();
 
+    public static NewCheckpointsPageViewModel NewCheckpointsPageViewModel =>
+        Services.GetRequiredService<NewCheckpointsPageViewModel>();
+
     public static SettingsViewModel SettingsViewModel =>
         Services.GetRequiredService<SettingsViewModel>();
 

StabilityMatrix.Avalonia/ViewModels/CheckpointManager/CheckpointFile.cs

@@ -14,6 +14,7 @@ using StabilityMatrix.Avalonia.ViewModels.Base;
 using StabilityMatrix.Core.Extensions;
 using StabilityMatrix.Core.Helper;
 using StabilityMatrix.Core.Models;
+using StabilityMatrix.Core.Models.Api;
 using StabilityMatrix.Core.Models.Progress;
 using StabilityMatrix.Core.Processes;
 
@@ -46,6 +47,7 @@ public partial class CheckpointFile : ViewModelBase
     public bool IsConnectedModel => ConnectedModel != null;
 
     [ObservableProperty] private bool isLoading;
+    [ObservableProperty] private CivitModelType modelType;
     
     public string FileName => Path.GetFileName((string?) FilePath);
 
@@ -223,6 +225,39 @@ public partial class CheckpointFile : ViewModelBase
             yield return checkpointFile;
         }
     }
+    
+    public static IEnumerable<CheckpointFile> GetAllCheckpointFiles(string modelsDirectory)
+    {
+        foreach (var file in Directory.EnumerateFiles(modelsDirectory, "*.*", SearchOption.AllDirectories))
+        {
+            if (!SupportedCheckpointExtensions.Any(ext => file.Contains(ext)))
+                continue;
+            
+            var checkpointFile = new CheckpointFile
+            {
+                Title = Path.GetFileNameWithoutExtension(file),
+                FilePath = file,
+            };
+
+            var jsonPath = Path.Combine(Path.GetDirectoryName(file),
+                Path.GetFileNameWithoutExtension(file) + ".cm-info.json");
+            
+            if (File.Exists(jsonPath)) 
+            {
+                var json = File.ReadAllText(jsonPath);
+                var connectedModelInfo = ConnectedModelInfo.FromJson(json);
+                checkpointFile.ConnectedModel = connectedModelInfo;
+                checkpointFile.ModelType = GetCivitModelType(file);
+            }
+
+            checkpointFile.PreviewImagePath = SupportedImageExtensions
+                .Select(ext => Path.Combine(Path.GetDirectoryName(file),
+                    $"{Path.GetFileNameWithoutExtension(file)}.preview{ext}")).Where(File.Exists)
+                .FirstOrDefault();
+
+            yield return checkpointFile;
+        }
+    }
 
     /// <summary>
     /// Index with progress reporting.
@@ -238,4 +273,39 @@ public partial class CheckpointFile : ViewModelBase
             yield return checkpointFile;
         }
     }
+
+    private static CivitModelType GetCivitModelType(string filePath)
+    {
+        if (filePath.Contains(SharedFolderType.StableDiffusion.ToString()))
+        {
+            return CivitModelType.Checkpoint;
+        }
+        
+        if (filePath.Contains(SharedFolderType.ControlNet.ToString()))
+        {
+            return CivitModelType.Controlnet;
+        }
+
+        if (filePath.Contains(SharedFolderType.Lora.ToString()))
+        {
+            return CivitModelType.LORA;
+        }
+
+        if (filePath.Contains(SharedFolderType.TextualInversion.ToString()))
+        {
+            return CivitModelType.TextualInversion;
+        }
+
+        if (filePath.Contains(SharedFolderType.Hypernetwork.ToString()))
+        {
+            return CivitModelType.Hypernetwork;
+        }
+
+        if (filePath.Contains(SharedFolderType.LyCORIS.ToString()))
+        {
+            return CivitModelType.LoCon;
+        }
+
+        return CivitModelType.Unknown;
+    }
 }

StabilityMatrix.Avalonia/ViewModels/CheckpointsPageViewModel.cs

@@ -70,7 +70,7 @@ public partial class CheckpointsPageViewModel : PageViewModelBase
         this.downloadService = downloadService;
         this.modelFinder = modelFinder;
     }
-
+    
     public override async Task OnLoadedAsync()
     {
         DisplayedCheckpointFolders = CheckpointFolders;
@@ -147,7 +147,7 @@ public partial class CheckpointsPageViewModel : PageViewModelBase
         var indexTasks = folders.Select(async f =>
         {
             var checkpointFolder =
-                new CheckpointManager.CheckpointFolder(settingsManager, downloadService, modelFinder)
+                new CheckpointFolder(settingsManager, downloadService, modelFinder)
                 {
                     Title = Path.GetFileName(f),
                     DirectoryPath = f,

StabilityMatrix.Avalonia/ViewModels/NewCheckpointsPageViewModel.cs

@@ -0,0 +1,187 @@
+using System;
+using System.Collections.Generic;
+using System.Collections.Immutable;
+using System.Collections.ObjectModel;
+using System.IO;
+using System.Linq;
+using System.Threading.Tasks;
+using AsyncAwaitBestPractices;
+using Avalonia.Controls;
+using Avalonia.Controls.Notifications;
+using AvaloniaEdit.Utils;
+using CommunityToolkit.Mvvm.ComponentModel;
+using FluentAvalonia.UI.Controls;
+using StabilityMatrix.Avalonia.Controls;
+using StabilityMatrix.Avalonia.Services;
+using StabilityMatrix.Avalonia.ViewModels.Base;
+using StabilityMatrix.Avalonia.ViewModels.CheckpointManager;
+using StabilityMatrix.Avalonia.ViewModels.Dialogs;
+using StabilityMatrix.Avalonia.Views;
+using StabilityMatrix.Avalonia.Views.Dialogs;
+using StabilityMatrix.Core.Api;
+using StabilityMatrix.Core.Attributes;
+using StabilityMatrix.Core.Database;
+using StabilityMatrix.Core.Extensions;
+using StabilityMatrix.Core.Helper;
+using StabilityMatrix.Core.Models;
+using StabilityMatrix.Core.Models.Api;
+using StabilityMatrix.Core.Services;
+using Symbol = FluentIcons.Common.Symbol;
+using SymbolIconSource = FluentIcons.FluentAvalonia.SymbolIconSource;
+
+namespace StabilityMatrix.Avalonia.ViewModels;
+
+[View(typeof(NewCheckpointsPage))]
+public partial class NewCheckpointsPageViewModel : PageViewModelBase
+{
+    private readonly ISettingsManager settingsManager;
+    private readonly ILiteDbContext liteDbContext;
+    private readonly ICivitApi civitApi;
+    private readonly ServiceManager<ViewModelBase> dialogFactory;
+    private readonly INotificationService notificationService;
+    public override string Title => "Checkpoint Manager";
+    public override IconSource IconSource => new SymbolIconSource
+        {Symbol = Symbol.Cellular5g, IsFilled = true};
+
+    public NewCheckpointsPageViewModel(ISettingsManager settingsManager, ILiteDbContext liteDbContext,
+        ICivitApi civitApi, ServiceManager<ViewModelBase> dialogFactory, INotificationService notificationService)
+    {
+        this.settingsManager = settingsManager;
+        this.liteDbContext = liteDbContext;
+        this.civitApi = civitApi;
+        this.dialogFactory = dialogFactory;
+        this.notificationService = notificationService;
+    }
+
+    [ObservableProperty]
+    [NotifyPropertyChangedFor(nameof(ConnectedCheckpoints))]
+    [NotifyPropertyChangedFor(nameof(NonConnectedCheckpoints))]
+    private ObservableCollection<CheckpointFile> allCheckpoints = new();
+    
+    [ObservableProperty]
+    private ObservableCollection<CivitModel> civitModels = new();
+
+    public ObservableCollection<CheckpointFile> ConnectedCheckpoints => new(
+        AllCheckpoints.Where(x => x.IsConnectedModel)
+            .OrderBy(x => x.ConnectedModel!.ModelName)
+            .ThenBy(x => x.ModelType)
+            .GroupBy(x => x.ConnectedModel!.ModelId)
+            .Select(x => x.First()));
+
+    public ObservableCollection<CheckpointFile> NonConnectedCheckpoints => new(
+        AllCheckpoints.Where(x => !x.IsConnectedModel).OrderBy(x => x.ModelType));
+
+    public override async Task OnLoadedAsync()
+    {
+        if (Design.IsDesignMode) return;
+
+        var files = CheckpointFile.GetAllCheckpointFiles(settingsManager.ModelsDirectory);
+        AllCheckpoints = new ObservableCollection<CheckpointFile>(files);
+
+        var connectedModelIds = ConnectedCheckpoints.Select(x => x.ConnectedModel.ModelId);
+        var modelRequest = new CivitModelsRequest
+        {
+            CommaSeparatedModelIds = string.Join(',', connectedModelIds)
+        };
+        
+        // See if query is cached
+        var cachedQuery = await liteDbContext.CivitModelQueryCache
+            .IncludeAll()
+            .FindByIdAsync(ObjectHash.GetMd5Guid(modelRequest));
+        
+        // If cached, update model cards
+        if (cachedQuery is not null)
+        {
+            CivitModels = new ObservableCollection<CivitModel>(cachedQuery.Items);
+            
+            // Start remote query (background mode)
+            // Skip when last query was less than 2 min ago
+            var timeSinceCache = DateTimeOffset.UtcNow - cachedQuery.InsertedAt;
+            if (timeSinceCache?.TotalMinutes >= 2)
+            {
+                CivitQuery(modelRequest).SafeFireAndForget();
+            }
+        }
+        else
+        {
+            await CivitQuery(modelRequest);
+        }
+    }
+    
+    public async Task ShowVersionDialog(int modelId)
+    {
+        var model = CivitModels.FirstOrDefault(m => m.Id == modelId);
+        if (model == null)
+        {
+            notificationService.Show(new Notification("Model has no versions available",
+                "This model has no versions available for download", NotificationType.Warning));
+            return;
+        }
+        var versions = model.ModelVersions;
+        if (versions is null || versions.Count == 0)
+        {
+            notificationService.Show(new Notification("Model has no versions available",
+                "This model has no versions available for download", NotificationType.Warning));
+            return;
+        }
+        
+        var dialog = new BetterContentDialog
+        {
+            Title = model.Name,
+            IsPrimaryButtonEnabled = false,
+            IsSecondaryButtonEnabled = false,
+            IsFooterVisible = false,
+            MaxDialogWidth = 750,
+        };
+        
+        var viewModel = dialogFactory.Get<SelectModelVersionViewModel>();
+        viewModel.Dialog = dialog;
+        viewModel.Versions = versions.Select(version =>
+                new ModelVersionViewModel(
+                    settingsManager.Settings.InstalledModelHashes ?? new HashSet<string>(), version))
+            .ToImmutableArray();
+        viewModel.SelectedVersionViewModel = viewModel.Versions[0];
+        
+        dialog.Content = new SelectModelVersionDialog
+        {
+            DataContext = viewModel
+        };
+
+        var result = await dialog.ShowAsync();
+
+        if (result != ContentDialogResult.Primary)
+        {
+            return;
+        }
+
+        var selectedVersion = viewModel?.SelectedVersionViewModel?.ModelVersion;
+        var selectedFile = viewModel?.SelectedFile?.CivitFile;
+    }
+
+    private async Task CivitQuery(CivitModelsRequest request)
+    {
+        var modelResponse = await civitApi.GetModels(request);
+        var models = modelResponse.Items;
+        // Filter out unknown model types and archived/taken-down models
+        models = models.Where(m => m.Type.ConvertTo<SharedFolderType>() > 0)
+            .Where(m => m.Mode == null).ToList();
+
+        // Database update calls will invoke `OnModelsUpdated`
+        // Add to database
+        await liteDbContext.UpsertCivitModelAsync(models);
+        // Add as cache entry
+        var cacheNew = await liteDbContext.UpsertCivitModelQueryCacheEntryAsync(new CivitModelQueryCacheEntry
+        {
+            Id = ObjectHash.GetMd5Guid(request),
+            InsertedAt = DateTimeOffset.UtcNow,
+            Request = request,
+            Items = models,
+            Metadata = modelResponse.Metadata
+        });
+
+        if (cacheNew)
+        {
+            CivitModels = new ObservableCollection<CivitModel>(models);
+        }
+    }
+}

StabilityMatrix.Avalonia/Views/CheckpointBrowserPage.axaml

@@ -40,7 +40,7 @@
                             Margin="0,8,0,8"
                             Height="300"
                             StretchDirection="Both"
-                            CornerRadius="4"
+                            CornerRadius="8"
                             VerticalContentAlignment="Top"
                             HorizontalContentAlignment="Center"
                             Source="{Binding CardImage}"

StabilityMatrix.Avalonia/Views/NewCheckpointsPage.axaml

@@ -0,0 +1,105 @@
+<controls:UserControlBase xmlns="https://github.com/avaloniaui"
+                          xmlns:x="http://schemas.microsoft.com/winfx/2006/xaml"
+                          xmlns:d="http://schemas.microsoft.com/expression/blend/2008"
+                          xmlns:mc="http://schemas.openxmlformats.org/markup-compatibility/2006"
+                          xmlns:viewModels="clr-namespace:StabilityMatrix.Avalonia.ViewModels"
+                          xmlns:mocks="clr-namespace:StabilityMatrix.Avalonia.DesignData"
+                          xmlns:checkpointManager="clr-namespace:StabilityMatrix.Avalonia.ViewModels.CheckpointManager"
+                          xmlns:controls="clr-namespace:StabilityMatrix.Avalonia.Controls"
+                          mc:Ignorable="d" d:DesignWidth="800" d:DesignHeight="450"
+                          x:Class="StabilityMatrix.Avalonia.Views.NewCheckpointsPage"
+                          d:DataContext="{x:Static mocks:DesignData.NewCheckpointsPageViewModel}"
+                          x:CompileBindings="True"
+                          x:Name="ControlBase"
+                          x:DataType="viewModels:NewCheckpointsPageViewModel">
+    <ScrollViewer>
+        <Grid RowDefinitions="Auto, Auto" Margin="8">
+            <TextBlock Text="Installed Checkpoints"
+                       FontSize="24"
+                       Margin="8"/>
+            
+            <ItemsRepeater Grid.Row="1" ItemsSource="{Binding ConnectedCheckpoints}">
+                <ItemsRepeater.Layout>
+                    <UniformGridLayout />
+                </ItemsRepeater.Layout>
+                <ItemsRepeater.ItemTemplate>
+                    <DataTemplate DataType="{x:Type checkpointManager:CheckpointFile}">
+                        <controls:Card
+                            Margin="8"
+                            MaxHeight="450"
+                            Width="300"
+                            CornerRadius="8">
+                            <Grid RowDefinitions="Auto, Auto, Auto, Auto">
+                                <TextBlock Grid.Row="0"
+                                           Margin="0,0,0,8"
+                                           Text="{Binding ConnectedModel.ModelName}" />
+
+                                <controls:BetterAdvancedImage
+                                    Grid.Row="1"
+                                    Height="250"
+                                    Stretch="UniformToFill"
+                                    CornerRadius="8"
+                                    Source="{Binding PreviewImagePath}" />
+
+                                <Grid Grid.Row="1"
+                                      Margin="8,8,0,0"
+                                      ColumnDefinitions="Auto, Auto">
+                                    <controls:Card
+                                        Grid.Column="0"
+                                        Classes="info"
+                                        Height="24"
+                                        HorizontalAlignment="Left"
+                                        Padding="4"
+                                        Margin="0,0,4,0"
+                                        VerticalAlignment="Top">
+
+                                        <TextBlock
+                                            FontSize="11"
+                                            FontWeight="Medium"
+                                            HorizontalAlignment="Center"
+                                            Text="{Binding ModelType}"
+                                            VerticalAlignment="Center" />
+                                    </controls:Card>
+                                    <controls:Card
+                                        Grid.Column="1"
+                                        Classes="info"
+                                        Height="24"
+                                        Margin="4,0"
+                                        HorizontalAlignment="Left"
+                                        Padding="4"
+                                        VerticalAlignment="Top">
+
+                                        <TextBlock
+                                            FontSize="11"
+                                            FontWeight="Medium"
+                                            HorizontalAlignment="Center"
+                                            Text="{Binding ConnectedModel.BaseModel}"
+                                            VerticalAlignment="Center" />
+                                    </controls:Card>
+                                </Grid>
+
+                                <Grid Grid.Row="2"
+                                      Margin="0,16,0,0"
+                                      ColumnDefinitions="*, *">
+                                    <Button Grid.Column="0"
+                                            Content="Update"
+                                            Classes="accent"
+                                            Margin="0,0,4,0"
+                                            HorizontalAlignment="Stretch" />
+                                    <Button Grid.Column="1"
+                                            Margin="4,0,0,0"
+                                            Content="All Versions"
+                                            Classes="accent"
+                                            Command="{Binding $parent[ItemsRepeater].((viewModels:NewCheckpointsPageViewModel)DataContext).ShowVersionDialog}"
+                                            CommandParameter="{Binding ConnectedModel.ModelId}"
+                                            HorizontalAlignment="Stretch" />
+                                </Grid>
+
+                            </Grid>
+                        </controls:Card>
+                    </DataTemplate>
+                </ItemsRepeater.ItemTemplate>
+            </ItemsRepeater>
+        </Grid>
+    </ScrollViewer>
+</controls:UserControlBase>

StabilityMatrix.Avalonia/Views/NewCheckpointsPage.axaml.cs

@@ -0,0 +1,11 @@
+using StabilityMatrix.Avalonia.Controls;
+
+namespace StabilityMatrix.Avalonia.Views;
+
+public partial class NewCheckpointsPage : UserControlBase
+{
+    public NewCheckpointsPage()
+    {
+        InitializeComponent();
+    }
+}

StabilityMatrix.Core/Models/Api/CivitModelsRequest.cs

@@ -117,4 +117,7 @@ public class CivitModelsRequest
     /// </summary>
     [AliasAs("baseModels")]
     public string? BaseModel { get; set; }
+    
+    [AliasAs("ids")]
+    public string CommaSeparatedModelIds { get; set; }
 }