Add ImageIndexService and Inference Image Viewer

StabilityMatrix.Avalonia/App.axaml

@@ -50,5 +50,6 @@
         <StyleInclude Source="Controls/FrameCarousel.axaml"/>
         <StyleInclude Source="Controls/CodeCompletion/CompletionWindow.axaml"/>
         <StyleInclude Source="Controls/SelectImageCard.axaml"/>
+        <StyleInclude Source="Controls/ImageFolderCard.axaml"/>
     </Application.Styles>
 </Application>

StabilityMatrix.Avalonia/App.axaml.cs

@@ -307,6 +307,7 @@ public sealed class App : Application
         services.AddTransient<SamplerCardViewModel>();
         services.AddTransient<UpscalerCardViewModel>();
         services.AddTransient<ImageGalleryCardViewModel>();
+        services.AddTransient<ImageFolderCardViewModel>();
         services.AddTransient<PromptCardViewModel>();
         services.AddTransient<StackCardViewModel>();
         services.AddTransient<StackExpanderViewModel>();
@@ -337,6 +338,7 @@ public sealed class App : Application
                     .Register(provider.GetRequiredService<SeedCardViewModel>)
                     .Register(provider.GetRequiredService<SamplerCardViewModel>)
                     .Register(provider.GetRequiredService<ImageGalleryCardViewModel>)
+                    .Register(provider.GetRequiredService<ImageFolderCardViewModel>)
                     .Register(provider.GetRequiredService<PromptCardViewModel>)
                     .Register(provider.GetRequiredService<StackCardViewModel>)
                     .Register(provider.GetRequiredService<StackExpanderViewModel>)
@@ -369,6 +371,7 @@ public sealed class App : Application
 
         // Inference controls
         services.AddTransient<ImageGalleryCard>();
+        services.AddTransient<ImageFolderCard>();
         services.AddTransient<SeedCard>();
         services.AddTransient<SamplerCard>();
         services.AddTransient<PromptCard>();
@@ -434,6 +437,7 @@ public sealed class App : Application
         services.AddSingleton<ICompletionProvider, CompletionProvider>();
         services.AddSingleton<ITokenizerProvider, TokenizerProvider>();
         services.AddSingleton<IModelIndexService, ModelIndexService>();
+        services.AddSingleton<IImageIndexService, ImageIndexService>();
 
         services.AddSingleton<ITrackedDownloadService, TrackedDownloadService>();
         services.AddSingleton<IDisposable>(

StabilityMatrix.Avalonia/Controls/ImageFolderCard.axaml

@@ -0,0 +1,62 @@
+<Styles
+    xmlns="https://github.com/avaloniaui"
+    xmlns:x="http://schemas.microsoft.com/winfx/2006/xaml"
+    xmlns:controls="using:StabilityMatrix.Avalonia.Controls"
+    xmlns:mocks="clr-namespace:StabilityMatrix.Avalonia.DesignData"
+    xmlns:models="clr-namespace:StabilityMatrix.Avalonia.Models"
+    xmlns:vmInference="clr-namespace:StabilityMatrix.Avalonia.ViewModels.Inference"
+    x:DataType="vmInference:ImageFolderCardViewModel">
+
+    <Design.PreviewWith>
+        <Panel Width="600" Height="800">
+            <Grid MaxWidth="500" MaxHeight="700">
+                <controls:ImageFolderCard 
+                    DataContext="{x:Static mocks:DesignData.ImageFolderCardViewModel}" />
+            </Grid>
+        </Panel>
+    </Design.PreviewWith>
+
+
+    <!--<Style Selector="ListBox /template/ VirtualizingStackPanel">
+        <Setter Property="Orientation" Value="Horizontal" />
+    </Style>-->
+
+    <Style Selector="controls|ImageFolderCard">
+        <!--  Set Defaults  -->
+        <Setter Property="Template">
+            <ControlTemplate>
+                <controls:Card
+                    Padding="8"
+                    HorizontalAlignment="{TemplateBinding HorizontalAlignment}"
+                    VerticalAlignment="{TemplateBinding VerticalAlignment}"
+                    HorizontalContentAlignment="{TemplateBinding HorizontalAlignment}"
+                    VerticalContentAlignment="{TemplateBinding VerticalAlignment}">
+
+                    <ScrollViewer>
+                        <ItemsRepeater
+                            HorizontalAlignment="Center"
+                            VerticalAlignment="Stretch"
+                            ItemsSource="{Binding Items}">
+                            <ItemsRepeater.Layout>
+                                <UniformGridLayout
+                                    MinColumnSpacing="8"
+                                    MinRowSpacing="8"
+                                    MaximumRowsOrColumns="2"/>
+                            </ItemsRepeater.Layout>
+                            <ItemsRepeater.ItemTemplate>
+                                <DataTemplate DataType="{x:Type vmInference:ImageFolderCardItemViewModel}">
+                                    <controls:BetterAdvancedImage
+                                        MaxHeight="256"
+                                        CornerRadius="4"
+                                        Source="{Binding ImagePath}"
+                                        Stretch="Uniform"
+                                        StretchDirection="Both" />
+                                </DataTemplate>
+                            </ItemsRepeater.ItemTemplate>
+                        </ItemsRepeater>
+                    </ScrollViewer>
+                </controls:Card>
+            </ControlTemplate>
+        </Setter>
+    </Style>
+</Styles>

StabilityMatrix.Avalonia/Controls/ImageFolderCard.axaml.cs

@@ -0,0 +1,27 @@
+using AsyncAwaitBestPractices;
+using Avalonia.Controls.Primitives;
+using Avalonia.Interactivity;
+using Avalonia.Threading;
+using StabilityMatrix.Avalonia.ViewModels.Base;
+
+namespace StabilityMatrix.Avalonia.Controls;
+
+public class ImageFolderCard : TemplatedControl
+{
+    /// <inheritdoc />
+    protected override void OnLoaded(RoutedEventArgs e)
+    {
+        base.OnLoaded(e);
+
+        if (DataContext is ViewModelBase vm)
+        {
+            vm.OnLoaded();
+            Dispatcher.UIThread
+                .InvokeAsync(async () =>
+                {
+                    await vm.OnLoadedAsync();
+                })
+                .SafeFireAndForget();
+        }
+    }
+}

StabilityMatrix.Avalonia/DesignData/DesignData.cs

@@ -104,6 +104,7 @@ public static class DesignData
             .AddSingleton<IDiscordRichPresenceService, MockDiscordRichPresenceService>()
             .AddSingleton<ICompletionProvider, MockCompletionProvider>()
             .AddSingleton<IModelIndexService, MockModelIndexService>()
+            .AddSingleton<IImageIndexService, MockImageIndexService>()
             .AddSingleton<ITrackedDownloadService, MockTrackedDownloadService>();
 
         // Placeholder services that nobody should need during design time
@@ -511,6 +512,9 @@ The gallery images are often inpainted, but you will get something very similar
             );
         });
 
