Merge pull request LykosAI#947 from ionite34/backport/main/pr-946

[dev to main] backport: Fix & redesign checkpoint manager base model filtering (946)

Author: mohnjiles

CHANGELOG.md

@@ -6,8 +6,11 @@ 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.13.1
+### Changed
+- Redesigned the Checkpoint Manager Filter flyout to include more options and improve the layout
 ### Fixed
 - Fixed [#1078](https://github.com/LykosAI/StabilityMatrix/issues/1078) - "Call from invalid thread" error after one-click install finishes
+- Fixed [#1080](https://github.com/LykosAI/StabilityMatrix/issues/1080) - Some models not displayed in Checkpoint Manager
 - Fixed Inference image selector card buttons taking up the whole height of the card
 
 ## v2.13.0

StabilityMatrix.Avalonia/Services/CivitBaseModelTypeService.cs

@@ -0,0 +1,92 @@
+using System.Text.Json.Nodes;
+using Injectio.Attributes;
+using Microsoft.Extensions.Logging;
+using StabilityMatrix.Core.Api;
+using StabilityMatrix.Core.Database;
+using StabilityMatrix.Core.Models.Api;
+using StabilityMatrix.Core.Models.Database;
+
+namespace StabilityMatrix.Avalonia.Services;
+
+[RegisterSingleton<ICivitBaseModelTypeService, CivitBaseModelTypeService>]
+public class CivitBaseModelTypeService(
+    ILogger<CivitBaseModelTypeService> logger,
+    ICivitApi civitApi,
+    ILiteDbContext dbContext
+) : ICivitBaseModelTypeService
+{
+    private const string CacheId = "BaseModelTypes";
+    private static readonly TimeSpan CacheExpiration = TimeSpan.FromHours(24);
+
+    /// <summary>
+    /// Gets the list of base model types, using cache if available and not expired
+    /// </summary>
+    public async Task<List<string>> GetBaseModelTypes(bool forceRefresh = false, bool includeAllOption = true)
+    {
+        List<string> civitBaseModels = [];
+
+        if (!forceRefresh)
+        {
+            var cached = await dbContext.GetCivitBaseModelTypeCacheEntry(CacheId);
+            if (cached != null && DateTimeOffset.UtcNow.Subtract(cached.CreatedAt) < CacheExpiration)
+            {
+                civitBaseModels = cached.ModelTypes;
+            }
+        }
+
+        try
+        {
+            if (civitBaseModels.Count <= 0)
+            {
+                var baseModelsResponse = await civitApi.GetBaseModelList();
+                var jsonContent = await baseModelsResponse.Content.ReadAsStringAsync();
+                var baseModels = JsonNode.Parse(jsonContent);
+
+                var jArray =
+                    baseModels?["error"]?["issues"]?[0]?["unionErrors"]?[0]?["issues"]?[0]?["options"]
+                    as JsonArray;
+
+                civitBaseModels = jArray?.GetValues<string>().ToList() ?? [];
+
+                // Cache the results
+                var cacheEntry = new CivitBaseModelTypeCacheEntry
+                {
+                    Id = CacheId,
+                    ModelTypes = civitBaseModels,
+                    CreatedAt = DateTimeOffset.UtcNow
+                };
+
+                await dbContext.UpsertCivitBaseModelTypeCacheEntry(cacheEntry);
+            }
+
+            if (includeAllOption)
+            {
+                civitBaseModels.Insert(0, CivitBaseModelType.All.ToString());
+            }
+
+            // Filter and sort results
+            var filteredResults = civitBaseModels
+                .Where(s => !s.Equals("odor", StringComparison.OrdinalIgnoreCase))
+                .OrderBy(s => s)
+                .ToList();
+
+            return filteredResults;
+        }
+        catch (Exception e)
+        {
+            logger.LogError(e, "Failed to get base model list");
+
+            // Return cached results if available, even if expired
+            var expiredCache = await dbContext.GetCivitBaseModelTypeCacheEntry(CacheId);
+            return expiredCache?.ModelTypes ?? [];
+        }
+    }
+
+    /// <summary>
+    /// Clears the cached base model types
+    /// </summary>
+    public void ClearCache()
+    {
+        dbContext.CivitBaseModelTypeCache.DeleteAllAsync();
+    }
+}

StabilityMatrix.Avalonia/Services/ICivitBaseModelTypeService.cs

@@ -0,0 +1,7 @@
+namespace StabilityMatrix.Avalonia.Services;
+
+public interface ICivitBaseModelTypeService
+{
+    Task<List<string>> GetBaseModelTypes(bool forceRefresh = false, bool includeAllOption = true);
+    void ClearCache();
+}

StabilityMatrix.Avalonia/ViewModels/Base/PageViewModelBase.cs

@@ -5,7 +5,7 @@ namespace StabilityMatrix.Avalonia.ViewModels.Base;
 /// <summary>
 /// An abstract class for enabling page navigation.
 /// </summary>
-public abstract class PageViewModelBase : ViewModelBase
+public abstract class PageViewModelBase : DisposableViewModelBase
 {
     /// <summary>
     /// Gets if the user can navigate to the next page
@@ -16,7 +16,7 @@ public abstract class PageViewModelBase : ViewModelBase
     /// Gets if the user can navigate to the previous page
     /// </summary>
     public virtual bool CanNavigatePrevious { get; protected set; }
-    
+
     public abstract string Title { get; }
     public abstract IconSource IconSource { get; }
 }

StabilityMatrix.Avalonia/ViewModels/CheckpointBrowser/CivitAiBrowserViewModel.cs

@@ -45,6 +45,7 @@ public sealed partial class CivitAiBrowserViewModel : TabViewModelBase, IInfinit
     private readonly ISettingsManager settingsManager;
     private readonly ILiteDbContext liteDbContext;
     private readonly INotificationService notificationService;
+    private readonly ICivitBaseModelTypeService baseModelTypeService;
     private bool dontSearch = false;
 
     private readonly SourceCache<OrderedValue<CivitModel>, int> modelCache = new(static ov => ov.Value.Id);
@@ -126,13 +127,15 @@ public CivitAiBrowserViewModel(
         ISettingsManager settingsManager,
         ServiceManager<ViewModelBase> dialogFactory,
         ILiteDbContext liteDbContext,
-        INotificationService notificationService
+        INotificationService notificationService,
+        ICivitBaseModelTypeService baseModelTypeService
     )
     {
         this.civitApi = civitApi;
         this.settingsManager = settingsManager;
         this.liteDbContext = liteDbContext;
         this.notificationService = notificationService;
+        this.baseModelTypeService = baseModelTypeService;
 
         EventManager.Instance.NavigateAndFindCivitModelRequested += OnNavigateAndFindCivitModelRequested;
 
@@ -727,31 +730,7 @@ private void UpdateResultsText()
     [Localizable(false)]
     private async Task<List<string>> GetBaseModelList()
     {
-        try
-        {
-            var baseModelsResponse = await civitApi.GetBaseModelList();
-            var jsonContent = await baseModelsResponse.Content.ReadAsStringAsync();
-            var baseModels = JsonNode.Parse(jsonContent);
-
-            var jArray =
-                baseModels?["error"]?["issues"]?[0]?["unionErrors"]?[0]?["issues"]?[0]?["options"]
-                as JsonArray;
-            var civitBaseModels = jArray?.GetValues<string>().ToList() ?? [];
-
-            civitBaseModels.Insert(0, CivitBaseModelType.All.ToString());
-
-            var filteredResults = civitBaseModels
-                .Where(s => s.Equals("odor", StringComparison.OrdinalIgnoreCase) == false)
-                .OrderBy(s => s)
-                .ToList();
-
-            return filteredResults;
-        }
-        catch (Exception e)
-        {
-            Logger.Error(e, "Failed to get base model list");
-            return [];
-        }
+        return await baseModelTypeService.GetBaseModelTypes();
     }
 
     public override string Header => Resources.Label_CivitAi;