Browse Source

Add faster memory based image indexing

pull/165/head
Ionite 1 year ago
parent
commit
d12e1c7809
No known key found for this signature in database
  1. 33
      StabilityMatrix.Avalonia/Controls/ImageFolderCard.axaml
  2. 20
      StabilityMatrix.Avalonia/DesignData/MockImageIndexService.cs
  3. 29
      StabilityMatrix.Avalonia/ViewModels/Inference/ImageFolderCardViewModel.cs
  4. 60
      StabilityMatrix.Avalonia/ViewModels/Inference/InferenceTextToImageViewModel.cs
  5. 18
      StabilityMatrix.Avalonia/ViewModels/Inference/PromptCardViewModel.cs
  6. 2
      StabilityMatrix.Avalonia/ViewModels/InferenceViewModel.cs
  7. 80
      StabilityMatrix.Core/Models/Database/LocalImageFile.cs
  8. 8
      StabilityMatrix.Core/Models/GenerationParameters.cs
  9. 59
      StabilityMatrix.Core/Models/IndexCollection.cs
  10. 15
      StabilityMatrix.Core/Services/IImageIndexService.cs
  11. 104
      StabilityMatrix.Core/Services/ImageIndexService.cs

33
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}"

20
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()
{

29
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");
var imageFiles = await imageIndexService.GetLocalImagesByPrefix("Inference");
localImagesSource.Edit(x =>
{
x.Load(imageFiles);
});
imageIndexService.RefreshIndexForAllCollections().SafeFireAndForget();
}
/// <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)

60
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,7 +421,23 @@ 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;
promptTask?.Dispose();
client.PreviewImageReceived -= OnPreviewImageReceived;
}
}
private async Task ProcessOutputs(IReadOnlyList<ComfyImage> images)
{
List<ImageSource> outputImages;
// Use local file path if available, otherwise use remote URL
if (client.OutputImagesDir is { } outputPath)
@ -425,21 +446,30 @@ public partial class InferenceTextToImageViewModel : InferenceTabViewModelBase
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,
await filePath.ReadAllBytesAsync(),
generationInfo,
smproj
);
fileStream.Close();
fileStream.Dispose();
await using var outputStream = filePath.Info.OpenWrite();
/*await using (var readStream = filePath.Info.OpenWrite())
{
using (var reader = new BinaryReader(readStream))
{
}
}*/
await using (var outputStream = filePath.Info.OpenWrite())
{
await outputStream.WriteAsync(bytesWithMetadata);
await outputStream.FlushAsync();
}
outputImages.Add(new ImageSource(filePath));
imageIndexService.OnImageAdded(filePath);
}
}
else
@ -478,6 +508,8 @@ public partial class InferenceTextToImageViewModel : InferenceTabViewModelBase
// Preload
await gridImage.GetBitmapAsync();
ImageGalleryCardViewModel.ImageSources.Add(gridImage);
imageIndexService.OnImageAdded(gridPath);
}
// Add rest of images
@ -488,19 +520,6 @@ public partial class InferenceTextToImageViewModel : InferenceTabViewModelBase
ImageGalleryCardViewModel.ImageSources.Add(img);
}
}
finally
{
// Disconnect progress handler
OutputProgress.Value = 0;
OutputProgress.Text = "";
ImageGalleryCardViewModel.PreviewImage?.Dispose();
ImageGalleryCardViewModel.PreviewImage = null;
ImageGalleryCardViewModel.IsPreviewOverlayEnabled = false;
promptTask?.Dispose();
client.PreviewImageReceived -= OnPreviewImageReceived;
}
}
[RelayCommand(IncludeCancelCommand = true, FlowExceptionsToTaskScheduler = true)]
private async Task GenerateImage(
@ -517,7 +536,6 @@ public partial class InferenceTextToImageViewModel : InferenceTabViewModelBase
};
await GenerateImageImpl(overrides, cancellationToken);
await ImageFolderCardViewModel.OnLoadedAsync();
}
catch (OperationCanceledException e)
{

18
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

2
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>

80
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();
}

8
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)

59
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);
}
}

15
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

104
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 />

Loading…
Cancel
Save