Improve progress change behavior

StabilityMatrix.Avalonia/ViewModels/Base/InferenceGenerationViewModelBase.cs

@@ -1,6 +1,7 @@
 using System;
 using System.Collections.Generic;
 using System.Collections.Immutable;
+using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
 using System.Linq;
 using System.Text.Json.Serialization;
@@ -119,6 +120,8 @@ public abstract partial class InferenceGenerationViewModelBase
 
         try
         {
+            var timer = Stopwatch.StartNew();
+
             try
             {
                 promptTask = await client.QueuePromptAsync(nodes, cancellationToken);
@@ -138,9 +141,13 @@ public abstract partial class InferenceGenerationViewModelBase
             Task.Run(
                     async () =>
                     {
-                        await Task.Delay(200, cancellationToken);
+                        var delayTime = 250 - (int)timer.ElapsedMilliseconds;
+                        if (delayTime > 0)
+                        {
+                            await Task.Delay(delayTime, cancellationToken);
+                        }
                         // ReSharper disable once AccessToDisposedClosure
-                        promptTask.RunningNodeChanged += OnRunningNodeChanged;
+                        AttachRunningNodeChangedHandler(promptTask);
                     },
                     cancellationToken
                 )
@@ -351,13 +358,24 @@ public abstract partial class InferenceGenerationViewModelBase
         });
     }
 
+    private void AttachRunningNodeChangedHandler(ComfyTask comfyTask)
+    {
+        // Do initial update
+        if (comfyTask.RunningNodesHistory.TryPeek(out var lastNode))
+        {
+            OnRunningNodeChanged(comfyTask, lastNode);
+        }
+
+        comfyTask.RunningNodeChanged += OnRunningNodeChanged;
+    }
+
     /// <summary>
     /// Handles the node executing updates received event from the websocket.
     /// </summary>
     protected virtual void OnRunningNodeChanged(object? sender, string? nodeName)
     {
         // Ignore if regular progress updates started
-        if (sender is not ComfyTask { LastProgressUpdate: null })
+        if (sender is not ComfyTask { HasProgressUpdateStarted: false })
         {
             return;
         }

StabilityMatrix.Core/Inference/ComfyClient.cs

@@ -282,7 +282,7 @@ public class ComfyClient : InferenceClientBase
 
         // Add task to dictionary and set it as the current task
         var task = new ComfyTask(result.PromptId);
-        PromptTasks[result.PromptId] = task;
+        PromptTasks.TryAdd(result.PromptId, task);
         currentPromptTask = task;
 
         return task;

StabilityMatrix.Core/Inference/ComfyTask.cs

@@ -1,11 +1,10 @@
-using System.ComponentModel;
-using StabilityMatrix.Core.Models.Api.Comfy.WebSocketData;
+using StabilityMatrix.Core.Models.Api.Comfy.WebSocketData;
 
 namespace StabilityMatrix.Core.Inference;
 
 public class ComfyTask : TaskCompletionSource, IDisposable
 {
-    public string Id { get; set; }
+    public string Id { get; }
 
     private string? runningNode;
     public string? RunningNode
@@ -15,11 +14,19 @@ public class ComfyTask : TaskCompletionSource, IDisposable
         {
             runningNode = value;
             RunningNodeChanged?.Invoke(this, value);
+            if (value != null)
+            {
+                RunningNodesHistory.Push(value);
+            }
         }
     }
 
+    public Stack<string> RunningNodesHistory { get; } = new();
+
     public ComfyProgressUpdateEventArgs? LastProgressUpdate { get; private set; }
 
+    public bool HasProgressUpdateStarted => LastProgressUpdate != null;
+
     public EventHandler<ComfyProgressUpdateEventArgs>? ProgressUpdate;
 
     public event EventHandler<string?>? RunningNodeChanged;
@@ -34,7 +41,8 @@ public class ComfyTask : TaskCompletionSource, IDisposable
     /// </summary>
     public void OnProgressUpdate(ComfyWebSocketProgressData update)
     {
-        var args = new ComfyProgressUpdateEventArgs(update.Value, update.Max, Id, RunningNode);
+        RunningNodesHistory.TryPeek(out var lastNode);
+        var args = new ComfyProgressUpdateEventArgs(update.Value, update.Max, Id, lastNode);
         ProgressUpdate?.Invoke(this, args);
         LastProgressUpdate = args;
     }