Add model index functions

StabilityMatrix.Avalonia/DesignData/MockModelIndexService.cs

@@ -1,5 +1,6 @@
 using System;
 using System.Collections.Generic;
+using System.Linq;
 using System.Threading.Tasks;
 using StabilityMatrix.Core.Models;
 using StabilityMatrix.Core.Models.Database;
@@ -25,9 +26,21 @@ public class MockModelIndexService : IModelIndexService
     }
 
     /// <inheritdoc />
-    public Task<IReadOnlyList<LocalModelFile>> GetModelsOfType(SharedFolderType type)
+    public Task<IEnumerable<LocalModelFile>> FindAsync(SharedFolderType type)
     {
-        return Task.FromResult<IReadOnlyList<LocalModelFile>>(new List<LocalModelFile>());
+        return Task.FromResult(Enumerable.Empty<LocalModelFile>());
+    }
+
+    /// <inheritdoc />
+    public Task<IEnumerable<LocalModelFile>> FindByHashAsync(string hashBlake3)
+    {
+        return Task.FromResult(Enumerable.Empty<LocalModelFile>());
+    }
+
+    /// <inheritdoc />
+    public Task<bool> RemoveModelAsync(LocalModelFile model)
+    {
+        return Task.FromResult(false);
     }
 
     /// <inheritdoc />

StabilityMatrix.Avalonia/Models/CommandItem.cs

@@ -0,0 +1,33 @@
+using System.Diagnostics.Contracts;
+using System.Runtime.CompilerServices;
+using System.Text.RegularExpressions;
+using System.Windows.Input;
+using StabilityMatrix.Core.Extensions;
+
+namespace StabilityMatrix.Avalonia.Models;
+
+public partial record CommandItem
+{
+    public ICommand Command { get; init; }
+
+    public string DisplayName { get; init; }
+
+    public CommandItem(ICommand command, [CallerArgumentExpression("command")] string? commandName = null)
+    {
+        Command = command;
+        DisplayName = commandName == null ? "" : ProcessName(commandName);
+    }
+
+    [Pure]
+    private static string ProcessName(string name)
+    {
+        name = name.StripEnd("Command");
+
+        name = SpaceTitleCaseRegex().Replace(name, "$1 $2");
+
+        return name;
+    }
+
+    [GeneratedRegex("([a-z])_?([A-Z])")]
+    private static partial Regex SpaceTitleCaseRegex();
+}

StabilityMatrix.Avalonia/Views/Settings/MainSettingsPage.axaml

@@ -20,6 +20,8 @@
     xmlns:update="clr-namespace:StabilityMatrix.Core.Models.Update;assembly=StabilityMatrix.Core"
     xmlns:vm="clr-namespace:StabilityMatrix.Avalonia.ViewModels"
     xmlns:vmSettings="clr-namespace:StabilityMatrix.Avalonia.ViewModels.Settings"
+    xmlns:input="clr-namespace:System.Windows.Input;assembly=System.ObjectModel"
+    xmlns:ct="clr-namespace:CommunityToolkit.Mvvm.Input;assembly=CommunityToolkit.Mvvm"
     d:DataContext="{x:Static mocks:DesignData.MainSettingsViewModel}"
     d:DesignHeight="700"
     d:DesignWidth="800"
@@ -566,6 +568,22 @@
                         </ui:SettingsExpanderItem.Footer>
                     </ui:SettingsExpanderItem>
 
+                    <ItemsRepeater
+                        Margin="4,0,4,4"
+                        ItemsSource="{Binding DebugCommands}">
+                        <ItemsRepeater.Layout>
+                            <UniformGridLayout />
+                        </ItemsRepeater.Layout>
+                        <ItemsRepeater.ItemTemplate>
+                            <DataTemplate DataType="models:CommandItem">
+                                <Button
+                                    Margin="4,8"
+                                    Command="{Binding Command}"
+                                    Content="{Binding DisplayName}" />
+                            </DataTemplate>
+                        </ItemsRepeater.ItemTemplate>
+                    </ItemsRepeater>
+                    
                 </ui:SettingsExpander>
             </Grid>
 

StabilityMatrix.Core/Models/Database/LocalModelFile.cs

@@ -6,13 +6,13 @@ namespace StabilityMatrix.Core.Models.Database;
 /// <summary>
 /// Represents a locally indexed model file.
 /// </summary>