+    public static ImageFolderCardViewModel ImageFolderCardViewModel =>
+        DialogFactory.Get<ImageFolderCardViewModel>();
+
     public static PromptCardViewModel PromptCardViewModel =>
         DialogFactory.Get<PromptCardViewModel>(vm =>
         {

StabilityMatrix.Avalonia/DesignData/MockImageIndexService.cs

@@ -0,0 +1,47 @@
+using System.Collections.Generic;
+using System.Threading.Tasks;
+using StabilityMatrix.Core.Models.Database;
+using StabilityMatrix.Core.Services;
+
+namespace StabilityMatrix.Avalonia.DesignData;
+
+public class MockImageIndexService : IImageIndexService
+{
+    /// <inheritdoc />
+    public Task<IReadOnlyList<LocalImageFile>> GetLocalImagesByPrefix(string pathPrefix)
+    {
+        return Task.FromResult(
+            (IReadOnlyList<LocalImageFile>)
+                new LocalImageFile[]
+                {
+                    new()
+                    {
+                        RelativePath =
+                            "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/4a7e00a7-6f18-42d4-87c0-10e792df2640/width=1152",
+                    },
+                    new()
+                    {
+                        RelativePath =
+                            "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/a318ac1f-3ad0-48ac-98cc-79126febcc17/width=1024",
+                    },
+                    new()
+                    {
+                        RelativePath =
+                            "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/16588c94-6595-4be9-8806-d7e6e22d198c/width=1152",
+                    }
+                }
+        );
+    }
+
+    /// <inheritdoc />
+    public Task RefreshIndex(string subPath = "")
+    {
+        return Task.CompletedTask;
+    }
+
+    /// <inheritdoc />
+    public void BackgroundRefreshIndex()
+    {
+        throw new System.NotImplementedException();
+    }
+}

StabilityMatrix.Avalonia/DesignData/MockLiteDbContext.cs

@@ -20,6 +20,8 @@ public class MockLiteDbContext : ILiteDbContext
         throw new NotImplementedException();
     public ILiteCollectionAsync<InferenceProjectEntry> InferenceProjects =>
         throw new NotImplementedException();
+    public ILiteCollectionAsync<LocalImageFile> LocalImageFiles =>
+        throw new NotImplementedException();
 
     public Task<(CivitModel?, CivitModelVersion?)> FindCivitModelFromFileHashAsync(
         string hashBlake3

StabilityMatrix.Avalonia/Extensions/ComfyNodeBuilderExtensions.cs

@@ -266,7 +266,7 @@ public static class ComfyNodeBuilderExtensions
                 ClassType = "SaveImage",
                 Inputs = new Dictionary<string, object?>
                 {
-                    ["filename_prefix"] = "SM-Inference",
+                    ["filename_prefix"] = "Inference/TextToImage",
                     ["images"] = builder.Connections.Image
                 }
             }

StabilityMatrix.Avalonia/Services/InferenceClientManager.cs

@@ -1,11 +1,6 @@
 using System;
-using System.Collections.Generic;
-using System.Collections.Immutable;
-using System.Collections.ObjectModel;
 using System.Diagnostics.CodeAnalysis;
-using System.IO;
 using System.Linq;
-using System.Reactive.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 using AsyncAwaitBestPractices;
@@ -13,12 +8,10 @@ using CommunityToolkit.Mvvm.ComponentModel;
 using DynamicData;
 using DynamicData.Binding;
 using Microsoft.Extensions.Logging;
-using StabilityMatrix.Avalonia.ViewModels.PackageManager;
 using StabilityMatrix.Core.Api;
 using StabilityMatrix.Core.Inference;
 using StabilityMatrix.Core.Models;
 using StabilityMatrix.Core.Models.Api.Comfy;
-using StabilityMatrix.Core.Models.Database;
 using StabilityMatrix.Core.Models.FileInterfaces;
 using StabilityMatrix.Core.Models.Packages;
 using StabilityMatrix.Core.Services;
@@ -258,11 +251,17 @@ public partial class InferenceClientManager : ObservableObject, IInferenceClient
         if (IsConnected)
             return;
 
-        if (packagePair.BasePackage is not ComfyUI)
+        if (packagePair.BasePackage is not ComfyUI comfyPackage)
         {
             throw new ArgumentException("Base package is not ComfyUI", nameof(packagePair));
         }
 
+        // Setup image folder links
+        await comfyPackage.SetupInferenceOutputFolderLinks(
+            packagePair.InstalledPackage.FullPath
+                ?? throw new InvalidOperationException("Package does not have a Path")
+        );
+
         // Get user defined host and port
         var host = packagePair.InstalledPackage.GetLaunchArgsHost() ?? "127.0.0.1";
         host = host.Replace("localhost", "127.0.0.1");

StabilityMatrix.Avalonia/ViewModels/Inference/ImageFolderCardItemViewModel.cs

@@ -0,0 +1,14 @@
+using CommunityToolkit.Mvvm.ComponentModel;
+using StabilityMatrix.Avalonia.ViewModels.Base;
+using StabilityMatrix.Core.Models.Database;
+
+namespace StabilityMatrix.Avalonia.ViewModels.Inference;
+
+public partial class ImageFolderCardItemViewModel : ViewModelBase
+{
+    [ObservableProperty]
+    private LocalImageFile? localImageFile;
+
+    [ObservableProperty]
+    private string? imagePath;
+}

StabilityMatrix.Avalonia/ViewModels/Inference/ImageFolderCardViewModel.cs

@@ -0,0 +1,83 @@
+using System;
+using System.Reactive.Linq;
+using System.Threading;
+using System.Threading.Tasks;
+using Avalonia.Controls;
+using DynamicData;
+using DynamicData.Binding;
+using Microsoft.Extensions.Logging;
+using StabilityMatrix.Avalonia.Controls;
+using StabilityMatrix.Avalonia.ViewModels.Base;
+using StabilityMatrix.Core.Attributes;
+using StabilityMatrix.Core.Models.Database;
+using StabilityMatrix.Core.Services;
+
+namespace StabilityMatrix.Avalonia.ViewModels.Inference;
+
+[View(typeof(ImageFolderCard))]
+public partial class ImageFolderCardViewModel : ViewModelBase
+{
+    private readonly ILogger<ImageFolderCardViewModel> logger;
+    private readonly IImageIndexService imageIndexService;
+    private readonly ISettingsManager settingsManager;
+
+    /// <summary>
+    /// Source of image files to display
+    /// </summary>
+    private readonly SourceCache<LocalImageFile, string> localImagesSource =
+        new(imageFile => imageFile.RelativePath);
+
+    /// <summary>
+    /// Collection of image files to display
+    /// </summary>
+    public IObservableCollection<LocalImageFile> LocalImages { get; } =
+        new ObservableCollectionExtended<LocalImageFile>();
+
+    /// <summary>
+    /// Collection of image items to display
+    /// </summary>
+    public IObservableCollection<ImageFolderCardItemViewModel> Items { get; } =
+        new ObservableCollectionExtended<ImageFolderCardItemViewModel>();
+
+    public ImageFolderCardViewModel(
+        ILogger<ImageFolderCardViewModel> logger,
+        IImageIndexService imageIndexService,
+        ISettingsManager settingsManager
+    )
+    {
+        this.logger = logger;
+        this.imageIndexService = imageIndexService;
+        this.settingsManager = settingsManager;
+
+        localImagesSource
+            .Connect()
+            .DeferUntilLoaded()
+            .Transform(
+                imageFile =>
+                    new ImageFolderCardItemViewModel
+                    {
+                        LocalImageFile = imageFile,
+                        ImagePath = Design.IsDesignMode
+                            ? imageFile.RelativePath
+                            : imageFile.GetFullPath(settingsManager.ImagesDirectory)
+                    }
+            )
+            .Bind(Items)
+            .Subscribe();
+    }
+
+    /// <inheritdoc />
+    public override async Task OnLoadedAsync()
+    {
+        await base.OnLoadedAsync();
+
+        await imageIndexService.RefreshIndex("Inference");
+
+        var imageFiles = await imageIndexService.GetLocalImagesByPrefix("Inference");
+
+        localImagesSource.Edit(x =>
+        {
+            x.Load(imageFiles);
+        });
+    }
+}

StabilityMatrix.Avalonia/ViewModels/Inference/InferenceTextToImageViewModel.cs

@@ -57,6 +57,9 @@ public partial class InferenceTextToImageViewModel : InferenceTabViewModelBase
     [JsonPropertyName("ImageGallery")]
     public ImageGalleryCardViewModel ImageGalleryCardViewModel { get; }
 
+    [JsonPropertyName("ImageFolder")]
+    public ImageFolderCardViewModel ImageFolderCardViewModel { get; }
+
     [JsonPropertyName("Prompt")]
     public PromptCardViewModel PromptCardViewModel { get; }
 
@@ -122,6 +125,7 @@ public partial class InferenceTextToImageViewModel : InferenceTabViewModelBase
         });
 
         ImageGalleryCardViewModel = vmFactory.Get<ImageGalleryCardViewModel>();
