Add faster memory based image indexing

StabilityMatrix.Avalonia/Controls/ImageFolderCard.axaml

@@ -9,6 +9,7 @@
     xmlns:input="using:FluentAvalonia.UI.Input"
     xmlns:vmInference="clr-namespace:StabilityMatrix.Avalonia.ViewModels.Inference"
     xmlns:dbModels="clr-namespace:StabilityMatrix.Core.Models.Database;assembly=StabilityMatrix.Core"
+    xmlns:animations="clr-namespace:StabilityMatrix.Avalonia.Animations"
     x:DataType="vmInference:ImageFolderCardViewModel">
 
     <Design.PreviewWith>
@@ -38,6 +39,12 @@
                     HorizontalContentAlignment="{TemplateBinding HorizontalAlignment}"
                     VerticalContentAlignment="{TemplateBinding VerticalAlignment}">
                     
+                    <controls:Card.Styles>
+                        <Style Selector="ItemsRepeater">
+                            <Setter Property="animations:ItemsRepeaterArrangeAnimation.EnableItemsArrangeAnimation" Value="True"/>
+                        </Style>
+                    </controls:Card.Styles>
+
                     <Grid RowDefinitions="Auto,*">
                         <TextBox
                             x:Name="SearchBox"
@@ -53,10 +60,16 @@
                             </TextBox.InnerRightContent>
                         </TextBox>
 
-                        <ScrollViewer Grid.Row="1">
+                        <ScrollViewer 
+                            Grid.Row="1" 
+                            IsScrollInertiaEnabled="True"
+                            VerticalSnapPointsType="Mandatory"
+                            BringIntoViewOnFocusChange="False"
+                            HorizontalScrollBarVisibility="Disabled">
                             <ItemsRepeater
                                 HorizontalAlignment="Center"
                                 VerticalAlignment="Stretch"
+                                VerticalCacheLength="16"
                                 ItemsSource="{Binding LocalImages}">
                                 
                                 <ItemsRepeater.Resources>
@@ -107,6 +120,7 @@
                                 <ItemsRepeater.ItemTemplate>
                                     <DataTemplate DataType="{x:Type dbModels:LocalImageFile}">
                                         <Button
+                                            ToolTip.Placement="LeftEdgeAlignedTop"
                                             Command="{StaticResource ImageClickCommand}"
                                             CommandParameter="{Binding $self.DataContext}"
                                             HotKey="{x:Null}"
@@ -115,23 +129,6 @@
                                             CornerRadius="8">
                                             
                                             <Button.ContextFlyout>
-                                                <!--<ui:CommandBarFlyout>
-                                                    <ui:CommandBarButton
-                                                        Command="{StaticResource ImageCopyCommand}"
-                                                        CommandParameter="{Binding $self.DataContext}"/>
-                                                    <ui:CommandBarButton
-                                                        Command="{StaticResource ImageDeleteCommand}"
-                                                        CommandParameter="{Binding $self.DataContext}"/>
-                                                    
-                                                    <ui:CommandBarFlyout.SecondaryCommands>
-                                                        <ui:CommandBarSeparator Margin="0,-4,0,0"/>
-                                                        <ui:CommandBarButton>
-                                                            <ui:CommandBarButton
-                                                                Command="{StaticResource ImageDeleteCommand}"
-                                                                CommandParameter="{Binding $self.DataContext}"/>
-                                                        </ui:CommandBarButton>
-                                                    </ui:CommandBarFlyout.SecondaryCommands>
-                                                </ui:CommandBarFlyout>-->
                                                 <ui:FAMenuFlyout>
                                                     <ui:MenuFlyoutItem 
                                                         HotKey="{x:Null}"

StabilityMatrix.Avalonia/DesignData/MockImageIndexService.cs

@@ -1,12 +1,21 @@
 using System.Collections.Generic;
 using System.Threading.Tasks;
+using StabilityMatrix.Core.Models;
 using StabilityMatrix.Core.Models.Database;
