Merge pull request #427 from ionite34/fast-png

Fast png

CHANGELOG.md

@@ -8,6 +8,8 @@ and this project adheres to [Semantic Versioning 2.0](https://semver.org/spec/v2
 ## v2.7.5
 ### Fixed
 - Fixed Python Packages manager crash when pip list returns warnings in json
+- Fixed slowdown when loading PNGs with large amounts of metadata
+- Fixed crash when scanning directories for missing metadata
 
 ## v2.7.4
 ### Changed

StabilityMatrix.Avalonia/FallbackRamCachedWebImageLoader.cs

@@ -7,6 +7,7 @@ using AsyncAwaitBestPractices;
 using AsyncImageLoader.Loaders;
 using Avalonia.Media.Imaging;
 using StabilityMatrix.Core.Extensions;
+using StabilityMatrix.Core.Helper;
 
 namespace StabilityMatrix.Avalonia;
 
@@ -42,7 +43,11 @@ public class FallbackRamCachedWebImageLoader : RamCachedWebImageLoader
         {
             try
             {
+                if (!url.EndsWith("png", StringComparison.OrdinalIgnoreCase))
                     return new Bitmap(url);
+
+                using var stream = ImageMetadata.BuildImageWithoutMetadata(url);
+                return stream == null ? new Bitmap(url) : new Bitmap(stream);
             }
             catch (Exception e)
             {

StabilityMatrix.Avalonia/ViewModels/OutputsPageViewModel.cs

@@ -54,11 +54,9 @@ public partial class OutputsPageViewModel : PageViewModelBase
     private readonly ILogger<OutputsPageViewModel> logger;
     public override string Title => Resources.Label_OutputsPageTitle;
 
-    public override IconSource IconSource =>
-        new SymbolIconSource { Symbol = Symbol.Grid, IsFilled = true };
+    public override IconSource IconSource => new SymbolIconSource { Symbol = Symbol.Grid, IsFilled = true };
 
-    public SourceCache<LocalImageFile, string> OutputsCache { get; } =
-        new(file => file.AbsolutePath);
+    public SourceCache<LocalImageFile, string> OutputsCache { get; } = new(file => file.AbsolutePath);
 
     public IObservableCollection<OutputImageViewModel> Outputs { get; set; } =
         new ObservableCollectionExtended<OutputImageViewModel>();
@@ -88,8 +86,7 @@ public partial class OutputsPageViewModel : PageViewModelBase
     [ObservableProperty]
     private bool isConsolidating;
 
-    public bool CanShowOutputTypes =>
-        SelectedCategory?.Name?.Equals("Shared Output Folder") ?? false;
+    public bool CanShowOutputTypes => SelectedCategory?.Name?.Equals("Shared Output Folder") ?? false;
 
     public string NumImagesSelected =>
         NumItemsSelected == 1
@@ -138,7 +135,7 @@ public partial class OutputsPageViewModel : PageViewModelBase
             delay: TimeSpan.FromMilliseconds(250)
         );
 
-        RefreshCategories();
+        RefreshCategories(false);
     }
 
     public override void OnLoaded()
@@ -163,10 +160,7 @@ public partial class OutputsPageViewModel : PageViewModelBase
         GetOutputs(path);
     }
 
