From 586b0029a51102eb4e939c46e126a9aca06654ae Mon Sep 17 00:00:00 2001
From: Ionite <dev@ionite.io>
Date: Sat, 26 Aug 2023 18:47:25 -0400
Subject: [PATCH] Image viewer fixes

---
 .../Controls/FrameCarousel.axaml              |  1 +
 .../Controls/FrameCarousel.axaml.cs           | 52 ++++++-------------
 .../Controls/ImageGalleryCard.axaml           | 41 +++------------
 .../InferenceTextToImageViewModel.cs          | 34 ++++++------
 4 files changed, 41 insertions(+), 87 deletions(-)

diff --git a/StabilityMatrix.Avalonia/Controls/FrameCarousel.axaml b/StabilityMatrix.Avalonia/Controls/FrameCarousel.axaml
index 5d062a39..fe399fce 100644
--- a/StabilityMatrix.Avalonia/Controls/FrameCarousel.axaml
+++ b/StabilityMatrix.Avalonia/Controls/FrameCarousel.axaml
@@ -16,6 +16,7 @@
                         CornerRadius="{TemplateBinding CornerRadius}">
                     <ui:Frame 
                         Name="PART_Frame"
+                        CornerRadius="{TemplateBinding CornerRadius}"
                         ContentTemplate="{TemplateBinding ContentTemplate}"/>
                 </Border>
             </ControlTemplate>
diff --git a/StabilityMatrix.Avalonia/Controls/FrameCarousel.axaml.cs b/StabilityMatrix.Avalonia/Controls/FrameCarousel.axaml.cs
index 39887880..c6637b31 100644
--- a/StabilityMatrix.Avalonia/Controls/FrameCarousel.axaml.cs
+++ b/StabilityMatrix.Avalonia/Controls/FrameCarousel.axaml.cs
@@ -1,14 +1,10 @@
 using System;
 using System.Collections;
-using System.Collections.Generic;
-using System.Collections.Specialized;
 using System.Diagnostics.CodeAnalysis;
 using Avalonia;
 using Avalonia.Controls;
 using Avalonia.Controls.Primitives;
 using Avalonia.Controls.Templates;
-using Avalonia.Interactivity;
-using DynamicData.Binding;
 using FluentAvalonia.UI.Controls;
 using FluentAvalonia.UI.Media.Animation;
 using FluentAvalonia.UI.Navigation;
@@ -67,24 +63,6 @@ public class FrameCarousel : SelectingItemsControl
             TransitionInfoOverride = new SuppressNavigationTransitionInfo()
         };
     
-    private void ItemsView_CollectionChanged(object? sender, NotifyCollectionChangedEventArgs e)
-    {
-        if (frame is null) return;
-        
-        // If current size is 0, reset the frame cache
-        var count = ItemCount;
-        if (count == 0)
-        {
-            frame.ClearValue(Frame.CacheSizeProperty);
-            // Setting this to false clears the page cache and stack caches
-            frame.IsNavigationStackEnabled = false;
-        }
-        else
-        {
-            frame.SetValue(Frame.CacheSizeProperty, count);
-        }
-    }
-    
     /// <inheritdoc />
     protected override void OnApplyTemplate(TemplateAppliedEventArgs e)
     {
@@ -95,32 +73,21 @@ public class FrameCarousel : SelectingItemsControl
         
         frame.NavigationPageFactory = new FrameNavigationFactory(SourcePageType);
 
-        ItemsView.CollectionChanged += ItemsView_CollectionChanged;
-
         if (SelectedItem is not null)
         {
             frame.NavigateFromObject(SelectedItem, DirectionlessNavigationOptions);
         }
     }
 
-    /// <inheritdoc />
-    protected override void OnLoaded(RoutedEventArgs e)
-    {
-        base.OnLoaded(e);
-        
-        if (SelectedItem is not null)
-        {
-            frame!.NavigateFromObject(SelectedItem, DirectionlessNavigationOptions);
-        }
-    }
-
     protected override void OnPropertyChanged(AvaloniaPropertyChangedEventArgs change)
     {
         base.OnPropertyChanged(change);
 
-        if (change.Property == SelectedItemProperty && frame is not null)
+        if (frame is null) return;
+        
+        if (change.Property == SelectedItemProperty)
         {
-            var value = change.GetNewValue<object?>();
+            if (change.GetNewValue<object?>() is not { } value) return;
             
             if (SelectedIndex > previousIndex)
             {
@@ -138,6 +105,17 @@ public class FrameCarousel : SelectingItemsControl
             }
             
             previousIndex = SelectedIndex;
+        } 
+        else if (change.Property == ItemCountProperty)
+        {
+            // On item count change to 0, clear the frame cache
+            var value = change.GetNewValue<int>();
+
+            if (value == 0)
+            {
+                var pageCache = frame.GetPrivateField<IList>("_pageCache");
+                pageCache?.Clear();
+            }
         }
     }
     
