Merge pull request #427 from ionite34/fast-png

Fast png

Author: mohnjiles

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 AsyncImageLoader.Loaders;
 using Avalonia.Media.Imaging;
 using StabilityMatrix.Core.Extensions;
+using StabilityMatrix.Core.Helper;
 
 namespace StabilityMatrix.Avalonia;
 
@@ -42,7 +43,11 @@ protected void OnLoadFailed(string url, Exception exception) =>
         {
             try
             {
-                return new Bitmap(url);
+                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 @@ ILogger<OutputsPageViewModel> logger
             delay: TimeSpan.FromMilliseconds(250)
         );
 
-        RefreshCategories();
+        RefreshCategories(false);
     }
 
     public override void OnLoaded()
@@ -163,10 +160,7 @@ public override void OnLoaded()
         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 async Task ShowImageDialog(OutputImageViewModel item)
             )
             .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 Task CopyImage(string imagePath)
 
     public void Refresh()
     {
-        Dispatcher.UIThread.Post(RefreshCategories);
+        Dispatcher.UIThread.Post(() => RefreshCategories());
         Dispatcher.UIThread.Post(OnLoaded);
     }
 
@@ -430,9 +424,7 @@ public async Task ConsolidateImages()
 
         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 @@ var category in stackPanel.Children.OfType<CheckBox>().Where(c => c.IsChecked ==
 
             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 @@ private void GetOutputs(string directory)
         }
     }
 
-    private void RefreshCategories()
+    private void RefreshCategories(bool updateProperty = true)
     {
         if (Design.IsDesignMode)
             return;
@@ -534,15 +520,11 @@ private void RefreshCategories()
 
         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 @@ private void RefreshCategories()
 
         packageCategories.Insert(
             1,
-            new PackageOutputCategory
-            {
-                Path = settingsManager.ImagesInferenceDirectory,
-                Name = "Inference"
-            }
+            new PackageOutputCategory { Path = settingsManager.ImagesInferenceDirectory, Name = "Inference" }
         );
 
         Categories = new ObservableCollection<PackageOutputCategory>(packageCategories);
-        SelectedCategory =
-            Categories.FirstOrDefault(x => x.Name == previouslySelectedCategory?.Name)
-            ?? Categories.First();
+
+        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 static string ReadTextChunk(BinaryReader byteStream, string key)
 
         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;
+    }
 }