Move image source logic to SelectImageCard Step

StabilityMatrix.Avalonia/Extensions/ComfyNodeBuilderExtensions.cs

@@ -2,6 +2,7 @@
 using System.ComponentModel.DataAnnotations;
 using System.Drawing;
 using System.IO;
+using StabilityMatrix.Avalonia.Models;
 using StabilityMatrix.Avalonia.Models.Inference;
 using StabilityMatrix.Avalonia.ViewModels.Inference;
 using StabilityMatrix.Core.Models.Api.Comfy.Nodes;
@@ -49,24 +50,18 @@ public static class ComfyNodeBuilderExtensions
         }
     }
 
-    public static void SetupImageLatentSource(
+    /// <summary>
+    /// Setup an image as the <see cref="ComfyNodeBuilder.NodeBuilderConnections.Primary"/> connection
+    /// </summary>
+    public static void SetupImagePrimarySource(
         this ComfyNodeBuilder builder,
-        SelectImageCardViewModel selectImageCardViewModel,
+        ImageSource image,
+        Size imageSize,
         int? batchIndex = null
     )
     {
         // Get source image
-        var sourceImage = selectImageCardViewModel.ImageSource;
-        var sourceImageRelativePath = Path.Combine(
-            "Inference",
-            sourceImage!.GetHashGuidFileNameCached()
-        );
-        var sourceImageSize =
-            selectImageCardViewModel.CurrentBitmapSize
-            ?? throw new InvalidOperationException("Source image size is null");
-
-        // Set source size
-        builder.Connections.PrimarySize = sourceImageSize;
+        var sourceImageRelativePath = Path.Combine("Inference", image.GetHashGuidFileNameCached());
 
         // Load source
         var loadImage = builder.Nodes.AddTypedNode(
@@ -74,6 +69,7 @@ public static class ComfyNodeBuilderExtensions
         );
 
         builder.Connections.Primary = loadImage.Output1;
+        builder.Connections.PrimarySize = imageSize;
 
         // If batch index is selected, add a LatentFromBatch
         if (batchIndex is not null)

StabilityMatrix.Avalonia/ViewModels/Inference/BatchSizeCardViewModel.cs

@@ -1,14 +1,16 @@
 using CommunityToolkit.Mvvm.ComponentModel;
 using StabilityMatrix.Avalonia.Controls;
+using StabilityMatrix.Avalonia.Models.Inference;
 using StabilityMatrix.Avalonia.ViewModels.Base;
 using StabilityMatrix.Core.Attributes;
+using StabilityMatrix.Core.Models.Api.Comfy.Nodes;
 
 namespace StabilityMatrix.Avalonia.ViewModels.Inference;
 
 [View(typeof(BatchSizeCard))]
 [ManagedService]
 [Transient]
-public partial class BatchSizeCardViewModel : LoadableViewModelBase
+public partial class BatchSizeCardViewModel : LoadableViewModelBase, IComfyStep
 {
     [ObservableProperty]
     private int batchSize = 1;
@@ -21,4 +23,18 @@ public partial class BatchSizeCardViewModel : LoadableViewModelBase
 
     [ObservableProperty]
     private int batchIndex = 1;
+
+    /// <summary>
+    /// Sets batch size to connections.
+    /// Provides:
+    /// <list type="number">
+    /// <item><see cref="ComfyNodeBuilder.NodeBuilderConnections.BatchSize"/></item>
+    /// <item><see cref="ComfyNodeBuilder.NodeBuilderConnections.BatchIndex"/></item>
+    /// </list>
+    /// </summary>
+    public void ApplyStep(ModuleApplyStepEventArgs e)
+    {
+        e.Builder.Connections.BatchSize = BatchSize;
+        e.Builder.Connections.BatchIndex = IsBatchIndexEnabled ? BatchIndex : null;
+    }
 }

StabilityMatrix.Avalonia/ViewModels/Inference/InferenceImageToImageViewModel.cs

@@ -1,5 +1,6 @@
 using System;
 using System.Collections.Generic;
+using System.ComponentModel.DataAnnotations;
 using System.Linq;
 using System.Text.Json.Serialization;
 using System.Threading;
@@ -114,17 +115,20 @@ public partial class InferenceImageToImageViewModel
 
         var builder = args.Builder;
 
+        // Setup constants
         builder.Connections.Seed = args.SeedOverride switch
         {
             { } seed => Convert.ToUInt64(seed),
             _ => Convert.ToUInt64(SeedCardViewModel.Seed)
         };
 
+        BatchSizeCardViewModel.ApplyStep(args);
+
         // Load models
         ModelCardViewModel.ApplyStep(args);
 
-        // Setup image latent
-        builder.SetupImageLatentSource(SelectImageCardViewModel, BatchSizeCardViewModel.BatchIndex);
+        // Setup image latent source
+        SelectImageCardViewModel.ApplyStep(args);
 
         // Prompts and loras
         PromptCardViewModel.ApplyStep(args);
@@ -132,7 +136,7 @@ public partial class InferenceImageToImageViewModel
         // Setup Sampler and Refiner if enabled
         SamplerCardViewModel.ApplyStep(args);
 
-        // Hires fix if enabled
+        // Apply module steps
         foreach (var module in ModulesCardViewModel.Cards.OfType<ModuleBase>())
         {
             module.ApplyStep(args);

StabilityMatrix.Avalonia/ViewModels/Inference/InferenceImageUpscaleViewModel.cs

@@ -1,12 +1,9 @@
-using System;
-using System.Collections.Generic;
+using System.Collections.Generic;
 using System.Diagnostics.CodeAnalysis;
 using System.Linq;
 using System.Text.Json.Serialization;
 using System.Threading;
 using System.Threading.Tasks;
-using AsyncAwaitBestPractices;
-using DynamicData.Binding;
 using NLog;
 using StabilityMatrix.Avalonia.Extensions;
 using StabilityMatrix.Avalonia.Models;
@@ -19,7 +16,6 @@ using StabilityMatrix.Core.Extensions;
 using StabilityMatrix.Core.Models;
 using StabilityMatrix.Core.Models.Api.Comfy.Nodes;
 using StabilityMatrix.Core.Services;
-using Path = System.IO.Path;
 
 #pragma warning disable CS0657 // Not a valid attribute location for this declaration
 
@@ -86,17 +82,6 @@ public class InferenceImageUpscaleViewModel : InferenceGenerationViewModelBase
                 stackExpander.AddCards(SharpenCardViewModel);
             })
         );