diff --git a/StabilityMatrix.Avalonia/Controls/ImageGalleryCard.axaml b/StabilityMatrix.Avalonia/Controls/ImageGalleryCard.axaml
index bf9a23ca..b6d0265c 100644
--- a/StabilityMatrix.Avalonia/Controls/ImageGalleryCard.axaml
+++ b/StabilityMatrix.Avalonia/Controls/ImageGalleryCard.axaml
@@ -3,8 +3,6 @@
     xmlns:x="http://schemas.microsoft.com/winfx/2006/xaml"
     xmlns:controls="using:StabilityMatrix.Avalonia.Controls"
     xmlns:mocks="clr-namespace:StabilityMatrix.Avalonia.DesignData"
-    xmlns:system="clr-namespace:System;assembly=System.Runtime"
-    xmlns:ui="using:FluentAvalonia.UI.Controls"
     xmlns:vmInference="clr-namespace:StabilityMatrix.Avalonia.ViewModels.Inference"
     xmlns:models="clr-namespace:StabilityMatrix.Avalonia.Models"
     x:DataType="vmInference:ImageGalleryCardViewModel">
@@ -40,40 +38,13 @@
                         <!-- Main image view -->
                         <Grid Margin="4">
                             <controls:FrameCarousel 
-                                x:Name="ImageCarousel"
+                                x:Name="PART_ImageCarousel" 
+                                CornerRadius="4"
+                                IsVisible="{Binding !IsPreviewOverlayEnabled}"
                                 SourcePageType="controls:AdvancedImageBoxView"
                                 SelectedItem="{Binding SelectedImage}"
                                 SelectedIndex="{Binding SelectedImageIndex}"
-                                ItemsSource="{Binding ImageSources}">
-                                <!--<controls:FrameCarousel.ContentTemplate>
-                                    <DataTemplate DataType="{x:Type models:ImageSource}">
-                                        <controls:AdvancedImageBox
-                                            x:Name="ImageView"
-                                            CornerRadius="4"
-                                            SizeMode="Fit"
-                                            Image="{Binding BitmapAsync^}">
-                                            <controls:AdvancedImageBox.ContextFlyout>
-                                                <ui:FAMenuFlyout>
-                                                    ~2~<ui:MenuFlyoutItem 
-                                                        Command="{ReflectionBinding #ImageCarousel.DataContext.FlyoutPreviewCommand}" 
-                                                        CommandParameter="{Binding #ImageView.Image}"
-                                                        Text="Expanded Preview"
-                                                        HotKey="Space"
-                                                        IconSource="ZoomInFilled" />
-                                                    <ui:MenuFlyoutSeparator/>@2@
-                                                    <ui:MenuFlyoutItem 
-                                                        IsEnabled="{OnPlatform Windows=True, Default=False}"
-                                                        Command="{ReflectionBinding #ImageCarousel.DataContext.FlyoutCopyCommand}" 
-                                                        CommandParameter="{Binding #ImageView.Image}"
-                                                        Text="Copy"
-                                                        HotKey="Ctrl+C"
-                                                        IconSource="Copy" />
-                                                </ui:FAMenuFlyout>
-                                            </controls:AdvancedImageBox.ContextFlyout>
-                                        </controls:AdvancedImageBox>
-                                    </DataTemplate>
-                                </controls:FrameCarousel.ContentTemplate>-->
-                            </controls:FrameCarousel>
+                                ItemsSource="{Binding ImageSources}" />
 
                             <!--<Carousel
                                 ScrollViewer.IsScrollChainingEnabled="False"
@@ -152,7 +123,7 @@
                                     IsEnabled="{Binding CanNavigateBack}"
                                     Classes="transparent"
                                     Padding="10,20"
-                                    Command="{Binding #ImageCarousel.Previous}">
+                                    Command="{Binding #PART_ImageCarousel.Previous}">
                                     <Path Data="M20,11V13H8L13.5,18.5L12.08,19.92L4.16,12L12.08,4.08L13.5,5.5L8,11H20Z" 
                                           Fill="{DynamicResource ButtonForeground}" />
                                 </Button>
@@ -167,7 +138,7 @@
                                     IsEnabled="{Binding CanNavigateForward}"
                                     Classes="transparent"
                                     Padding="10,20"
-                                    Command="{Binding #ImageCarousel.Next}">
+                                    Command="{Binding #PART_ImageCarousel.Next}">
                                     <Path Data="M4,11V13H16L10.5,18.5L11.92,19.92L19.84,12L11.92,4.08L10.5,5.5L16,11H4Z" 
                                           Fill="{DynamicResource ButtonForeground}" />
                                 </Button>