+        ImageFolderCardViewModel = vmFactory.Get<ImageFolderCardViewModel>();
         PromptCardViewModel = vmFactory.Get<PromptCardViewModel>();
         HiresSamplerCardViewModel = vmFactory.Get<SamplerCardViewModel>(samplerCard =>
         {
@@ -475,6 +479,7 @@ public partial class InferenceTextToImageViewModel : InferenceTabViewModelBase
             };
 
             await GenerateImageImpl(overrides, cancellationToken);
+            await ImageFolderCardViewModel.OnLoadedAsync();
         }
         catch (OperationCanceledException e)
         {

StabilityMatrix.Avalonia/Views/Inference/InferenceTextToImageView.axaml

@@ -65,7 +65,7 @@
                         x:Name="ConfigPane"
                         Alignment="Left"
                         Id="ConfigPane"
-                        Proportion="0.25">
+                        Proportion="0.2">
                         <Tool
                             x:Name="ConfigTool"
                             Title="Config"
@@ -85,7 +85,7 @@
                         x:Name="PromptPane"
                         Alignment="Right"
                         Id="PromptPane"
-                        Proportion="0.25">
+                        Proportion="0.3">
                         <Tool
                             x:Name="PromptTool"
                             Title="Prompt"
@@ -167,14 +167,14 @@
                         </Tool>
                     </ToolDock>
 
-                    <ProportionalDockSplitter x:Name="RightSplitter" Id="RightSplitter" />
+                    <ProportionalDockSplitter x:Name="MiddleRightSplitter" Id="MiddleRightSplitter" />
 
-                    <!--  Right Pane  -->
+                    <!--  Middle Right Pane  -->
                     <ToolDock
                         x:Name="ImageGalleryPane"
                         Alignment="Right"
                         Id="ImageGalleryPane"
-                        Proportion="0.35">
+                        Proportion="0.3">
                         <Tool
                             x:Name="ImageGalleryTool"
                             Title="Image Output"
@@ -208,6 +208,29 @@
                             </Grid>
                         </Tool>
                     </ToolDock>
+                    
+                    <ProportionalDockSplitter x:Name="RightSplitter" Id="RightSplitter" />
+                    
+                    <!--  Right Pane  -->
+                    <ToolDock
+                        x:Name="ImageFolderPane"
+                        Alignment="Right"
+                        Id="ImageFolderPane"
+                        Proportion="0.2">
+                        <Tool
+                            x:Name="ImageFolderTool"
+                            x:DataType="Tool"
+                            Id="ImageFolderTool"
+                            CanClose="False">
+                            
+                            <Grid
+                                x:CompileBindings="False"
+                                DataContext="{Binding ElementName=Dock, Path=DataContext}">
+                                <controls:ImageFolderCard
+                                    DataContext="{Binding ImageFolderCardViewModel}" />
+                            </Grid>
+                        </Tool>
+                    </ToolDock>
 
                 </ProportionalDock>
             </RootDock>

StabilityMatrix.Core/Database/ILiteDbContext.cs

@@ -13,6 +13,7 @@ public interface ILiteDbContext : IDisposable
     ILiteCollectionAsync<CivitModelQueryCacheEntry> CivitModelQueryCache { get; }
     ILiteCollectionAsync<LocalModelFile> LocalModelFiles { get; }
     ILiteCollectionAsync<InferenceProjectEntry> InferenceProjects { get; }
+    ILiteCollectionAsync<LocalImageFile> LocalImageFiles { get; }
 
     Task<(CivitModel?, CivitModelVersion?)> FindCivitModelFromFileHashAsync(string hashBlake3);
     Task<bool> UpsertCivitModelAsync(CivitModel civitModel);

StabilityMatrix.Core/Database/LiteDbContext.cs

@@ -35,6 +35,8 @@ public class LiteDbContext : ILiteDbContext
         Database.GetCollection<LocalModelFile>("LocalModelFiles");
     public ILiteCollectionAsync<InferenceProjectEntry> InferenceProjects =>
         Database.GetCollection<InferenceProjectEntry>("InferenceProjects");
+    public ILiteCollectionAsync<LocalImageFile> LocalImageFiles =>
+        Database.GetCollection<LocalImageFile>("LocalImageFiles");
 
     public LiteDbContext(
         ILogger<LiteDbContext> logger,

StabilityMatrix.Core/Helper/SharedFolders.cs

@@ -1,4 +1,5 @@
-using NLog;
+using System.Diagnostics.CodeAnalysis;
+using NLog;
 using StabilityMatrix.Core.Extensions;
 using StabilityMatrix.Core.Helper.Factory;
 using StabilityMatrix.Core.Models;
@@ -20,9 +21,9 @@ public class SharedFolders : ISharedFolders
         this.settingsManager = settingsManager;
         this.packageFactory = packageFactory;
     }
-    
+
     // Platform redirect for junctions / symlinks
-    private static void CreateLinkOrJunction(string junctionDir, string targetDir, bool overwrite)
+    public static void CreateLinkOrJunction(string junctionDir, string targetDir, bool overwrite)
     {
         if (Compat.IsWindows)
         {
@@ -36,8 +37,54 @@ public class SharedFolders : ISharedFolders
         }
     }
 
-    public static void SetupLinks(Dictionary<SharedFolderType, IReadOnlyList<string>> definitions, 
-        DirectoryPath modelsDirectory, DirectoryPath installDirectory)
+    /// <summary>
+    /// Creates a junction link from the source to the destination.
+    /// Moves destination files to source if they exist.
+    /// </summary>
+    /// <param name="sourceDir">Shared source (i.e. "Models/")</param>
+    /// <param name="destinationDir">Destination (i.e. "webui/models/lora")</param>
+    public static void CreateLinkOrJunctionWithMove(
+        DirectoryPath sourceDir,
+        DirectoryPath destinationDir
+    )
+    {
+        // Create source folder if it doesn't exist
+        if (!sourceDir.Exists)
+        {
+            Logger.Info($"Creating junction source {sourceDir}");
+            sourceDir.Create();
+        }
+        // Delete the destination folder if it exists
+        if (destinationDir.Exists)
+        {
+            // Copy all files from destination to source
+            Logger.Info($"Copying files from {destinationDir} to {sourceDir}");
+            foreach (var file in destinationDir.Info.EnumerateFiles())
+            {
+                var sourceFile = sourceDir + file;
+                var destinationFile = destinationDir + file;
+                // Skip name collisions
+                if (File.Exists(sourceFile))
+                {
+                    Logger.Warn(
+                        $"Skipping file {file.FullName} because it already exists in {sourceDir}"
+                    );
+                    continue;
+                }
+                destinationFile.Info.MoveTo(sourceFile);
+            }
+            Logger.Info($"Deleting junction target {destinationDir}");
+            destinationDir.Delete(true);
+        }
+        Logger.Info($"Creating junction link from {sourceDir} to {destinationDir}");
+        CreateLinkOrJunction(destinationDir, sourceDir, true);
+    }
+
+    public static void SetupLinks(
+        Dictionary<SharedFolderType, IReadOnlyList<string>> definitions,
+        DirectoryPath modelsDirectory,
+        DirectoryPath installDirectory
+    )
     {
         foreach (var (folderType, relativePaths) in definitions)
         {
@@ -63,7 +110,9 @@ public class SharedFolders : ISharedFolders
                         // Skip name collisions
                         if (File.Exists(sourceFile))
                         {
-                            Logger.Warn($"Skipping file {file.FullName} because it already exists in {sourceDir}");
+                            Logger.Warn(
+                                $"Skipping file {file.FullName} because it already exists in {sourceDir}"
+                            );
                             continue;
                         }
                         destinationFile.Info.MoveTo(sourceFile);
@@ -81,19 +130,25 @@ public class SharedFolders : ISharedFolders
     {
         var modelsDirectory = new DirectoryPath(settingsManager.ModelsDirectory);
         var sharedFolders = basePackage.SharedFolders;
-        if (sharedFolders == null) return;
+        if (sharedFolders == null)
+            return;
         SetupLinks(sharedFolders, modelsDirectory, installDirectory);
     }
-    
+
     /// <summary>
-    /// Deletes junction links and remakes them. Unlike SetupLinksForPackage, 
+    /// Deletes junction links and remakes them. Unlike SetupLinksForPackage,
     /// this will not copy files from the destination to the source.
     /// </summary>
-    public static async Task UpdateLinksForPackage(BasePackage basePackage, DirectoryPath modelsDirectory, DirectoryPath installDirectory)
+    public static async Task UpdateLinksForPackage(
+        BasePackage basePackage,
+        DirectoryPath modelsDirectory,
+        DirectoryPath installDirectory
+    )
     {
         var sharedFolders = basePackage.SharedFolders;
-        if (sharedFolders is null) return;
-        
+        if (sharedFolders is null)
+            return;
+
         foreach (var (folderType, relativePaths) in sharedFolders)
         {
             foreach (var relativePath in relativePaths)
@@ -117,7 +172,8 @@ public class SharedFolders : ISharedFolders
                         if (destinationDir.Info.LinkTarget == sourceDir)
                         {
                             Logger.Info(
-                                $"Skipped updating matching folder link ({destinationDir} -> ({sourceDir})");
+                                $"Skipped updating matching folder link ({destinationDir} -> ({sourceDir})"
+                            );
                             return;
                         }
 
@@ -131,8 +187,12 @@ public class SharedFolders : ISharedFolders
                         if (destinationDir.Info.EnumerateFileSystemInfos().Any())
                         {
                             Logger.Info($"Moving files from {destinationDir} to {sourceDir}");
-                            await FileTransfers.MoveAllFilesAndDirectories(
-                                    destinationDir, sourceDir, overwriteIfHashMatches: true)
+                            await FileTransfers
+                                .MoveAllFilesAndDirectories(
+                                    destinationDir,
+                                    sourceDir,
+                                    overwriteIfHashMatches: true
+                                )
                                 .ConfigureAwait(false);
                         }
 
@@ -154,15 +214,16 @@ public class SharedFolders : ISharedFolders
         {
             return;
         }
-        
+
         foreach (var (_, relativePaths) in sharedFolders)
         {
             foreach (var relativePath in relativePaths)
             {
                 var destination = Path.GetFullPath(Path.Combine(installPath, relativePath));
                 // Delete the destination folder if it exists
-                if (!Directory.Exists(destination)) continue;
-            
+                if (!Directory.Exists(destination))
+                    continue;
+
                 Logger.Info($"Deleting junction target {destination}");
                 Directory.Delete(destination, false);
             }
@@ -174,19 +235,26 @@ public class SharedFolders : ISharedFolders
         var packages = settingsManager.Settings.InstalledPackages;
         foreach (var package in packages)
         {
-            if (package.PackageName == null) continue;
+            if (package.PackageName == null)
+                continue;
             var basePackage = packageFactory[package.PackageName];
-            if (basePackage == null) continue;
-            if (package.LibraryPath == null) continue;
- 
+            if (basePackage == null)
+                continue;
+            if (package.LibraryPath == null)
+                continue;
+
             try
             {
                 basePackage.RemoveModelFolderLinks(package.FullPath).GetAwaiter().GetResult();
             }
             catch (Exception e)
             {
-                Logger.Warn("Failed to remove links for package {Package} " +
-                            "({DisplayName}): {Message}", package.PackageName, package.DisplayName, e.Message);
+                Logger.Warn(
+                    "Failed to remove links for package {Package} " + "({DisplayName}): {Message}",
+                    package.PackageName,
+                    package.DisplayName,
+                    e.Message
+                );
             }
         }
     }
@@ -194,8 +262,9 @@ public class SharedFolders : ISharedFolders
     public void SetupSharedModelFolders()
     {
         var modelsDir = settingsManager.ModelsDirectory;
-        if (string.IsNullOrWhiteSpace(modelsDir)) return;
-        
+        if (string.IsNullOrWhiteSpace(modelsDir))
+            return;
+
         Directory.CreateDirectory(modelsDir);
         var allSharedFolderTypes = Enum.GetValues<SharedFolderType>();
         foreach (var sharedFolder in allSharedFolderTypes)

StabilityMatrix.Core/Models/Database/LocalImageFile.cs

@@ -0,0 +1,48 @@
+using LiteDB;
+
+namespace StabilityMatrix.Core.Models.Database;
+
+/// <summary>
+/// Represents a locally indexed image file.
+/// </summary>
+public class LocalImageFile
+{
+    /// <summary>
+    /// Relative path of the file from the root images directory ("%LIBRARY%/Images").
+    /// </summary>
+    [BsonId]
+    public required string RelativePath { get; set; }
+
+    /// <summary>
+    /// Type of the model file.
+    /// </summary>
+    public LocalImageFileType ImageType { get; set; }
+
+    /// <summary>
+    /// Creation time of the file.
+    /// </summary>
+    public DateTimeOffset CreatedAt { get; set; }
+
+    /// <summary>
+    /// Last modified time of the file.
+    /// </summary>
+    public DateTimeOffset LastModifiedAt { get; set; }
+
+    /// <summary>
+    /// File name of the relative path.
+    /// </summary>
+    public string FileName => Path.GetFileName(RelativePath);
+
+    /// <summary>
+    /// File name of the relative path without extension.
+    /// </summary>
+    public string FileNameWithoutExtension => Path.GetFileNameWithoutExtension(RelativePath);
+
+    public string GetFullPath(string rootImageDirectory)
+    {
+        return Path.Combine(rootImageDirectory, RelativePath);
+    }
+
+    public static readonly HashSet<string> SupportedImageExtensions =
+        new() { ".png", ".jpg", ".jpeg", ".webp" };
+}

StabilityMatrix.Core/Models/Database/LocalImageFileType.cs

@@ -0,0 +1,14 @@
+namespace StabilityMatrix.Core.Models.Database;
+
+[Flags]
+public enum LocalImageFileType : ulong
+{
+    // Source
+    Automatic = 1 << 1,
+    Comfy = 1 << 2,
+    Inference = 1 << 3,
+
+    // Generation Type
+    TextToImage = 1 << 10,
+    ImageToImage = 1 << 11
+}

StabilityMatrix.Core/Models/Packages/ComfyUI.cs

@@ -20,7 +20,7 @@ public class ComfyUI : BaseGitPackage
     public override string DisplayName { get; set; } = "ComfyUI";
     public override string Author => "comfyanonymous";
     public override string LicenseType => "GPL-3.0";
-    public override string LicenseUrl => 
+    public override string LicenseUrl =>
         "https://github.com/comfyanonymous/ComfyUI/blob/master/LICENSE";
     public override string Blurb => "A powerful and modular stable diffusion GUI and backend";
     public override string LaunchCommand => "main.py";
@@ -29,81 +29,88 @@ public class ComfyUI : BaseGitPackage
         new("https://github.com/comfyanonymous/ComfyUI/raw/master/comfyui_screenshot.png");
     public override bool ShouldIgnoreReleases => true;
 
-    public ComfyUI(IGithubApiCache githubApi, ISettingsManager settingsManager, IDownloadService downloadService,
-        IPrerequisiteHelper prerequisiteHelper) :
-        base(githubApi, settingsManager, downloadService, prerequisiteHelper)
-    {
-    }
+    public ComfyUI(
+        IGithubApiCache githubApi,
+        ISettingsManager settingsManager,
+        IDownloadService downloadService,
+        IPrerequisiteHelper prerequisiteHelper
+    )
+        : base(githubApi, settingsManager, downloadService, prerequisiteHelper) { }
 
     // https://github.com/comfyanonymous/ComfyUI/blob/master/folder_paths.py#L11
-    public override Dictionary<SharedFolderType, IReadOnlyList<string>> SharedFolders => new()
-    {
-        [SharedFolderType.StableDiffusion] = new[] {"models/checkpoints"},
-        [SharedFolderType.Diffusers] = new[] {"models/diffusers"},
-        [SharedFolderType.Lora] = new[] {"models/loras"},
-        [SharedFolderType.CLIP] = new[] {"models/clip"},
-        [SharedFolderType.TextualInversion] = new[] {"models/embeddings"},
-        [SharedFolderType.VAE] = new[] {"models/vae"},
-        [SharedFolderType.ApproxVAE] = new[] {"models/vae_approx"},
-        [SharedFolderType.ControlNet] = new[] {"models/controlnet"},
-        [SharedFolderType.GLIGEN] = new[] {"models/gligen"},
-        [SharedFolderType.ESRGAN] = new[] {"models/upscale_models"},
-        [SharedFolderType.Hypernetwork] = new[] {"models/hypernetworks"},
-    };
-    
-    public override List<LaunchOptionDefinition> LaunchOptions => new List<LaunchOptionDefinition>
-    {
+    public override Dictionary<SharedFolderType, IReadOnlyList<string>> SharedFolders =>
         new()
         {
-            Name = "VRAM",
-            Type = LaunchOptionType.Bool,
-            InitialValue = HardwareHelper.IterGpuInfo().Select(gpu => gpu.MemoryLevel).Max() switch
+            [SharedFolderType.StableDiffusion] = new[] { "models/checkpoints" },
+            [SharedFolderType.Diffusers] = new[] { "models/diffusers" },
+            [SharedFolderType.Lora] = new[] { "models/loras" },
+            [SharedFolderType.CLIP] = new[] { "models/clip" },
+            [SharedFolderType.TextualInversion] = new[] { "models/embeddings" },
+            [SharedFolderType.VAE] = new[] { "models/vae" },
+            [SharedFolderType.ApproxVAE] = new[] { "models/vae_approx" },
+            [SharedFolderType.ControlNet] = new[] { "models/controlnet" },
+            [SharedFolderType.GLIGEN] = new[] { "models/gligen" },
+            [SharedFolderType.ESRGAN] = new[] { "models/upscale_models" },
+            [SharedFolderType.Hypernetwork] = new[] { "models/hypernetworks" },
+        };
+
+    public override List<LaunchOptionDefinition> LaunchOptions =>
+        new List<LaunchOptionDefinition>
+        {
+            new()
             {
-                Level.Low => "--lowvram",
-                Level.Medium => "--normalvram",
-                _ => null
+                Name = "VRAM",
+                Type = LaunchOptionType.Bool,
+                InitialValue = HardwareHelper
+                    .IterGpuInfo()
+                    .Select(gpu => gpu.MemoryLevel)
+                    .Max() switch
+                {
+                    Level.Low => "--lowvram",
+                    Level.Medium => "--normalvram",
+                    _ => null
+                },
+                Options = { "--highvram", "--normalvram", "--lowvram", "--novram" }
             },
-            Options = { "--highvram", "--normalvram", "--lowvram", "--novram" }
-        },
-        new()
-        {
-            Name = "Use CPU only",
-            Type = LaunchOptionType.Bool,
-            InitialValue = !HardwareHelper.HasNvidiaGpu(),
-            Options = {"--cpu"}
-        },
-        new()
-        {
-            Name = "Disable Xformers",
-            Type = LaunchOptionType.Bool,
-            InitialValue = !HardwareHelper.HasNvidiaGpu(),
-            Options = { "--disable-xformers" }
-        },
-        new()
-        {
-            Name = "Auto-Launch",
-            Type = LaunchOptionType.Bool,
-            Options = { "--auto-launch" }
-        },
-        LaunchOptionDefinition.Extras
-    };
+            new()
+            {
+                Name = "Use CPU only",
+                Type = LaunchOptionType.Bool,
+                InitialValue = !HardwareHelper.HasNvidiaGpu(),
+                Options = { "--cpu" }
+            },
+            new()
+            {
+                Name = "Disable Xformers",
+                Type = LaunchOptionType.Bool,
+                InitialValue = !HardwareHelper.HasNvidiaGpu(),
+                Options = { "--disable-xformers" }
+            },
+            new()
+            {
+                Name = "Auto-Launch",
+                Type = LaunchOptionType.Bool,
+                Options = { "--auto-launch" }
+            },
+            LaunchOptionDefinition.Extras
+        };
 
     public override Task<string> GetLatestVersion() => Task.FromResult("master");
 
-    public override async Task<IEnumerable<PackageVersion>> GetAllVersions(bool isReleaseMode = true)
+    public override async Task<IEnumerable<PackageVersion>> GetAllVersions(
+        bool isReleaseMode = true
+    )
     {
         var allBranches = await GetAllBranches().ConfigureAwait(false);
-        return allBranches.Select(b => new PackageVersion
-        {
-            TagName = $"{b.Name}", 
-            ReleaseNotesMarkdown = string.Empty
-        });
+        return allBranches.Select(
+            b => new PackageVersion { TagName = $"{b.Name}", ReleaseNotesMarkdown = string.Empty }
+        );
     }
 
     public override async Task InstallPackage(IProgress<ProgressReport>? progress = null)
     {
         await UnzipPackage(progress);
-        
+
         progress?.Report(new ProgressReport(-1, "Setting up venv", isIndeterminate: true));
         // Setup venv
         await using var venvRunner = new PyVenvRunner(Path.Combine(InstallLocation, "venv"));
@@ -114,46 +121,62 @@ public class ComfyUI : BaseGitPackage
         var gpus = HardwareHelper.IterGpuInfo().ToList();
         if (gpus.Any(g => g.IsNvidia))
         {
-            progress?.Report(new ProgressReport(-1, "Installing PyTorch for CUDA", isIndeterminate: true));
-            
+            progress?.Report(
+                new ProgressReport(-1, "Installing PyTorch for CUDA", isIndeterminate: true)
+            );
+
             Logger.Info("Starting torch install (CUDA)...");
-            await venvRunner.PipInstall(PyVenvRunner.TorchPipInstallArgsCuda, OnConsoleOutput)
+            await venvRunner
+                .PipInstall(PyVenvRunner.TorchPipInstallArgsCuda, OnConsoleOutput)
                 .ConfigureAwait(false);
-            
+
             Logger.Info("Installing xformers...");
             await venvRunner.PipInstall("xformers", OnConsoleOutput).ConfigureAwait(false);
         }
         else if (HardwareHelper.PreferRocm())
         {
-            progress?.Report(new ProgressReport(-1, "Installing PyTorch for ROCm", isIndeterminate: true));
+            progress?.Report(
+                new ProgressReport(-1, "Installing PyTorch for ROCm", isIndeterminate: true)
+            );
             await venvRunner
                 .PipInstall(PyVenvRunner.TorchPipInstallArgsRocm542, OnConsoleOutput)
                 .ConfigureAwait(false);
         }
         else
         {
-            progress?.Report(new ProgressReport(-1, "Installing PyTorch for CPU", isIndeterminate: true));
+            progress?.Report(
+                new ProgressReport(-1, "Installing PyTorch for CPU", isIndeterminate: true)
+            );
             Logger.Info("Starting torch install (CPU)...");
-            await venvRunner.PipInstall(PyVenvRunner.TorchPipInstallArgsCpu, OnConsoleOutput)
+            await venvRunner
+                .PipInstall(PyVenvRunner.TorchPipInstallArgsCpu, OnConsoleOutput)
                 .ConfigureAwait(false);
         }
 
         // Install requirements file
-        progress?.Report(new ProgressReport(-1, "Installing Package Requirements", isIndeterminate: true));
+        progress?.Report(
+            new ProgressReport(-1, "Installing Package Requirements", isIndeterminate: true)
+        );
         Logger.Info("Installing requirements.txt");
         await venvRunner.PipInstall($"-r requirements.txt", OnConsoleOutput).ConfigureAwait(false);
-        
-        progress?.Report(new ProgressReport(1, "Installing Package Requirements", isIndeterminate: false));
+
+        progress?.Report(
+            new ProgressReport(1, "Installing Package Requirements", isIndeterminate: false)
+        );
     }
-    
-    public override async Task RunPackage(string installedPackagePath, string command, string arguments)
+
+    public override async Task RunPackage(
+        string installedPackagePath,
+        string command,
+        string arguments
+    )
     {
         await SetupVenv(installedPackagePath).ConfigureAwait(false);
 
         void HandleConsoleOutput(ProcessOutput s)
         {
             OnConsoleOutput(s);
-            
+
             if (s.Text.Contains("To see the GUI go to", StringComparison.OrdinalIgnoreCase))
             {
                 var regex = new Regex(@"(https?:\/\/)([^:\s]+):(\d+)");
@@ -174,10 +197,7 @@ public class ComfyUI : BaseGitPackage
 
         var args = $"\"{Path.Combine(installedPackagePath, command)}\" {arguments}";
 
-        VenvRunner?.RunDetached(
-            args.TrimEnd(), 
-            HandleConsoleOutput, 
-            HandleExit);
+        VenvRunner?.RunDetached(args.TrimEnd(), HandleConsoleOutput, HandleExit);
     }
 
     public override Task SetupModelFolders(DirectoryPath installDirectory)
@@ -197,19 +217,22 @@ public class ComfyUI : BaseGitPackage
             File.WriteAllText(extraPathsYamlPath, string.Empty);
         }
         var yaml = File.ReadAllText(extraPathsYamlPath);
-        var comfyModelPaths = deserializer.Deserialize<ComfyModelPathsYaml>(yaml) ??
-                              // ReSharper disable once NullCoalescingConditionIsAlwaysNotNullAccordingToAPIContract
-                              // cuz it can actually be null lol 
-                              new ComfyModelPathsYaml();
-        
+        var comfyModelPaths =
+            deserializer.Deserialize<ComfyModelPathsYaml>(yaml)
+            ??
+            // ReSharper disable once NullCoalescingConditionIsAlwaysNotNullAccordingToAPIContract
+            // cuz it can actually be null lol
+            new ComfyModelPathsYaml();
+
         comfyModelPaths.StabilityMatrix ??= new ComfyModelPathsYaml.SmData();
         comfyModelPaths.StabilityMatrix.Checkpoints = Path.Combine(modelsDir, "StableDiffusion");
         comfyModelPaths.StabilityMatrix.Vae = Path.Combine(modelsDir, "VAE");
-        comfyModelPaths.StabilityMatrix.Loras = $"{Path.Combine(modelsDir, "Lora")}\n" +
-                                       $"{Path.Combine(modelsDir, "LyCORIS")}";
-        comfyModelPaths.StabilityMatrix.UpscaleModels = $"{Path.Combine(modelsDir, "ESRGAN")}\n" +
-                                              $"{Path.Combine(modelsDir, "RealESRGAN")}\n" +
-                                              $"{Path.Combine(modelsDir, "SwinIR")}";
+        comfyModelPaths.StabilityMatrix.Loras =
+            $"{Path.Combine(modelsDir, "Lora")}\n" + $"{Path.Combine(modelsDir, "LyCORIS")}";
+        comfyModelPaths.StabilityMatrix.UpscaleModels =
+            $"{Path.Combine(modelsDir, "ESRGAN")}\n"
+            + $"{Path.Combine(modelsDir, "RealESRGAN")}\n"
+            + $"{Path.Combine(modelsDir, "SwinIR")}";
         comfyModelPaths.StabilityMatrix.Embeddings = Path.Combine(modelsDir, "TextualInversion");
         comfyModelPaths.StabilityMatrix.Hypernetworks = Path.Combine(modelsDir, "Hypernetwork");
         comfyModelPaths.StabilityMatrix.Controlnet = Path.Combine(modelsDir, "ControlNet");
@@ -217,7 +240,7 @@ public class ComfyUI : BaseGitPackage
         comfyModelPaths.StabilityMatrix.Diffusers = Path.Combine(modelsDir, "Diffusers");
         comfyModelPaths.StabilityMatrix.Gligen = Path.Combine(modelsDir, "GLIGEN");
         comfyModelPaths.StabilityMatrix.VaeApprox = Path.Combine(modelsDir, "ApproxVAE");
-        
+
         var serializer = new SerializerBuilder()
             .WithNamingConvention(UnderscoredNamingConvention.Instance)
             .Build();
@@ -233,6 +256,19 @@ public class ComfyUI : BaseGitPackage
     public override Task RemoveModelFolderLinks(DirectoryPath installDirectory) =>
         Task.CompletedTask;
 
+    public async Task SetupInferenceOutputFolderLinks(DirectoryPath installDirectory)
+    {
+        var inferenceDir = installDirectory.JoinDir("output", "Inference");
+
+        var sharedInferenceDir = SettingsManager.ImagesInferenceDirectory;
+
+        await Task.Run(() =>
+            {
+                Helper.SharedFolders.CreateLinkOrJunctionWithMove(sharedInferenceDir, inferenceDir);
+            })
+            .ConfigureAwait(false);
+    }
+
     public class ComfyModelPathsYaml
     {
         public class SmData

StabilityMatrix.Core/Services/IImageIndexService.cs

@@ -0,0 +1,21 @@
+using StabilityMatrix.Core.Models.Database;
+
+namespace StabilityMatrix.Core.Services;
+
+public interface IImageIndexService
+{
+    /// <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
+    /// </summary>
+    Task RefreshIndex(string subPath = "");
+
+    /// <summary>
+    /// Refreshes the index of local images in the background
+    /// </summary>
+    void BackgroundRefreshIndex();
+}

StabilityMatrix.Core/Services/ImageIndexService.cs

@@ -0,0 +1,111 @@
+using System.Diagnostics;
+using AsyncAwaitBestPractices;
+using Microsoft.Extensions.Logging;
+using StabilityMatrix.Core.Database;
+using StabilityMatrix.Core.Models.Database;
+using StabilityMatrix.Core.Models.FileInterfaces;
+
+namespace StabilityMatrix.Core.Services;
+
+public class ImageIndexService : IImageIndexService
+{
+    private readonly ILogger<ImageIndexService> logger;
+    private readonly ILiteDbContext liteDbContext;
+    private readonly ISettingsManager settingsManager;
+
+    public ImageIndexService(
+        ILogger<ImageIndexService> logger,
+        ILiteDbContext liteDbContext,
+        ISettingsManager settingsManager
+    )
+    {
+        this.logger = logger;
+        this.liteDbContext = liteDbContext;
+        this.settingsManager = settingsManager;
+    }
+
+    /// <inheritdoc />
+    public async Task<IReadOnlyList<LocalImageFile>> GetLocalImagesByPrefix(string pathPrefix)
+    {
+        return await liteDbContext.LocalImageFiles
+            .Query()
+            .Where(imageFile => imageFile.RelativePath.StartsWith(pathPrefix))
+            .ToArrayAsync()
+            .ConfigureAwait(false);
+    }
+
+    /// <inheritdoc />
+    public async Task RefreshIndex(string subPath = "")
+    {
+        var imagesDir = settingsManager.ImagesDirectory;
+
+        // Start
+        var stopwatch = Stopwatch.StartNew();
+        logger.LogInformation("Refreshing images index...");
+
+        using var db = await liteDbContext.Database.BeginTransactionAsync().ConfigureAwait(false);
+
+        var localImageFiles = db.GetCollection<LocalImageFile>("LocalImageFiles")!;
+
+        await localImageFiles.DeleteAllAsync().ConfigureAwait(false);
+
+        // Record start of actual indexing
+        var indexStart = stopwatch.Elapsed;
+
+        var added = 0;
+
+        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);
+
+            // Skip if not in sub-path
+            if (!string.IsNullOrEmpty(subPath) && !relativePath.StartsWith(subPath))
+            {
+                continue;
+            }
+
+            // TODO: Support other types
+            const LocalImageFileType imageType =
+                LocalImageFileType.Inference | LocalImageFileType.TextToImage;
+
+            var localImage = new LocalImageFile
+            {
+                RelativePath = relativePath,
+                ImageType = imageType
+            };
+
+            // Insert into database
+            await localImageFiles.InsertAsync(localImage).ConfigureAwait(false);
+
+            added++;
+        }
+        // Record end of actual indexing
+        var indexEnd = stopwatch.Elapsed;
+
+        await db.CommitAsync().ConfigureAwait(false);
+
+        // End
+        stopwatch.Stop();
+        var indexDuration = indexEnd - indexStart;
+        var dbDuration = stopwatch.Elapsed - indexDuration;
+
+        logger.LogInformation(
+            "Image index updated for {Prefix} with {Entries} files, took {IndexDuration:F1}ms ({DbDuration:F1}ms db)",
+            subPath,
+            added,
+            indexDuration.TotalMilliseconds,
+            dbDuration.TotalMilliseconds
+        );
+    }
+
+    /// <inheritdoc />
+    public void BackgroundRefreshIndex()
+    {
+        RefreshIndex().SafeFireAndForget();
+    }
+}

StabilityMatrix.Core/StabilityMatrix.Core.csproj

@@ -21,6 +21,7 @@
       <PackageReference Include="DeviceId.Mac" Version="6.2.0" />
       <PackageReference Include="DeviceId.Windows" Version="6.2.0" />
       <PackageReference Include="DeviceId.Windows.Wmi" Version="6.2.1" />
+      <PackageReference Include="DynamicData" Version="7.14.2" />
       <PackageReference Include="LiteDB" Version="5.0.16" />
       <PackageReference Include="LiteDB.Async" Version="0.1.6" />
       <PackageReference Include="MetadataExtractor" Version="2.8.1" />