+using StabilityMatrix.Core.Models.FileInterfaces;
 using StabilityMatrix.Core.Services;
 
 namespace StabilityMatrix.Avalonia.DesignData;
 
 public class MockImageIndexService : IImageIndexService
 {
+    /// <inheritdoc />
+    public IndexCollection<LocalImageFile, string> InferenceImages { get; } =
+        new IndexCollection<LocalImageFile, string>(null!, file => file.RelativePath)
+        {
+            RelativePath = "inference"
+        };
+
     /// <inheritdoc />
     public Task<IReadOnlyList<LocalImageFile>> GetLocalImagesByPrefix(string pathPrefix)
     {
@@ -34,11 +43,20 @@ public class MockImageIndexService : IImageIndexService
     }
 
     /// <inheritdoc />
-    public Task RefreshIndex(string subPath = "")
+    public Task RefreshIndexForAllCollections()
     {
         return Task.CompletedTask;
     }
 
+    /// <inheritdoc />
+    public Task RefreshIndex(IndexCollection<LocalImageFile, string> indexCollection)
+    {
+        return Task.CompletedTask;
+    }
+
+    /// <inheritdoc />
+    public void OnImageAdded(FilePath filePath) { }
+
     /// <inheritdoc />
     public void BackgroundRefreshIndex()
     {

StabilityMatrix.Avalonia/ViewModels/Inference/ImageFolderCardViewModel.cs

@@ -1,6 +1,7 @@
 using System;
 using System.IO;
 using System.Threading.Tasks;
+using AsyncAwaitBestPractices;
 using AsyncImageLoader;
 using Avalonia.Controls.Notifications;
 using Avalonia.Controls.Primitives;
@@ -35,18 +36,6 @@ public partial class ImageFolderCardViewModel : ViewModelBase
     private readonly ISettingsManager settingsManager;
     private readonly INotificationService notificationService;
 
-    /// <summary>
-    /// Source of image files to display
-    /// </summary>
-    private readonly SourceCache<LocalImageFile, string> localImagesSource =
-        new(imageFile => imageFile.RelativePath);
-
-    /// <summary>
-    /// Collection of image items to display
-    /// </summary>
-    public IObservableCollection<ImageFolderCardItemViewModel> Items { get; } =
-        new ObservableCollectionExtended<ImageFolderCardItemViewModel>();
-
     /// <summary>
     /// Collection of local image files
     /// </summary>
@@ -65,7 +54,7 @@ public partial class ImageFolderCardViewModel : ViewModelBase
         this.settingsManager = settingsManager;
         this.notificationService = notificationService;
 
-        localImagesSource
+        imageIndexService.InferenceImages.ItemsSource
             .Connect()
             .DeferUntilLoaded()
             .SortBy(file => file.LastModifiedAt, SortDirection.Descending)
@@ -78,14 +67,7 @@ public partial class ImageFolderCardViewModel : ViewModelBase
     {
         await base.OnLoadedAsync();
 
-        await imageIndexService.RefreshIndex("Inference");
+        imageIndexService.RefreshIndexForAllCollections().SafeFireAndForget();
-
-        var imageFiles = await imageIndexService.GetLocalImagesByPrefix("Inference");
-
-        localImagesSource.Edit(x =>
-        {
-            x.Load(imageFiles);
-        });
     }
 
     /// <summary>
@@ -141,11 +123,8 @@ public partial class ImageFolderCardViewModel : ViewModelBase
             return;
         }
 
-        // Remove from source
-        localImagesSource.Remove(item);
-
         // Remove from index
-        await imageIndexService.RemoveImage(item);
+        imageIndexService.InferenceImages.Remove(item);
 
         // Invalidate cache
         if (ImageLoader.AsyncImageLoader is FallbackRamCachedWebImageLoader loader)

StabilityMatrix.Avalonia/ViewModels/Inference/InferenceTextToImageViewModel.cs

@@ -28,6 +28,8 @@ using StabilityMatrix.Core.Models.Api.Comfy;
 using StabilityMatrix.Core.Models.Api.Comfy.Nodes;
 using StabilityMatrix.Core.Models.Api.Comfy.NodeTypes;
 using StabilityMatrix.Core.Models.Api.Comfy.WebSocketData;
+using StabilityMatrix.Core.Models.Database;
+using StabilityMatrix.Core.Models.FileInterfaces;
 using StabilityMatrix.Core.Services;
 using InferenceTextToImageView = StabilityMatrix.Avalonia.Views.Inference.InferenceTextToImageView;
 
@@ -43,6 +45,7 @@ public partial class InferenceTextToImageViewModel : InferenceTabViewModelBase
     private readonly INotificationService notificationService;
     private readonly ServiceManager<ViewModelBase> vmFactory;
     private readonly IModelIndexService modelIndexService;
+    private readonly IImageIndexService imageIndexService;
 
     [JsonIgnore]
     public IInferenceClientManager ClientManager { get; }
@@ -103,12 +106,14 @@ public partial class InferenceTextToImageViewModel : InferenceTabViewModelBase
         INotificationService notificationService,
         IInferenceClientManager inferenceClientManager,
         ServiceManager<ViewModelBase> vmFactory,
-        IModelIndexService modelIndexService
+        IModelIndexService modelIndexService,
+        IImageIndexService imageIndexService
     )
     {
         this.notificationService = notificationService;
         this.vmFactory = vmFactory;
         this.modelIndexService = modelIndexService;
+        this.imageIndexService = imageIndexService;
         ClientManager = inferenceClientManager;
 
         // Get sub view models from service manager
@@ -416,89 +421,103 @@ public partial class InferenceTextToImageViewModel : InferenceTabViewModelBase
                 notificationService.Show("No output", "Did not receive any output images");
                 return;
             }