diff --git a/StabilityMatrix.Avalonia/ViewModels/Inference/InferenceTextToImageViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/Inference/InferenceTextToImageViewModel.cs
index ebd29af4..3851ea61 100644
--- a/StabilityMatrix.Avalonia/ViewModels/Inference/InferenceTextToImageViewModel.cs
+++ b/StabilityMatrix.Avalonia/ViewModels/Inference/InferenceTextToImageViewModel.cs
@@ -1,23 +1,18 @@
 using System;
 using System.Collections.Generic;
 using System.Collections.Immutable;
-using System.Diagnostics.CodeAnalysis;
 using System.IO;
 using System.Linq;
-using System.Text.Json.Nodes;
 using System.Text.Json.Serialization;
 using System.Threading;
 using System.Threading.Tasks;
 using AsyncAwaitBestPractices;
 using Avalonia.Media.Imaging;
-using AvaloniaEdit.Document;
 using CommunityToolkit.Mvvm.ComponentModel;
 using CommunityToolkit.Mvvm.Input;
 using NLog;
 using Refit;
 using SkiaSharp;
-using StabilityMatrix.Avalonia.Controls;
-using StabilityMatrix.Avalonia.Extensions;
 using StabilityMatrix.Avalonia.Helpers;
 using StabilityMatrix.Avalonia.Models;
 using StabilityMatrix.Avalonia.Services;
@@ -30,6 +25,7 @@ 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;
+
 #pragma warning disable CS0657 // Not a valid attribute location for this declaration
 
 namespace StabilityMatrix.Avalonia.ViewModels.Inference;
@@ -129,7 +125,7 @@ public partial class InferenceTextToImageViewModel : InferenceTabViewModelBase
             }
         );
 
-        GenerateImageCommand.WithNotificationErrorHandler(notificationService);
+        // GenerateImageCommand.WithNotificationErrorHandler(notificationService);
     }
 
     private (NodeDictionary prompt, string[] outputs) BuildPrompt()
@@ -335,7 +331,7 @@ public partial class InferenceTextToImageViewModel : InferenceTabViewModelBase
         // Decode to bitmap
         using var stream = new MemoryStream(args.ImageBytes);
         var bitmap = new Bitmap(stream);
-
+        
         ImageGalleryCardViewModel.PreviewImage?.Dispose();
         ImageGalleryCardViewModel.PreviewImage = bitmap;
         ImageGalleryCardViewModel.IsPreviewOverlayEnabled = true;
@@ -364,6 +360,7 @@ public partial class InferenceTextToImageViewModel : InferenceTabViewModelBase
         // client.ProgressUpdateReceived += OnProgressUpdateReceived;
         client.PreviewImageReceived += OnPreviewImageReceived;
 
+        
         ComfyTask? promptTask = null;
         try
         {
@@ -434,18 +431,25 @@ public partial class InferenceTextToImageViewModel : InferenceTabViewModelBase
                 var lastName = outputImages.Last().LocalFile?.Info.Name;
                 var gridPath = client.OutputImagesDir!.JoinFile($"grid-{lastName}");
 
-                await using var fileStream = gridPath.Info.OpenWrite();
-                await fileStream.WriteAsync(grid.Encode().ToArray(), cancellationToken);
+                await using (var fileStream = gridPath.Info.OpenWrite())
+                {
+                    await fileStream.WriteAsync(grid.Encode().ToArray(), cancellationToken);
+                }
 
                 // Insert to start of images
-                ImageGalleryCardViewModel.ImageSources.Add(new ImageSource(gridPath));
+                var gridImage = new ImageSource(gridPath);
+                // Preload
+                await gridImage.GetBitmapAsync();
+                ImageGalleryCardViewModel.ImageSources.Add(gridImage);
             }
 
             // Add rest of images
-            ImageGalleryCardViewModel.ImageSources.AddRange(outputImages);
-            
-            // Preload all images
-            await Task.WhenAll(outputImages.Select(i => i.BitmapAsync));
+            foreach (var img in outputImages)
+            {
+                // Preload
+                await img.GetBitmapAsync();
+                ImageGalleryCardViewModel.ImageSources.Add(img);
+            }
         }
         finally
         {
@@ -462,7 +466,7 @@ public partial class InferenceTextToImageViewModel : InferenceTabViewModelBase
         }
     }
 
-    [RelayCommand(IncludeCancelCommand = true, FlowExceptionsToTaskScheduler = true)]
+    [RelayCommand(IncludeCancelCommand = true)]
     private async Task GenerateImage(CancellationToken cancellationToken = default)
     {
         try