-public class LocalModelFile
+public record LocalModelFile
 {
     /// <summary>
     /// Relative path to the file from the root model directory.
     /// </summary>
     [BsonId]
-    public required string RelativePath { get; set; }
+    public required string RelativePath { get; init; }
 
     /// <summary>
     /// Type of the model file.
@@ -53,20 +53,10 @@ public class LocalModelFile
     public string RelativePathFromSharedFolder =>
         Path.GetRelativePath(SharedFolderType.GetStringValue(), RelativePath);
 
-    public string GetFullPath(string rootModelDirectory)
-    {
-        return Path.Combine(rootModelDirectory, RelativePath);
-    }
-
-    public string? GetPreviewImageFullPath(string rootModelDirectory)
-    {
-        if (PreviewImageFullPath != null)
-            return PreviewImageFullPath;
-
-        return PreviewImageRelativePath == null
-            ? null
-            : Path.Combine(rootModelDirectory, PreviewImageRelativePath);
-    }
+    /// <summary>
+    /// Blake3 hash of the file.
+    /// </summary>
+    public string? HashBlake3 => ConnectedModelInfo?.Hashes.BLAKE3;
 
     [BsonIgnore]
     public string FullPathGlobal => GetFullPath(GlobalConfig.LibraryDir.JoinDir("Models"));
@@ -75,6 +65,10 @@ public class LocalModelFile
     public string? PreviewImageFullPathGlobal =>
         PreviewImageFullPath ?? GetPreviewImageFullPath(GlobalConfig.LibraryDir.JoinDir("Models"));
 
+    [BsonIgnore]
+    public Uri? PreviewImageUriGlobal =>
+        PreviewImageFullPathGlobal == null ? null : new Uri(PreviewImageFullPathGlobal);
+
     [BsonIgnore]
     public string DisplayModelName => ConnectedModelInfo?.ModelName ?? FileNameWithoutExtension;
 
@@ -84,7 +78,22 @@ public class LocalModelFile
     [BsonIgnore]
     public string DisplayModelFileName => FileName;
 
-    protected bool Equals(LocalModelFile other)
+    public string GetFullPath(string rootModelDirectory)
+    {
+        return Path.Combine(rootModelDirectory, RelativePath);
+    }
+
+    public string? GetPreviewImageFullPath(string rootModelDirectory)
+    {
+        if (PreviewImageFullPath != null)
+            return PreviewImageFullPath;
+
+        return PreviewImageRelativePath == null
+            ? null
+            : Path.Combine(rootModelDirectory, PreviewImageRelativePath);
+    }
+
+    /*protected bool Equals(LocalModelFile other)
     {
         return RelativePath == other.RelativePath;
     }
@@ -105,7 +114,7 @@ public class LocalModelFile
     public override int GetHashCode()
     {
         return RelativePath.GetHashCode();
-    }
+    }*/
 
     public static readonly HashSet<string> SupportedCheckpointExtensions =
     [

StabilityMatrix.Core/Services/IModelIndexService.cs

@@ -12,18 +12,28 @@ public interface IModelIndexService
     /// </summary>
     Task RefreshIndex();
 
+    /// <summary>
+    /// Starts a background task to refresh the local model file index.
+    /// </summary>
+    void BackgroundRefreshIndex();
+
     /// <summary>
     /// Get all models of the specified type from the existing (in-memory) index.
     /// </summary>
     IEnumerable<LocalModelFile> GetFromModelIndex(SharedFolderType types);
 
     /// <summary>
-    /// Get all models of the specified type from the existing index.
+    /// Find all models of the specified SharedFolderType.
     /// </summary>
-    Task<IReadOnlyList<LocalModelFile>> GetModelsOfType(SharedFolderType type);
+    Task<IEnumerable<LocalModelFile>> FindAsync(SharedFolderType type);
 
     /// <summary>
-    /// Starts a background task to refresh the local model file index.
+    /// Find all models with the specified Blake3 hash.
     /// </summary>
-    void BackgroundRefreshIndex();
+    Task<IEnumerable<LocalModelFile>> FindByHashAsync(string hashBlake3);
+
+    /// <summary>
+    /// Remove a model from the index.
+    /// </summary>
+    Task<bool> RemoveModelAsync(LocalModelFile model);
 }

StabilityMatrix.Core/Services/ModelIndexService.cs

@@ -1,6 +1,9 @@
-using System.Diagnostics;
+using System.Collections.Concurrent;
+using System.Collections.Immutable;
+using System.Diagnostics;
 using System.Text;
 using AsyncAwaitBestPractices;
+using AutoCtor;
 using Microsoft.Extensions.Logging;
 using StabilityMatrix.Core.Attributes;
 using StabilityMatrix.Core.Database;
@@ -13,45 +16,44 @@ using StabilityMatrix.Core.Models.FileInterfaces;
 namespace StabilityMatrix.Core.Services;
 
 [Singleton(typeof(IModelIndexService))]
-public class ModelIndexService : IModelIndexService
+[AutoConstruct]
+public partial class ModelIndexService : IModelIndexService
 {
     private readonly ILogger<ModelIndexService> logger;
-    private readonly ILiteDbContext liteDbContext;
     private readonly ISettingsManager settingsManager;
+    private readonly ILiteDbContext liteDbContext;
 
     public Dictionary<SharedFolderType, List<LocalModelFile>> ModelIndex { get; private set; } = new();
 
-    public ModelIndexService(
-        ILogger<ModelIndexService> logger,
-        ILiteDbContext liteDbContext,
-        ISettingsManager settingsManager
-    )
+    [AutoPostConstruct]
+    private void Initialize()
     {
-        this.logger = logger;
-        this.liteDbContext = liteDbContext;
-        this.settingsManager = settingsManager;
+        // Start background index when library dir is set
+        settingsManager.RegisterOnLibraryDirSet(_ => BackgroundRefreshIndex());
     }
 
-    /// <summary>
-    /// Deletes all entries in the local model file index.
-    /// </summary>
-    private async Task ClearIndex()
+    public IEnumerable<LocalModelFile> GetFromModelIndex(SharedFolderType types)
     {
-        await liteDbContext.LocalModelFiles.DeleteAllAsync().ConfigureAwait(false);
+        return ModelIndex.Where(kvp => (kvp.Key & types) != 0).SelectMany(kvp => kvp.Value);
     }
 
-    public IEnumerable<LocalModelFile> GetFromModelIndex(SharedFolderType types)
+    /// <inheritdoc />
+    public async Task<IEnumerable<LocalModelFile>> FindAsync(SharedFolderType type)
     {
-        return ModelIndex.Where(kvp => (kvp.Key & types) != 0).SelectMany(kvp => kvp.Value);
+        await liteDbContext.LocalModelFiles.EnsureIndexAsync(m => m.SharedFolderType).ConfigureAwait(false);
+
+        return await liteDbContext
+            .LocalModelFiles.FindAsync(m => m.SharedFolderType == type)
+            .ConfigureAwait(false);
     }
 
     /// <inheritdoc />
-    public async Task<IReadOnlyList<LocalModelFile>> GetModelsOfType(SharedFolderType type)
+    public async Task<IEnumerable<LocalModelFile>> FindByHashAsync(string hashBlake3)
     {
+        await liteDbContext.LocalModelFiles.EnsureIndexAsync(m => m.HashBlake3).ConfigureAwait(false);
+
         return await liteDbContext
-            .LocalModelFiles.Query()
-            .Where(m => m.SharedFolderType == type)
-            .ToArrayAsync()
+            .LocalModelFiles.FindAsync(m => m.HashBlake3 == hashBlake3)
             .ConfigureAwait(false);
     }
 
@@ -86,6 +88,7 @@ public class ModelIndexService : IModelIndexService
         var added = 0;
 
         var newIndex = new Dictionary<SharedFolderType, List<LocalModelFile>>();
+
         foreach (
             var file in modelsDir
                 .Info.EnumerateFiles("*.*", SearchOption.AllDirectories)
@@ -161,6 +164,7 @@ public class ModelIndexService : IModelIndexService
             // Add to index
             var list = newIndex.GetOrAdd(sharedFolderType);
             list.Add(localModel);
+
             added++;
         }
 
@@ -195,4 +199,24 @@ public class ModelIndexService : IModelIndexService
                 logger.LogError(ex, "Error in background model indexing");
             });
     }
+
+    /// <inheritdoc />
+    public async Task<bool> RemoveModelAsync(LocalModelFile model)
+    {
+        // Remove from database
+        if (await liteDbContext.LocalModelFiles.DeleteAsync(model.RelativePath).ConfigureAwait(false))
+        {
+            // Remove from index
+            if (ModelIndex.TryGetValue(model.SharedFolderType, out var list))
+            {
+                list.Remove(model);
+            }
+
+            EventManager.Instance.OnModelIndexChanged();
+
+            return true;
+        }
+
+        return false;
+    }
 }