From 48b81afc6816fb0ce137ab141a409de3aa5cb884 Mon Sep 17 00:00:00 2001
From: Ionite <dev@ionite.io>
Date: Mon, 2 Oct 2023 01:59:38 -0400
Subject: [PATCH] Fix interrupt handler disposed too early

---
 .../Base/InferenceGenerationViewModelBase.cs   | 18 ++++++++++++++----
 .../ViewModels/Base/ProgressViewModel.cs       | 11 +++++++++--
 2 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/StabilityMatrix.Avalonia/ViewModels/Base/InferenceGenerationViewModelBase.cs b/StabilityMatrix.Avalonia/ViewModels/Base/InferenceGenerationViewModelBase.cs
index f32a93b8..7ee884e2 100644
--- a/StabilityMatrix.Avalonia/ViewModels/Base/InferenceGenerationViewModelBase.cs
+++ b/StabilityMatrix.Avalonia/ViewModels/Base/InferenceGenerationViewModelBase.cs
@@ -132,7 +132,19 @@ public abstract partial class InferenceGenerationViewModelBase
 
             // Register progress handler
             promptTask.ProgressUpdate += OnProgressUpdateReceived;
-            promptTask.RunningNodeChanged += OnRunningNodeChanged;
+
+            // Delay attaching running node change handler to not show indeterminate progress
+            // if progress updates are received before the prompt starts
+            Task.Run(
+                    async () =>
+                    {
+                        await Task.Delay(200, cancellationToken);
+                        // ReSharper disable once AccessToDisposedClosure
+                        promptTask.RunningNodeChanged += OnRunningNodeChanged;
+                    },
+                    cancellationToken
+                )
+                .SafeFireAndForget();
 
             // Wait for prompt to finish
             await promptTask.Task.WaitAsync(cancellationToken);
@@ -166,15 +178,13 @@ public abstract partial class InferenceGenerationViewModelBase
             client.PreviewImageReceived -= OnPreviewImageReceived;
 
             // Clear progress
-            OutputProgress.Value = 0;
-            OutputProgress.Text = "";
+            OutputProgress.ClearProgress();
             ImageGalleryCardViewModel.PreviewImage?.Dispose();
             ImageGalleryCardViewModel.PreviewImage = null;
             ImageGalleryCardViewModel.IsPreviewOverlayEnabled = false;
 
             // Cleanup tasks
             promptTask?.Dispose();
-            await promptInterrupt.DisposeAsync();
         }
     }
 
diff --git a/StabilityMatrix.Avalonia/ViewModels/Base/ProgressViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/Base/ProgressViewModel.cs
index cc0466e3..40375dfa 100644
--- a/StabilityMatrix.Avalonia/ViewModels/Base/ProgressViewModel.cs
+++ b/StabilityMatrix.Avalonia/ViewModels/Base/ProgressViewModel.cs
@@ -12,16 +12,23 @@ public partial class ProgressViewModel : ViewModelBase
 
     [ObservableProperty]
     private string? description;
-    
+
     [ObservableProperty, NotifyPropertyChangedFor(nameof(IsProgressVisible))]
     private double value;
 
     [ObservableProperty]
     private double maximum = 100;
-    
+
     [ObservableProperty, NotifyPropertyChangedFor(nameof(IsProgressVisible))]
     private bool isIndeterminate;
 
     public virtual bool IsProgressVisible => Value > 0 || IsIndeterminate;
     public virtual bool IsTextVisible => !string.IsNullOrWhiteSpace(Text);
+
+    public void ClearProgress()
+    {
+        Value = 0;
+        Text = null;
+        IsIndeterminate = false;
+    }
 }