+        }
+        finally
+        {
+            // Disconnect progress handler
+            OutputProgress.Value = 0;
+            OutputProgress.Text = "";
+            ImageGalleryCardViewModel.PreviewImage?.Dispose();
+            ImageGalleryCardViewModel.PreviewImage = null;
+            ImageGalleryCardViewModel.IsPreviewOverlayEnabled = false;
 
-            List<ImageSource> outputImages;
+            promptTask?.Dispose();
-            // Use local file path if available, otherwise use remote URL
+            client.PreviewImageReceived -= OnPreviewImageReceived;
-            if (client.OutputImagesDir is { } outputPath)
+        }
-            {
+    }
-                outputImages = new List<ImageSource>();
-                foreach (var image in images)
-                {
-                    var filePath = image.ToFilePath(outputPath);
-                    var fileStream = new BinaryReader(filePath.Info.OpenRead());
-                    var bytes = fileStream.ReadBytes((int)filePath.Info.Length);
-                    var bytesWithMetadata = PngDataHelper.AddMetadata(
-                        bytes,
-                        generationInfo,
-                        smproj
-                    );
-                    fileStream.Close();
-                    fileStream.Dispose();
-
-                    await using var outputStream = filePath.Info.OpenWrite();
-                    await outputStream.WriteAsync(bytesWithMetadata);
-                    await outputStream.FlushAsync();
 
-                    outputImages.Add(new ImageSource(filePath));
+    private async Task ProcessOutputs(IReadOnlyList<ComfyImage> images)
-                }
+    {
-            }
+        List<ImageSource> outputImages;
-            else
+        // Use local file path if available, otherwise use remote URL
+        if (client.OutputImagesDir is { } outputPath)
+        {
+            outputImages = new List<ImageSource>();
+            foreach (var image in images)
             {
-                outputImages = images!
+                var filePath = image.ToFilePath(outputPath);
-                    .Select(i => new ImageSource(i.ToUri(client.BaseAddress)))
-                    .ToList();
-            }
 
-            // Download all images to make grid, if multiple
+                var bytesWithMetadata = PngDataHelper.AddMetadata(
-            if (outputImages.Count > 1)
+                    await filePath.ReadAllBytesAsync(),
-            {
-                var loadedImages = outputImages
-                    .Select(i => SKImage.FromEncodedData(i.LocalFile?.Info.OpenRead()))
-                    .ToImmutableArray();
-
-                var grid = ImageProcessor.CreateImageGrid(loadedImages);
-                var gridBytes = grid.Encode().ToArray();
-                var gridBytesWithMetadata = PngDataHelper.AddMetadata(
-                    gridBytes,
                     generationInfo,
                     smproj
                 );
 
-                // Save to disk
+                /*await using (var readStream = filePath.Info.OpenWrite())
-                var lastName = outputImages.Last().LocalFile?.Info.Name;
+                {
-                var gridPath = client.OutputImagesDir!.JoinFile($"grid-{lastName}");
+                    using (var reader = new BinaryReader(readStream))
+                    {
+
+                    }
+                }*/
 
-                await using (var fileStream = gridPath.Info.OpenWrite())
+                await using (var outputStream = filePath.Info.OpenWrite())
                 {
-                    await fileStream.WriteAsync(gridBytesWithMetadata, cancellationToken);
+                    await outputStream.WriteAsync(bytesWithMetadata);
+                    await outputStream.FlushAsync();
                 }
 
-                // Insert to start of images
+                outputImages.Add(new ImageSource(filePath));
-                var gridImage = new ImageSource(gridPath);
+
-                // Preload
+                imageIndexService.OnImageAdded(filePath);
-                await gridImage.GetBitmapAsync();
-                ImageGalleryCardViewModel.ImageSources.Add(gridImage);
             }
+        }
+        else
+        {
+            outputImages = images!
+                .Select(i => new ImageSource(i.ToUri(client.BaseAddress)))
+                .ToList();
+        }
+
+        // Download all images to make grid, if multiple
+        if (outputImages.Count > 1)
+        {
+            var loadedImages = outputImages
+                .Select(i => SKImage.FromEncodedData(i.LocalFile?.Info.OpenRead()))
+                .ToImmutableArray();
+
+            var grid = ImageProcessor.CreateImageGrid(loadedImages);
+            var gridBytes = grid.Encode().ToArray();
+            var gridBytesWithMetadata = PngDataHelper.AddMetadata(
+                gridBytes,
+                generationInfo,
+                smproj
+            );
+
+            // Save to disk
+            var lastName = outputImages.Last().LocalFile?.Info.Name;
+            var gridPath = client.OutputImagesDir!.JoinFile($"grid-{lastName}");
 
-            // Add rest of images
+            await using (var fileStream = gridPath.Info.OpenWrite())
-            foreach (var img in outputImages)
             {
-                // Preload
+                await fileStream.WriteAsync(gridBytesWithMetadata, cancellationToken);
-                await img.GetBitmapAsync();
-                ImageGalleryCardViewModel.ImageSources.Add(img);
             }
+
+            // Insert to start of images
+            var gridImage = new ImageSource(gridPath);
+            // Preload
+            await gridImage.GetBitmapAsync();
+            ImageGalleryCardViewModel.ImageSources.Add(gridImage);
+
+            imageIndexService.OnImageAdded(gridPath);
         }
-        finally
-        {
-            // Disconnect progress handler
-            OutputProgress.Value = 0;
-            OutputProgress.Text = "";
-            ImageGalleryCardViewModel.PreviewImage?.Dispose();
-            ImageGalleryCardViewModel.PreviewImage = null;
-            ImageGalleryCardViewModel.IsPreviewOverlayEnabled = false;
 
-            promptTask?.Dispose();
+        // Add rest of images
-            client.PreviewImageReceived -= OnPreviewImageReceived;
+        foreach (var img in outputImages)
+        {
+            // Preload
+            await img.GetBitmapAsync();
+            ImageGalleryCardViewModel.ImageSources.Add(img);
         }
     }
 
@@ -517,7 +536,6 @@ public partial class InferenceTextToImageViewModel : InferenceTabViewModelBase
             };
 
             await GenerateImageImpl(overrides, cancellationToken);