-    partial void OnSelectedCategoryChanged(
-        PackageOutputCategory? oldValue,
-        PackageOutputCategory? newValue
-    )
+    partial void OnSelectedCategoryChanged(PackageOutputCategory? oldValue, PackageOutputCategory? newValue)
     {
         if (oldValue == newValue || newValue == null)
             return;
@@ -223,8 +217,8 @@ public partial class OutputsPageViewModel : PageViewModelBase
             )
             .Subscribe(ctx =>
             {
-                Dispatcher.UIThread
-                    .InvokeAsync(async () =>
+                Dispatcher
+                    .UIThread.InvokeAsync(async () =>
                     {
                         var sender = (ImageViewerViewModel)ctx.Sender!;
                         var newIndex = currentIndex + (ctx.EventArgs.IsNext ? 1 : -1);
@@ -261,7 +255,7 @@ public partial class OutputsPageViewModel : PageViewModelBase
 
     public void Refresh()
     {
-        Dispatcher.UIThread.Post(RefreshCategories);
+        Dispatcher.UIThread.Post(() => RefreshCategories());
         Dispatcher.UIThread.Post(OnLoaded);
     }
 
@@ -430,9 +424,7 @@ public partial class OutputsPageViewModel : PageViewModelBase
 
         Directory.CreateDirectory(settingsManager.ConsolidatedImagesDirectory);
 
-        foreach (
-            var category in stackPanel.Children.OfType<CheckBox>().Where(c => c.IsChecked == true)
-        )
+        foreach (var category in stackPanel.Children.OfType<CheckBox>().Where(c => c.IsChecked == true))
         {
             if (
                 string.IsNullOrWhiteSpace(category.Tag?.ToString())
@@ -442,13 +434,7 @@ public partial class OutputsPageViewModel : PageViewModelBase
 
             var directory = category.Tag.ToString();
 
-            foreach (
-                var path in Directory.EnumerateFiles(
-                    directory,
-                    "*.png",
-                    SearchOption.AllDirectories
-                )
-            )
+            foreach (var path in Directory.EnumerateFiles(directory, "*.png", SearchOption.AllDirectories))
             {
                 try
                 {
@@ -524,7 +510,7 @@ public partial class OutputsPageViewModel : PageViewModelBase
         }
     }
 
-    private void RefreshCategories()
+    private void RefreshCategories(bool updateProperty = true)
     {
         if (Design.IsDesignMode)
             return;
@@ -534,15 +520,11 @@ public partial class OutputsPageViewModel : PageViewModelBase
 
         var previouslySelectedCategory = SelectedCategory;
 
-        var packageCategories = settingsManager.Settings.InstalledPackages
-            .Where(x => !x.UseSharedOutputFolder)
+        var packageCategories = settingsManager
+            .Settings.InstalledPackages.Where(x => !x.UseSharedOutputFolder)
             .Select(packageFactory.GetPackagePair)
             .WhereNotNull()
-            .Where(
-                p =>
-                    p.BasePackage.SharedOutputFolders != null
-                    && p.BasePackage.SharedOutputFolders.Any()
-            )
+            .Where(p => p.BasePackage.SharedOutputFolders != null && p.BasePackage.SharedOutputFolders.Any())
             .Select(
                 pair =>
                     new PackageOutputCategory
@@ -567,16 +549,22 @@ public partial class OutputsPageViewModel : PageViewModelBase
 
         packageCategories.Insert(
             1,
-            new PackageOutputCategory
-            {
-                Path = settingsManager.ImagesInferenceDirectory,
-                Name = "Inference"
-            }
+            new PackageOutputCategory { Path = settingsManager.ImagesInferenceDirectory, Name = "Inference" }
         );
 
         Categories = new ObservableCollection<PackageOutputCategory>(packageCategories);
+
+        if (updateProperty)
+        {
             SelectedCategory =
                 Categories.FirstOrDefault(x => x.Name == previouslySelectedCategory?.Name)
                 ?? Categories.First();
         }
+        else
+        {
+            selectedCategory =
+                Categories.FirstOrDefault(x => x.Name == previouslySelectedCategory?.Name)
+                ?? Categories.First();
+        }
+    }
 }

StabilityMatrix.Core/Helper/ImageMetadata.cs

@@ -179,4 +179,52 @@ public class ImageMetadata
 
         return string.Empty;
     }
+
+    public static MemoryStream? BuildImageWithoutMetadata(FilePath imagePath)
+    {
+        using var byteStream = new BinaryReader(File.OpenRead(imagePath));
+        byteStream.BaseStream.Position = 0;
+
+        if (!byteStream.ReadBytes(8).SequenceEqual(PngHeader))
+        {
+            return null;
+        }
+
+        var memoryStream = new MemoryStream();
+        memoryStream.Write(PngHeader);
+
+        // add the IHDR chunk
+        var ihdrStuff = byteStream.ReadBytes(25);
+        memoryStream.Write(ihdrStuff);
+
+        // find IDATs
+        while (byteStream.BaseStream.Position < byteStream.BaseStream.Length - 4)
+        {
+            var chunkSizeBytes = byteStream.ReadBytes(4);
+            var chunkSize = BitConverter.ToInt32(chunkSizeBytes.Reverse().ToArray());
+            var chunkTypeBytes = byteStream.ReadBytes(4);
+            var chunkType = Encoding.UTF8.GetString(chunkTypeBytes);
+
+            if (chunkType != Encoding.UTF8.GetString(Idat))
+            {
+                // skip chunk data
+                byteStream.BaseStream.Position += chunkSize;
+                // skip crc
+                byteStream.BaseStream.Position += 4;
+                continue;
+            }
+
+            memoryStream.Write(chunkSizeBytes);
+            memoryStream.Write(chunkTypeBytes);
+            var idatBytes = byteStream.ReadBytes(chunkSize);
+            memoryStream.Write(idatBytes);
+            var crcBytes = byteStream.ReadBytes(4);
+            memoryStream.Write(crcBytes);
+        }
+
+        // Add IEND chunk
+        memoryStream.Write([0x00, 0x00, 0x00, 0x00, 0x49, 0x45, 0x4E, 0x44, 0xAE, 0x42, 0x60, 0x82]);
+        memoryStream.Position = 0;
+        return memoryStream;
+    }
 }

StabilityMatrix.Core/Services/MetadataImportService.cs

@@ -1,5 +1,4 @@
-using System.Diagnostics;
-using System.Text.Json;
+using System.Text.Json;
 using Microsoft.Extensions.Logging;
 using StabilityMatrix.Core.Attributes;
 using StabilityMatrix.Core.Helper;
@@ -18,7 +17,10 @@ public class MetadataImportService(
     ModelFinder modelFinder
 ) : IMetadataImportService
 {
-    public async Task ScanDirectoryForMissingInfo(DirectoryPath directory, IProgress<ProgressReport>? progress = null)
+    public async Task ScanDirectoryForMissingInfo(
+        DirectoryPath directory,
+        IProgress<ProgressReport>? progress = null
+    )
     {
         progress?.Report(new ProgressReport(-1f, "Scanning directory...", isIndeterminate: true));
 
@@ -54,7 +56,9 @@ public class MetadataImportService(
             }
 
             var fileNameWithoutExtension = checkpointFilePath.NameWithoutExtension;
-            var cmInfoPath = checkpointFilePath.Directory?.JoinFile($"{fileNameWithoutExtension}.cm-info.json");
+            var cmInfoPath = checkpointFilePath.Directory?.JoinFile(
+                $"{fileNameWithoutExtension}.cm-info.json"
+            );
             var cmInfoExists = File.Exists(cmInfoPath);
             if (cmInfoExists)
                 continue;
@@ -70,7 +74,10 @@ public class MetadataImportService(
                 );
             });
 
-            var blake3 = await GetBlake3Hash(cmInfoPath, checkpointFilePath, hashProgress).ConfigureAwait(false);
+            try
+            {
+                var blake3 = await GetBlake3Hash(cmInfoPath, checkpointFilePath, hashProgress)
+                    .ConfigureAwait(false);
                 if (string.IsNullOrWhiteSpace(blake3))
                 {
                     logger.LogWarning($"Blake3 hash was null for {checkpointFilePath}");
@@ -88,14 +95,19 @@ public class MetadataImportService(
 
                 var (model, modelVersion, modelFile) = modelInfo.Value;
 
-            var updatedCmInfo = new ConnectedModelInfo(model, modelVersion, modelFile, DateTimeOffset.UtcNow);
+                var updatedCmInfo = new ConnectedModelInfo(
+                    model,
+                    modelVersion,
+                    modelFile,
+                    DateTimeOffset.UtcNow
+                );
                 await updatedCmInfo
                     .SaveJsonToDirectory(checkpointFilePath.Directory, fileNameWithoutExtension)
                     .ConfigureAwait(false);
 
-            var image = modelVersion
-                .Images
-                ?.FirstOrDefault(img => LocalModelFile.SupportedImageExtensions.Contains(Path.GetExtension(img.Url)));
+                var image = modelVersion.Images?.FirstOrDefault(
+                    img => LocalModelFile.SupportedImageExtensions.Contains(Path.GetExtension(img.Url))
+                );
                 if (image == null)
                 {
                     scanned++;
@@ -108,6 +120,12 @@ public class MetadataImportService(
                 scanned++;
                 success++;
             }
+            catch (Exception e)
+            {
+                logger.LogError(e, "Error while scanning {checkpointFilePath}", checkpointFilePath);
+                scanned++;
+            }
+        }
 
         progress?.Report(
             new ProgressReport(
@@ -124,7 +142,10 @@ public class MetadataImportService(
             && !File.Exists(file.Directory?.JoinFile($"{file.NameWithoutExtension}.cm-info.json"));
     }
 
-    public async Task UpdateExistingMetadata(DirectoryPath directory, IProgress<ProgressReport>? progress = null)
+    public async Task UpdateExistingMetadata(
+        DirectoryPath directory,
+        IProgress<ProgressReport>? progress = null
+    )
     {
         progress?.Report(new ProgressReport(-1f, "Scanning directory...", isIndeterminate: true));
 
@@ -151,6 +172,8 @@ public class MetadataImportService(
                 )
             );
 
+            try
+            {
                 var hash = cmInfoValue.Hashes.BLAKE3;
                 if (string.IsNullOrWhiteSpace(hash))
                     continue;
@@ -164,14 +187,21 @@ public class MetadataImportService(
 
                 var (model, modelVersion, modelFile) = modelInfo.Value;
 
-            var updatedCmInfo = new ConnectedModelInfo(model, modelVersion, modelFile, DateTimeOffset.UtcNow);
+                var updatedCmInfo = new ConnectedModelInfo(
+                    model,
+                    modelVersion,
+                    modelFile,
+                    DateTimeOffset.UtcNow
+                );
 
                 var nameWithoutCmInfo = filePath.NameWithoutExtension.Replace(".cm-info", string.Empty);
-            await updatedCmInfo.SaveJsonToDirectory(filePath.Directory, nameWithoutCmInfo).ConfigureAwait(false);
+                await updatedCmInfo
+                    .SaveJsonToDirectory(filePath.Directory, nameWithoutCmInfo)
+                    .ConfigureAwait(false);
 
-            var image = modelVersion
-                .Images
-                ?.FirstOrDefault(img => LocalModelFile.SupportedImageExtensions.Contains(Path.GetExtension(img.Url)));
+                var image = modelVersion.Images?.FirstOrDefault(
+                    img => LocalModelFile.SupportedImageExtensions.Contains(Path.GetExtension(img.Url))
+                );
                 if (image == null)
                     continue;
 
@@ -179,6 +209,11 @@ public class MetadataImportService(
 
                 success++;
             }
+            catch (Exception e)
+            {
+                logger.LogError(e, "Error while updating {filePath}", filePath);
+            }
+        }
     }
 
     public async Task<ConnectedModelInfo?> GetMetadataForFile(
@@ -223,11 +258,13 @@ public class MetadataImportService(
         var (model, modelVersion, modelFile) = modelInfo.Value;
 
         var updatedCmInfo = new ConnectedModelInfo(model, modelVersion, modelFile, DateTimeOffset.UtcNow);
-        await updatedCmInfo.SaveJsonToDirectory(filePath.Directory, fileNameWithoutExtension).ConfigureAwait(false);
+        await updatedCmInfo
+            .SaveJsonToDirectory(filePath.Directory, fileNameWithoutExtension)
+            .ConfigureAwait(false);
 
-        var image = modelVersion
-            .Images
-            ?.FirstOrDefault(img => LocalModelFile.SupportedImageExtensions.Contains(Path.GetExtension(img.Url)));
+        var image = modelVersion.Images?.FirstOrDefault(
+            img => LocalModelFile.SupportedImageExtensions.Contains(Path.GetExtension(img.Url))
+        );
 
         if (image == null)
             return updatedCmInfo;