-
-        // On any new images, copy to input dir
-        /*SelectImageCardViewModel
-            .WhenPropertyChanged(x => x.ImageSource)
-            .Subscribe(e =>
-            {
-                if (e.Value?.LocalFile?.FullPath is { } path)
-                {
-                    ClientManager.CopyImageToInputAsync(path).SafeFireAndForget();
-                }
-            });*/
     }
 
     /// <inheritdoc />
@@ -116,24 +101,8 @@ public class InferenceImageUpscaleViewModel : InferenceGenerationViewModelBase
         var builder = args.Builder;
         var nodes = builder.Nodes;
 
-        // Get source image
-        var sourceImage = SelectImageCardViewModel.ImageSource;
-        var sourceImageRelativePath = Path.Combine(
-            "Inference",
-            sourceImage!.GetHashGuidFileNameCached()
-        );
-        var sourceImageSize =
-            SelectImageCardViewModel.CurrentBitmapSize
-            ?? throw new InvalidOperationException("Source image size is null");
-
-        // Set source size
-        builder.Connections.PrimarySize = sourceImageSize;
-
-        // Load source
-        var loadImage = nodes.AddTypedNode(
-            new ComfyNodeBuilder.LoadImage { Name = "LoadImage", Image = sourceImageRelativePath }
-        );
-        builder.Connections.Primary = loadImage.Output1;
+        // Setup image source
+        SelectImageCardViewModel.ApplyStep(args);
 
         // If upscale is enabled, add another upscale group
         if (IsUpscaleEnabled)

StabilityMatrix.Avalonia/ViewModels/Inference/SelectImageCardViewModel.cs

@@ -1,5 +1,6 @@
 using System;
 using System.Collections.Generic;
+using System.ComponentModel.DataAnnotations;
 using System.Linq;
 using System.Runtime.CompilerServices;
 using System.Threading.Tasks;
@@ -12,7 +13,9 @@ using CommunityToolkit.Mvvm.ComponentModel;
 using CommunityToolkit.Mvvm.Input;
 using NLog;
 using StabilityMatrix.Avalonia.Controls;
+using StabilityMatrix.Avalonia.Extensions;
 using StabilityMatrix.Avalonia.Models;
+using StabilityMatrix.Avalonia.Models.Inference;
 using StabilityMatrix.Avalonia.Services;
 using StabilityMatrix.Avalonia.ViewModels.Base;
 using StabilityMatrix.Core.Attributes;
@@ -27,12 +30,13 @@ namespace StabilityMatrix.Avalonia.ViewModels.Inference;
 [View(typeof(SelectImageCard))]
 [ManagedService]
 [Transient]
-public partial class SelectImageCardViewModel : ViewModelBase, IDropTarget
+public partial class SelectImageCardViewModel(INotificationService notificationService)
+    : ViewModelBase,
+        IDropTarget,
+        IComfyStep
 {
     private static readonly Logger Logger = LogManager.GetCurrentClassLogger();
 
-    private readonly INotificationService notificationService;
-
     [ObservableProperty]
     [NotifyPropertyChangedFor(nameof(IsSelectionAvailable))]
     private ImageSource? imageSource;
@@ -53,9 +57,14 @@ public partial class SelectImageCardViewModel : ViewModelBase, IDropTarget
             Convert.ToInt32(CurrentBitmap?.Size.Height ?? 0)
         );
 
-    public SelectImageCardViewModel(INotificationService notificationService)
+    /// <inheritdoc />
+    public void ApplyStep(ModuleApplyStepEventArgs e)
     {
-        this.notificationService = notificationService;
+        e.Builder.SetupImagePrimarySource(
+            ImageSource ?? throw new ValidationException("Input Image is required"),
+            CurrentBitmapSize ?? throw new ValidationException("Input Image is required"),
+            e.Builder.Connections.BatchIndex
+        );
     }
 
     [MethodImpl(MethodImplOptions.Synchronized)]