-            await ImageFolderCardViewModel.OnLoadedAsync();
         }
         catch (OperationCanceledException e)
         {

StabilityMatrix.Avalonia/ViewModels/Inference/PromptCardViewModel.cs

@@ -133,13 +133,13 @@ public partial class PromptCardViewModel : LoadableViewModelBase
     {
         var md = $"""
                   ## {Resources.Label_Emphasis}
-                  ```python
+                  ```prompt
                   (keyword)
-                  (keyword:weight)
+                  (keyword:1.0)
                   ```
                   
                   ## {Resources.Label_Deemphasis}
-                  ```python
+                  ```prompt
                   [keyword]
                   ```
                   
@@ -147,25 +147,25 @@ public partial class PromptCardViewModel : LoadableViewModelBase
                   They may be used in either the positive or negative prompts. 
                   Essentially they are text presets, so the position where you place them 
                   could make a difference. 
-                  ```python
+                  ```prompt
                   <embedding:model>
-                  <embedding:model:weight>
+                  <embedding:model:1.0>
                   ```
                   
                   ## {Resources.Label_NetworksLoraOrLycoris}
                   Unlike embeddings, network tags do not get tokenized to the model, 
                   so the position in the prompt where you place them does not matter.
-                  ```python
+                  ```prompt
                   <lora:model>
-                  <lora:model:weight>
+                  <lora:model:1.0>
                   <lyco:model>
-                  <lyco:model:weight>
+                  <lyco:model:1.0>
                   ```
                   
                   ## {Resources.Label_Comments}
                   Inline comments can be marked by a hashtag ‘#’. 
                   All text after a ‘#’ on a line will be disregarded during generation.
-                  ```c
+                  ```prompt
                   # comments
                   a red cat # also comments
                   detailed

StabilityMatrix.Avalonia/ViewModels/InferenceViewModel.cs

@@ -224,7 +224,7 @@ public partial class InferenceViewModel : PageViewModelBase
         }
 
         // Start a model index update
-        await modelIndexService.RefreshIndex();
+        modelIndexService.BackgroundRefreshIndex();
     }
 
     /// <summary>

StabilityMatrix.Core/Models/Database/LocalImageFile.cs

@@ -1,4 +1,7 @@
 using LiteDB;
+using StabilityMatrix.Core.Helper;
+using StabilityMatrix.Core.Models.FileInterfaces;
+using JsonSerializer = System.Text.Json.JsonSerializer;
 
 namespace StabilityMatrix.Core.Models.Database;
 
@@ -51,6 +54,83 @@ public class LocalImageFile
         return Path.Combine(rootImageDirectory, RelativePath);
     }
 
+    public static LocalImageFile FromPath(FilePath filePath)
+    {
+        var relativePath = Path.GetRelativePath(
+            GlobalConfig.LibraryDir.JoinDir("Images"),
+            filePath
+        );
+
+        // TODO: Support other types
+        const LocalImageFileType imageType =
+            LocalImageFileType.Inference | LocalImageFileType.TextToImage;
+
+        // Get metadata
+        using var stream = new FileStream(
+            filePath.FullPath,
+            FileMode.Open,
+            FileAccess.Read,
+            FileShare.Read
+        );
+        using var reader = new BinaryReader(stream);
+        var metadata = ImageMetadata.ReadTextChunk(reader, "parameters-json");
+
+        GenerationParameters? genParams = null;
+
+        if (!string.IsNullOrWhiteSpace(metadata))
+        {
+            genParams = JsonSerializer.Deserialize<GenerationParameters>(metadata);
+        }
+        else
+        {
+            metadata = ImageMetadata.ReadTextChunk(reader, "parameters");
+            GenerationParameters.TryParse(metadata, out genParams);
+        }
+
+        return new LocalImageFile
+        {
+            RelativePath = relativePath,
+            ImageType = imageType,
+            CreatedAt = filePath.Info.CreationTimeUtc,
+            LastModifiedAt = filePath.Info.LastWriteTimeUtc,
+            GenerationParameters = genParams
+        };
+    }
+
     public static readonly HashSet<string> SupportedImageExtensions =
         new() { ".png", ".jpg", ".jpeg", ".webp" };
+
+    private sealed class LocalImageFileEqualityComparer : IEqualityComparer<LocalImageFile>
+    {
+        public bool Equals(LocalImageFile? x, LocalImageFile? y)
+        {
+            if (ReferenceEquals(x, y))
+                return true;
+            if (ReferenceEquals(x, null))
+                return false;
+            if (ReferenceEquals(y, null))
+                return false;
+            if (x.GetType() != y.GetType())
+                return false;
+            return x.RelativePath == y.RelativePath
+                && x.ImageType == y.ImageType
+                && x.CreatedAt.Equals(y.CreatedAt)
+                && x.LastModifiedAt.Equals(y.LastModifiedAt)
+                && Equals(x.GenerationParameters, y.GenerationParameters);
+        }
+
+        public int GetHashCode(LocalImageFile obj)
+        {
+            return HashCode.Combine(
+                obj.RelativePath,
+                obj.ImageType,
+                obj.CreatedAt,
+                obj.LastModifiedAt,
+                obj.GenerationParameters
+            );
+        }
+    }
+
+    public static IEqualityComparer<LocalImageFile> Comparer { get; } =
+        new LocalImageFileEqualityComparer();
 }

StabilityMatrix.Core/Models/GenerationParameters.cs

@@ -19,10 +19,16 @@ public partial class GenerationParameters
     public string? ModelName { get; set; }
 
     public static bool TryParse(
-        string text,
+        string? text,
         [NotNullWhen(true)] out GenerationParameters? generationParameters
     )
     {
+        if (string.IsNullOrWhiteSpace(text))
+        {
+            generationParameters = null;
+            return false;
+        }
+
         var lines = text.Split('\n');
 
         if (lines.LastOrDefault() is not { } lastLine)

StabilityMatrix.Core/Models/IndexCollection.cs

@@ -0,0 +1,59 @@
+using DynamicData;
+using DynamicData.Binding;
+using StabilityMatrix.Core.Services;
+
+namespace StabilityMatrix.Core.Models;
+
+public class IndexCollection<TObject, TKey>
+    where TKey : notnull
+{
+    private readonly IImageIndexService imageIndexService;
+
+    public string? RelativePath { get; set; }
+
+    public SourceCache<TObject, TKey> ItemsSource { get; }
+
+    /// <summary>
+    /// Observable Collection of indexed items
+    /// </summary>
+    public IObservableCollection<TObject> Items { get; } =
+        new ObservableCollectionExtended<TObject>();
+
+    public IndexCollection(
+        IImageIndexService imageIndexService,
+        Func<TObject, TKey> keySelector,
+        Func<
+            IObservable<IChangeSet<TObject, TKey>>,
+            IObservable<IChangeSet<TObject, TKey>>
+        >? transform = null
+    )
+    {
+        this.imageIndexService = imageIndexService;
+
+        ItemsSource = new SourceCache<TObject, TKey>(keySelector);
+
+        var source = ItemsSource.Connect().DeferUntilLoaded();
+
+        if (transform is not null)
+        {
+            source = transform(source);
+        }
+
+        source.Bind(Items).Subscribe();
+    }
+
+    public void Add(TObject item)
+    {
+        ItemsSource.AddOrUpdate(item);
+    }
+
+    public void Remove(TObject item)
+    {
+        ItemsSource.Remove(item);
+    }
+
+    public void RemoveKey(TKey key)
+    {
+        ItemsSource.RemoveKey(key);
+    }
+}

StabilityMatrix.Core/Services/IImageIndexService.cs

@@ -1,18 +1,27 @@
-using StabilityMatrix.Core.Models.Database;
+using DynamicData.Binding;
+using StabilityMatrix.Core.Models;
+using StabilityMatrix.Core.Models.Database;
+using StabilityMatrix.Core.Models.FileInterfaces;
 
 namespace StabilityMatrix.Core.Services;
 
 public interface IImageIndexService
 {
+    IndexCollection<LocalImageFile, string> InferenceImages { get; }
+
     /// <summary>
     /// Gets a list of local images that start with the given path prefix
     /// </summary>
     Task<IReadOnlyList<LocalImageFile>> GetLocalImagesByPrefix(string pathPrefix);
 
     /// <summary>
-    /// Refreshes the index of local images
+    /// Refresh index for all collections
     /// </summary>
-    Task RefreshIndex(string subPath = "");
+    Task RefreshIndexForAllCollections();
+
+    Task RefreshIndex(IndexCollection<LocalImageFile, string> indexCollection);
+
+    void OnImageAdded(FilePath filePath);
 
     /// <summary>
     /// Refreshes the index of local images in the background

StabilityMatrix.Core/Services/ImageIndexService.cs

@@ -1,6 +1,8 @@
 using System.Diagnostics;
 using System.Text.Json;
 using AsyncAwaitBestPractices;
+using DynamicData;
+using DynamicData.Binding;
 using Microsoft.Extensions.Logging;
 using StabilityMatrix.Core.Database;
 using StabilityMatrix.Core.Helper;
@@ -16,6 +18,9 @@ public class ImageIndexService : IImageIndexService
     private readonly ILiteDbContext liteDbContext;
     private readonly ISettingsManager settingsManager;
 
+    /// <inheritdoc />
+    public IndexCollection<LocalImageFile, string> InferenceImages { get; }
+
     public ImageIndexService(
         ILogger<ImageIndexService> logger,
         ILiteDbContext liteDbContext,
@@ -25,6 +30,21 @@ public class ImageIndexService : IImageIndexService
         this.logger = logger;
         this.liteDbContext = liteDbContext;
         this.settingsManager = settingsManager;
+
+        InferenceImages = new IndexCollection<LocalImageFile, string>(
+            this,
+            file => file.RelativePath
+        )
+        {
+            RelativePath = "inference"
+        };
+
+        /*inferenceImagesSource
+            .Connect()
+            .DeferUntilLoaded()
+            .SortBy(file => file.LastModifiedAt, SortDirection.Descending)
+            .Bind(InferenceImages)
+            .Subscribe();*/
     }
 
     /// <inheritdoc />
@@ -37,8 +57,86 @@ public class ImageIndexService : IImageIndexService
             .ConfigureAwait(false);
     }
 
+    public async Task RefreshIndexForAllCollections()
+    {
+        await RefreshIndex(InferenceImages).ConfigureAwait(false);
+    }
+
+    public async Task RefreshIndex(IndexCollection<LocalImageFile, string> indexCollection)
+    {
+        if (indexCollection.RelativePath is not { } subPath)
+            return;
+
+        var imagesDir = settingsManager.ImagesDirectory;
+        var searchDir = imagesDir.JoinDir(indexCollection.RelativePath);
+        if (!searchDir.Exists)
+        {
+            return;
+        }
+
+        // Start
+        var stopwatch = Stopwatch.StartNew();
+        logger.LogInformation("Refreshing images index at {ImagesDir}...", imagesDir);
+
+        var added = 0;
+        var toAdd = new Queue<LocalImageFile>();
+
+        await Task.Run(() =>
+            {
+                foreach (
+                    var file in imagesDir.Info
+                        .EnumerateFiles("*.*", SearchOption.AllDirectories)
+                        .Where(
+                            info => LocalImageFile.SupportedImageExtensions.Contains(info.Extension)
+                        )
+                        .Select(info => new FilePath(info))
+                )
+                {
+                    /*var relativePath = Path.GetRelativePath(imagesDir, file);
+                    
+                    if (string.IsNullOrEmpty(relativePath))
+                    {
+                        continue;
+                    }*/
+
+                    var localImage = LocalImageFile.FromPath(file);
+
+                    toAdd.Enqueue(localImage);
+
+                    added++;
+                }
+            })
+            .ConfigureAwait(false);
+
+        var indexElapsed = stopwatch.Elapsed;
+
+        indexCollection.ItemsSource.EditDiff(toAdd, LocalImageFile.Comparer);
+
+        // End
+        stopwatch.Stop();
+        var editElapsed = stopwatch.Elapsed - indexElapsed;
+
+        logger.LogInformation(
+            "Image index updated for {Prefix} with {Entries} files, took {IndexDuration:F1}ms ({EditDuration:F1}ms edit)",
+            subPath,
+            added,
+            indexElapsed.TotalMilliseconds,
+            editElapsed.TotalMilliseconds
+        );
+    }
+
     /// <inheritdoc />
-    public async Task RefreshIndex(string subPath = "")
+    public void OnImageAdded(FilePath filePath)
+    {
+        var fullPath = settingsManager.ImagesDirectory.JoinDir(InferenceImages.RelativePath!);
+
+        if (!string.IsNullOrEmpty(Path.GetRelativePath(fullPath, filePath)))
+        {
+            InferenceImages.Add(LocalImageFile.FromPath(filePath));
+        }
+    }
+
+    /*public async Task RefreshIndex(IndexCollection<LocalImageFile, string> indexCollection)
     {
         var imagesDir = settingsManager.ImagesDirectory;
         if (!imagesDir.Exists)
@@ -129,12 +227,12 @@ public class ImageIndexService : IImageIndexService
             indexDuration.TotalMilliseconds,
             dbDuration.TotalMilliseconds
         );
-    }
+    }*/
 
     /// <inheritdoc />
     public void BackgroundRefreshIndex()
     {
-        RefreshIndex().SafeFireAndForget();
+        RefreshIndexForAllCollections().SafeFireAndForget();
     }
 
     /// <inheritdoc />