fix: sync WAN viewmodel and ComfyClient after cleanup

Author: Domica

StabilityMatrix.Avalonia/ViewModels/Inference/InferenceWanTextToVideoViewModel.cs

@@ -70,8 +70,8 @@ RunningPackageService runningPackageService
 
         BatchSizeCardViewModel = vmFactory.Get<BatchSizeCardViewModel>();
 
-        VideoOutputSettingsCardViewModel = vmFactory.Get<VideoOutputSettingsCardViewModel>(
-            vm => vm.Fps = 16.0d
+        VideoOutputSettingsCardViewModel = vmFactory.Get<VideoOutputSettingsCardViewModel>(vm =>
+            vm.Fps = 16.0d
         );
 
         StackCardViewModel = vmFactory.Get<StackCardViewModel>();
@@ -94,7 +94,7 @@ protected override void BuildPrompt(BuildPromptEventArgs args)
         builder.Connections.Seed = args.SeedOverride switch
         {
             { } seed => Convert.ToUInt64(seed),
-            _ => Convert.ToUInt64(SeedCardViewModel.Seed)
+            _ => Convert.ToUInt64(SeedCardViewModel.Seed),
         };
 
         // Load models
@@ -115,14 +115,13 @@ protected override void BuildPrompt(BuildPromptEventArgs args)
 
         SamplerCardViewModel.ApplyStep(args);
 
-        // ADD THIS - Apply TiledVAE module
+        // Apply TiledVAE module
         TiledVAEModule.ApplyStep(args);
 
-        // ADD THIS LINE - Invoke pre-output actions BEFORE video output!
+        // Invoke pre-output actions BEFORE video output!
         args.Builder.InvokeAllPreOutputActions();
 
         VideoOutputSettingsCardViewModel.ApplyStep(args);
-
     }
 
     /// <inheritdoc />
@@ -171,13 +170,13 @@ CancellationToken cancellationToken
                 OutputNodeNames = buildPromptArgs.Builder.Connections.OutputNodeNames.ToArray(),
                 Parameters = SaveStateToParameters(new GenerationParameters()) with
                 {
-                    Seed = Convert.ToUInt64(seed)
+                    Seed = Convert.ToUInt64(seed),
                 },
                 Project = inferenceProject,
                 FilesToTransfer = buildPromptArgs.FilesToTransfer,
                 BatchIndex = i,
                 // Only clear output images on the first batch
-                ClearOutputImages = i == 0
+                ClearOutputImages = i == 0,
             };
 
             batchArgs.Add(generationArgs);

StabilityMatrix.Core/Inference/ComfyClient.cs

@@ -29,16 +29,15 @@ public class ComfyClient : InferenceClientBase
     private readonly IComfyApi comfyApi;
     private bool isDisposed;
 
-    private readonly JsonSerializerOptions jsonSerializerOptions =
-        new()
+    private readonly JsonSerializerOptions jsonSerializerOptions = new()
+    {
+        PropertyNamingPolicy = JsonNamingPolicies.SnakeCaseLower,
+        Converters =
         {
-            PropertyNamingPolicy = JsonNamingPolicies.SnakeCaseLower,
-            Converters =
-            {
-                new NodeConnectionBaseJsonConverter(),
-                new OneOfJsonConverter<string, StringNodeConnection>()
-            }
-        };
+            new NodeConnectionBaseJsonConverter(),
+            new OneOfJsonConverter<string, StringNodeConnection>(),
+        },
+    };
 
     // ReSharper disable once MemberCanBePrivate.Global
     public string ClientId { get; } = Guid.NewGuid().ToString();
@@ -111,20 +110,20 @@ public ComfyClient(IApiFactory apiFactory, Uri baseAddress)
         {
             Scheme = "ws",
             Path = "/ws",
-            Query = $"clientId={ClientId}"
+            Query = $"clientId={ClientId}",
         }.Uri;
 
         webSocketClient = new WebsocketClient(wsUri)
         {
             Name = nameof(ComfyClient),
-            ReconnectTimeout = TimeSpan.FromSeconds(30)
+            ReconnectTimeout = TimeSpan.FromSeconds(30),
         };
 
-        webSocketClient.DisconnectionHappened.Subscribe(
-            info => Logger.Info("Websocket Disconnected, ({Type})", info.Type)
+        webSocketClient.DisconnectionHappened.Subscribe(info =>
+            Logger.Info("Websocket Disconnected, ({Type})", info.Type)
         );
-        webSocketClient.ReconnectionHappened.Subscribe(
-            info => Logger.Info("Websocket Reconnected, ({Type})", info.Type)
+        webSocketClient.ReconnectionHappened.Subscribe(info =>
+            Logger.Info("Websocket Reconnected, ({Type})", info.Type)
         );
 
         webSocketClient.MessageReceived.Subscribe(OnMessageReceived);
@@ -287,7 +286,7 @@ private void HandleBinaryMessage(byte[] data)
             Array.Reverse(typeBytes);
         }*/
 
-        PreviewImageReceived?.Invoke(this, new ComfyWebSocketImageData { ImageBytes = data[8..], });
+        PreviewImageReceived?.Invoke(this, new ComfyWebSocketImageData { ImageBytes = data[8..] });
     }
 
     public override async Task ConnectAsync(CancellationToken cancellationToken = default)
@@ -326,21 +325,23 @@ public override async Task CloseAsync(CancellationToken cancellationToken = defa
         await webSocketClient.Stop(WebSocketCloseStatus.NormalClosure, string.Empty).ConfigureAwait(false);
     }
 
-public async Task<ComfyTask> QueuePromptAsync(
-    Dictionary<string, ComfyNode> nodes,
-    CancellationToken cancellationToken = default
-)
-{
-    var request = new ComfyPromptRequest { ClientId = ClientId, Prompt = nodes };
+    public async Task<ComfyTask> QueuePromptAsync(
+        Dictionary<string, ComfyNode> nodes,
+        CancellationToken cancellationToken = default
+    )
+    {
+        var request = new ComfyPromptRequest { ClientId = ClientId, Prompt = nodes };
 
-    var result = await comfyApi.PostPrompt(request, cancellationToken).ConfigureAwait(false);
+        var result = await comfyApi.PostPrompt(request, cancellationToken).ConfigureAwait(false);
 
-    var task = new ComfyTask(result.PromptId);
-    PromptTasks.TryAdd(result.PromptId, task);
-    currentPromptTask = task;
+        // Add task to dictionary and set it as the current task
+        var task = new ComfyTask(result.PromptId);
+        PromptTasks.TryAdd(result.PromptId, task);
+        currentPromptTask = task;
+
+        return task;
+    }
 
-    return task;
-}
     public async Task InterruptPromptAsync(CancellationToken cancellationToken = default)
     {
         await comfyApi.PostInterrupt(cancellationToken).ConfigureAwait(false);