Merge pull request #287 from ionite34/fix-metadata

CHANGELOG.md

@@ -5,6 +5,18 @@ All notable changes to Stability Matrix will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project adheres to [Semantic Versioning 2.0](https://semver.org/spec/v2.0.0.html).
 
+## v2.5.2
+### Added
+- Right click Inference Batch options to enable selecting a "Batch Index". This can be used to reproduce a specific image from a batch generation. The field will be automatically populated in metadata of individual images from a batch generation.
+  - The index is 1-based, so the first image in a batch is index 1, and the last image is the batch size.
+  - Currently this generates different individual images for batches using Ancestral samplers, due to an upstream ComfyUI issue with noise masking. Looking into fixing this.
+- Inference Batches option now is implemented, previously the setting had no effect
+### Changed
+- Default upscale factor for Inference is now 2x instead of 1x
+### Fixed
+- Fixed batch combined image grids not showing metadata and not being importable
+- Fixed "Call from invalid thread" errors that sometimes occured during update notifications
+
 ## v2.5.1
 ### Added
 - `--skip-install` default launch argument for Automatic1111 Package

StabilityMatrix.Avalonia/Controls/BatchSizeCard.axaml

@@ -4,6 +4,8 @@
         xmlns:vmInference="using:StabilityMatrix.Avalonia.ViewModels.Inference"
         xmlns:controls="using:StabilityMatrix.Avalonia.Controls"
         xmlns:icons="clr-namespace:Projektanker.Icons.Avalonia;assembly=Projektanker.Icons.Avalonia"
+        xmlns:ui="clr-namespace:FluentAvalonia.UI.Controls;assembly=FluentAvalonia"
+        xmlns:lang="clr-namespace:StabilityMatrix.Avalonia.Languages"
         x:DataType="vmInference:BatchSizeCardViewModel">
     
     <Design.PreviewWith>
@@ -12,16 +14,23 @@
                 DataContext="{x:Static mocks:DesignData.BatchSizeCardViewModel}"/>
             <controls:BatchSizeCard 
                 Width="280"
-                DataContext="{x:Static mocks:DesignData.BatchSizeCardViewModel}"/>
+                DataContext="{x:Static mocks:DesignData.BatchSizeCardViewModelWithIndexOption}"/>
         </StackPanel>
     </Design.PreviewWith>
 
     <Style Selector="controls|BatchSizeCard">
         <!-- Set Defaults -->
+        <Setter Property="ContextFlyout">
+            <ui:FAMenuFlyout Placement="BottomEdgeAlignedLeft">
+                <ui:ToggleMenuFlyoutItem
+                    IsChecked="{Binding IsBatchIndexEnabled}"
+                    Text="{x:Static lang:Resources.Label_BatchIndex}"/>
+            </ui:FAMenuFlyout>
+        </Setter>
         <Setter Property="Template">
             <ControlTemplate>
                 <controls:Card Padding="8">
-                    <Grid Margin="8" ColumnDefinitions="*,*">
+                    <Grid Margin="8" RowDefinitions="Auto,Auto" ColumnDefinitions="*,*">
                         <StackPanel>
                             <StackPanel Orientation="Horizontal" >
                                 <TextBlock 
@@ -43,7 +52,9 @@
                                 ClipValueToMinMax="True"/>
                         </StackPanel>
 
-                        <StackPanel Grid.Column="1">
+                        <StackPanel
+                            Grid.Row="0"
+                            Grid.Column="1">
                             <StackPanel Orientation="Horizontal" >
                                 <TextBlock 
                                     VerticalAlignment="Center"
@@ -64,6 +75,34 @@
                                 ClipValueToMinMax="True"/>
                         </StackPanel>
                         
+                        <!-- Optional index selection -->
+                        <StackPanel 
+                            Grid.Row="1" 
+                            Grid.Column="0" 
+                            Grid.ColumnSpan="2"
+                            IsVisible="{Binding IsBatchIndexEnabled}"
+                            Margin="0,8,0,0">
+                            <StackPanel Orientation="Horizontal" >
+                                <TextBlock 
+                                    VerticalAlignment="Center"
+                                    Margin="0,0,4,0"
+                                    Text="{x:Static lang:Resources.Label_BatchIndex}"/>
+                                <icons:Icon
+                                    FontSize="12"
+                                    Value="fa-solid fa-crosshairs"/>
+                            </StackPanel>
+                            <NumericUpDown 
+                                HorizontalAlignment="Stretch"
+                                Margin="0,4,0,0"
+                                MinWidth="120"
+                                Minimum="1"
+                                Maximum="{Binding BatchSize}"
+                                Increment="1"
+                                ParsingNumberStyle="Integer"
+                                Value="{Binding BatchIndex}"
+                                ClipValueToMinMax="True"/>
+                        </StackPanel>
+
                     </Grid>
                 </controls:Card>
             </ControlTemplate>

StabilityMatrix.Avalonia/Controls/SeedCard.axaml

@@ -1,9 +1,7 @@
 <Styles xmlns="https://github.com/avaloniaui"
         xmlns:x="http://schemas.microsoft.com/winfx/2006/xaml"
-        xmlns:d="http://schemas.microsoft.com/expression/blend/2008"
         xmlns:controls="using:StabilityMatrix.Avalonia.Controls"
         xmlns:vmInference="using:StabilityMatrix.Avalonia.ViewModels.Inference"
-        xmlns:ui="clr-namespace:FluentAvalonia.UI.Controls;assembly=FluentAvalonia"
         xmlns:icons="clr-namespace:Projektanker.Icons.Avalonia;assembly=Projektanker.Icons.Avalonia"
         x:DataType="vmInference:SeedCardViewModel"
         xmlns:mocks="clr-namespace:StabilityMatrix.Avalonia.DesignData">

StabilityMatrix.Avalonia/DesignData/DesignData.cs

@@ -581,6 +581,12 @@ The gallery images are often inpainted, but you will get something very similar
     public static BatchSizeCardViewModel BatchSizeCardViewModel =>
         DialogFactory.Get<BatchSizeCardViewModel>();
 
+    public static BatchSizeCardViewModel BatchSizeCardViewModelWithIndexOption =>
+        DialogFactory.Get<BatchSizeCardViewModel>(vm =>
+        {
+            vm.IsBatchIndexEnabled = true;
+        });
+
     public static IList<ICompletionData> SampleCompletionData =>
         new List<ICompletionData>
         {

StabilityMatrix.Avalonia/Extensions/ComfyNodeBuilderExtensions.cs

@@ -31,11 +31,26 @@ public static class ComfyNodeBuilderExtensions
             samplerCardViewModel.Width,
             samplerCardViewModel.Height
         );
+
+        // If batch index is selected, add a LatentFromBatch
+        if (batchSizeCardViewModel.IsBatchIndexEnabled)
+        {
+            builder.Connections.Latent = builder.Nodes
+                .AddNamedNode(
+                    ComfyNodeBuilder.LatentFromBatch(
+                        "LatentFromBatch",
+                        builder.Connections.Latent,
+                        // remote expects a 0-based index, vm is 1-based
+                        batchSizeCardViewModel.BatchIndex - 1,
+                        1
+                    )
+                )
+                .Output;
+        }
     }
 
     public static void SetupBaseSampler(
         this ComfyNodeBuilder builder,
-        SeedCardViewModel seedCardViewModel,
         SamplerCardViewModel samplerCardViewModel,
         PromptCardViewModel promptCardViewModel,
         ModelCardViewModel modelCardViewModel,
@@ -109,7 +124,7 @@ public static class ComfyNodeBuilderExtensions
                 ComfyNodeBuilder.KSampler(
                     "Sampler",
                     builder.Connections.BaseModel,
-                    Convert.ToUInt64(seedCardViewModel.Seed),
+                    builder.Connections.Seed,
                     samplerCardViewModel.Steps,
                     samplerCardViewModel.CfgScale,
                     samplerCardViewModel.SelectedSampler
@@ -136,7 +151,7 @@ public static class ComfyNodeBuilderExtensions
                     "Sampler",
                     builder.Connections.BaseModel,
                     true,
-                    Convert.ToUInt64(seedCardViewModel.Seed),
+                    builder.Connections.Seed,
                     totalSteps,
                     samplerCardViewModel.CfgScale,
                     samplerCardViewModel.SelectedSampler
@@ -158,7 +173,6 @@ public static class ComfyNodeBuilderExtensions
 
     public static void SetupRefinerSampler(
         this ComfyNodeBuilder builder,
-        SeedCardViewModel seedCardViewModel,
         SamplerCardViewModel samplerCardViewModel,
         PromptCardViewModel promptCardViewModel,
         ModelCardViewModel modelCardViewModel,
@@ -232,7 +246,7 @@ public static class ComfyNodeBuilderExtensions
                 "Refiner_Sampler",
                 builder.Connections.RefinerModel,
                 false,
-                Convert.ToUInt64(seedCardViewModel.Seed),
+                builder.Connections.Seed,
                 totalSteps,
                 samplerCardViewModel.CfgScale,
                 samplerCardViewModel.SelectedSampler

StabilityMatrix.Avalonia/Languages/Resources.Designer.cs

@@ -527,6 +527,15 @@ namespace StabilityMatrix.Avalonia.Languages {
             }
         }
         
+        /// <summary>
+        ///   Looks up a localized string similar to Batch Index.
+        /// </summary>
+        public static string Label_BatchIndex {
+            get {
+                return ResourceManager.GetString("Label_BatchIndex", resourceCulture);
+            }
+        }
+        
         /// <summary>
         ///   Looks up a localized string similar to Become a Patron.
         /// </summary>

StabilityMatrix.Avalonia/Languages/Resources.resx

@@ -678,4 +678,7 @@
   <data name="Action_RestoreDefaultLayout" xml:space="preserve">
     <value>Restore Default Layout</value>
   </data>
+  <data name="Label_BatchIndex" xml:space="preserve">
+    <value>Batch Index</value>
+  </data>
 </root>

StabilityMatrix.Avalonia/Models/Inference/SeedCardModel.cs

@@ -3,8 +3,11 @@
 namespace StabilityMatrix.Avalonia.Models.Inference;
 
 [JsonSerializable(typeof(SeedCardModel))]
-public class SeedCardModel
+public record SeedCardModel
 {
-    public string? Seed { get; set; }
-    public bool IsRandomizeEnabled { get; set; }
+    [JsonNumberHandling(
+        JsonNumberHandling.WriteAsString | JsonNumberHandling.AllowReadingFromString
+    )]
+    public long Seed { get; init; }
+    public bool IsRandomizeEnabled { get; init; }
 }

StabilityMatrix.Avalonia/Models/InferenceProjectDocument.cs

@@ -2,6 +2,7 @@
 using System.Text.Json;
 using System.Text.Json.Nodes;
 using System.Text.Json.Serialization;
+using StabilityMatrix.Avalonia.Models.Inference;
 using StabilityMatrix.Avalonia.ViewModels.Inference;
 
 namespace StabilityMatrix.Avalonia.Models;
@@ -10,7 +11,7 @@ namespace StabilityMatrix.Avalonia.Models;
 /// This is the project file for inference tabs
 /// </summary>
 [JsonSerializable(typeof(InferenceProjectDocument))]
-public class InferenceProjectDocument
+public class InferenceProjectDocument : ICloneable
 {
     [JsonIgnore]
     private static readonly JsonSerializerOptions SerializerOptions =
@@ -50,4 +51,89 @@ public class InferenceProjectDocument
             );
         }
     }
+
+    public SeedCardModel? GetSeedModel()
+    {
+        if (State is null || !State.TryGetPropertyValue("Seed", out var seedCard))
+        {
+            return null;
+        }
+
+        return seedCard.Deserialize<SeedCardModel>();
+    }
+
+    /// <summary>
+    /// Returns a new <see cref="InferenceProjectDocument"/> with the State modified.
+    /// </summary>
+    /// <param name="stateModifier">Action that changes the state</param>
+    public InferenceProjectDocument WithState(Action<JsonObject?> stateModifier)
+    {
+        var document = (InferenceProjectDocument)Clone();
+        stateModifier(document.State);
+        return document;
+    }
+
+    public bool TryUpdateModel<T>(string key, Func<T, T> modifier)
+    {
+        if (State is not { } state)
+            return false;
+
+        if (!state.TryGetPropertyValue(key, out var modelNode))
+        {
+            return false;
+        }
+
+        if (modelNode.Deserialize<T>() is not { } model)
+        {
+            return false;
+        }
+
+        modelNode = JsonSerializer.SerializeToNode(modifier(model));
+
+        state[key] = modelNode;
+
+        return true;
+    }
+
+    public bool TryUpdateModel(string key, Func<JsonNode, JsonNode> modifier)
+    {
+        if (State is not { } state)
+            return false;
+
+        if (!state.TryGetPropertyValue(key, out var modelNode) || modelNode is null)
+        {
+            return false;
+        }
+
+        state[key] = modifier(modelNode);
+
+        return true;
+    }
+
+    public InferenceProjectDocument WithBatchSize(int batchSize, int batchCount)
+    {
+        if (State is null)
+            throw new InvalidOperationException("State is null");
+
+        var document = (InferenceProjectDocument)Clone();
+
+        var batchSizeCard =
+            document.State!["BatchSize"]
+            ?? throw new InvalidOperationException("BatchSize card is null");
+
+        batchSizeCard["BatchSize"] = batchSize;
+        batchSizeCard["BatchCount"] = batchCount;
+
+        return document;
+    }
+
+    /// <inheritdoc />
+    public object Clone()
+    {
+        var newObj = (InferenceProjectDocument)MemberwiseClone();
+        // Clone State also since its mutable
+        newObj.State =
+            State == null ? null : JsonSerializer.SerializeToNode(State).Deserialize<JsonObject>();
+        return newObj;
+    }
 }

StabilityMatrix.Avalonia/StabilityMatrix.Avalonia.csproj

@@ -8,7 +8,7 @@
         <ApplicationManifest>app.manifest</ApplicationManifest>
         <AvaloniaUseCompiledBindingsByDefault>true</AvaloniaUseCompiledBindingsByDefault>
         <ApplicationIcon>./Assets/Icon.ico</ApplicationIcon>
-        <Version>2.5.0</Version>
+        <Version>2.5.2-dev.1</Version>
         <InformationalVersion>$(Version)</InformationalVersion>
         <EnableWindowsTargeting>true</EnableWindowsTargeting>
         <AllowUnsafeBlocks>true</AllowUnsafeBlocks>

StabilityMatrix.Avalonia/ViewModels/Base/InferenceGenerationViewModelBase.cs

@@ -20,6 +20,7 @@ using StabilityMatrix.Avalonia.Models.Inference;
 using StabilityMatrix.Avalonia.Services;
 using StabilityMatrix.Avalonia.ViewModels.Dialogs;
 using StabilityMatrix.Avalonia.ViewModels.Inference;
+using StabilityMatrix.Core.Extensions;
 using StabilityMatrix.Core.Helper;
 using StabilityMatrix.Core.Inference;
 using StabilityMatrix.Core.Models;
@@ -166,7 +167,10 @@ public abstract partial class InferenceGenerationViewModelBase
             // Disable cancellation
             await promptInterrupt.DisposeAsync();
 
+            if (args.ClearOutputImages)
+            {
                 ImageGalleryCardViewModel.ImageSources.Clear();
+            }
 
             if (
                 !imageOutputs.TryGetValue(args.OutputNodeNames[0], out var images) || images is null
@@ -199,14 +203,16 @@ public abstract partial class InferenceGenerationViewModelBase
     /// Handles image output metadata for generation runs
     /// </summary>
     private async Task ProcessOutputImages(
-        IEnumerable<ComfyImage> images,
+        IReadOnlyCollection<ComfyImage> images,
         ImageGenerationEventArgs args
     )
     {
         // Write metadata to images
         var outputImages = new List<ImageSource>();
         foreach (
-            var filePath in images.Select(image => image.ToFilePath(args.Client.OutputImagesDir!))
+            var (i, filePath) in images
+                .Select(image => image.ToFilePath(args.Client.OutputImagesDir!))
+                .Enumerate()
         )
         {
             if (!filePath.Exists)
@@ -215,10 +221,37 @@ public abstract partial class InferenceGenerationViewModelBase
                 continue;
             }
 
+            var parameters = args.Parameters!;
+            var project = args.Project!;
+
+            // Lock seed
+            project.TryUpdateModel<SeedCardModel>(
+                "Seed",
+                model => model with { IsRandomizeEnabled = false }
+            );
+
+            // Seed and batch override for batches
+            if (images.Count > 1 && project.ProjectType is InferenceProjectType.TextToImage)
+            {
+                project = (InferenceProjectDocument)project.Clone();
+
+                // Set batch size indexes
+                project.TryUpdateModel(
+                    "BatchSize",
+                    node =>
+                    {
+                        node[nameof(BatchSizeCardViewModel.BatchCount)] = 1;
+                        node[nameof(BatchSizeCardViewModel.IsBatchIndexEnabled)] = true;
+                        node[nameof(BatchSizeCardViewModel.BatchIndex)] = i + 1;
+                        return node;
+                    }
+                );
+            }
+
             var bytesWithMetadata = PngDataHelper.AddMetadata(
                 await filePath.ReadAllBytesAsync(),
-                args.Parameters!,
-                args.Project!
+                parameters,
+                project
             );
 
             await using (var outputStream = filePath.Info.OpenWrite())
@@ -235,6 +268,8 @@ public abstract partial class InferenceGenerationViewModelBase
         // Download all images to make grid, if multiple
         if (outputImages.Count > 1)
         {
+            var outputDir = outputImages[0].LocalFile!.Directory;
+
             var loadedImages = outputImages
                 .Select(i => i.LocalFile)
                 .Where(f => f is { Exists: true })
@@ -245,6 +280,14 @@ public abstract partial class InferenceGenerationViewModelBase
                 })
                 .ToImmutableArray();
 
+            var project = args.Project!;
+
+            // Lock seed
+            project.TryUpdateModel<SeedCardModel>(
+                "Seed",
+                model => model with { IsRandomizeEnabled = false }
+            );
+
             var grid = ImageProcessor.CreateImageGrid(loadedImages);
             var gridBytes = grid.Encode().ToArray();
             var gridBytesWithMetadata = PngDataHelper.AddMetadata(
@@ -255,7 +298,7 @@ public abstract partial class InferenceGenerationViewModelBase
 
             // Save to disk
             var lastName = outputImages.Last().LocalFile?.Info.Name;
-            var gridPath = args.Client.OutputImagesDir!.JoinFile($"grid-{lastName}");
+            var gridPath = outputDir!.JoinFile($"grid-{lastName}");
 
             await using (var fileStream = gridPath.Info.OpenWrite())
             {
@@ -394,13 +437,15 @@ public abstract partial class InferenceGenerationViewModelBase
         public required ComfyClient Client { get; init; }
         public required NodeDictionary Nodes { get; init; }
         public required IReadOnlyList<string> OutputNodeNames { get; init; }
-        public GenerationParameters? Parameters { get; set; }
-        public InferenceProjectDocument? Project { get; set; }
+        public GenerationParameters? Parameters { get; init; }
+        public InferenceProjectDocument? Project { get; init; }
+        public bool ClearOutputImages { get; init; } = true;
     }
 
     public class BuildPromptEventArgs : EventArgs
     {
         public ComfyNodeBuilder Builder { get; } = new();
-        public GenerateOverrides Overrides { get; set; } = new();
+        public GenerateOverrides Overrides { get; init; } = new();
+        public long? SeedOverride { get; init; }
     }
 }

StabilityMatrix.Avalonia/ViewModels/Base/LoadableViewModelBase.cs

@@ -34,19 +34,19 @@ public abstract class LoadableViewModelBase : ViewModelBase, IJsonLoadableState
             && !typeof(IJsonLoadableState).IsAssignableFrom(property.PropertyType)
         )
         {
-            Logger.Trace("Skipping {Property} - read-only", property.Name);
+            Logger.ConditionalTrace("Skipping {Property} - read-only", property.Name);
             return true;
         }
         // Check not JsonIgnore
         if (property.GetCustomAttributes(typeof(JsonIgnoreAttribute), true).Length > 0)
         {
-            Logger.Trace("Skipping {Property} - has [JsonIgnore]", property.Name);
+            Logger.ConditionalTrace("Skipping {Property} - has [JsonIgnore]", property.Name);
             return true;
         }
         // Check not excluded type
         if (SerializerIgnoredTypes.Contains(property.PropertyType))
         {
-            Logger.Trace(
+            Logger.ConditionalTrace(
                 "Skipping {Property} - serializer ignored type {Type}",
                 property.Name,
                 property.PropertyType
@@ -56,7 +56,7 @@ public abstract class LoadableViewModelBase : ViewModelBase, IJsonLoadableState
         // Check not ignored name
         if (SerializerIgnoredNames.Contains(property.Name, StringComparer.Ordinal))
         {
-            Logger.Trace("Skipping {Property} - serializer ignored name", property.Name);
+            Logger.ConditionalTrace("Skipping {Property} - serializer ignored name", property.Name);
             return true;
         }
 
@@ -71,7 +71,7 @@ public abstract class LoadableViewModelBase : ViewModelBase, IJsonLoadableState
         // Has JsonIncludeAttribute
         if (property.GetCustomAttributes(typeof(JsonIncludeAttribute), true).Length > 0)
         {
-            Logger.Trace("Including {Property} - has [JsonInclude]", property.Name);
+            Logger.ConditionalTrace("Including {Property} - has [JsonInclude]", property.Name);
             return true;
         }
 
@@ -94,7 +94,11 @@ public abstract class LoadableViewModelBase : ViewModelBase, IJsonLoadableState
     {
         // Get all of our properties using reflection
         var properties = GetType().GetProperties();
-        Logger.Trace("Serializing {Type} with {Count} properties", GetType(), properties.Length);
+        Logger.ConditionalTrace(
+            "Serializing {Type} with {Count} properties",
+            GetType(),
+            properties.Length
+        );
 
         foreach (var property in properties)
         {
@@ -108,7 +112,7 @@ public abstract class LoadableViewModelBase : ViewModelBase, IJsonLoadableState
                 is JsonPropertyNameAttribute jsonPropertyName
             )
             {
-                Logger.Trace(
+                Logger.ConditionalTrace(
                     "Deserializing {Property} ({Type}) with JsonPropertyName {JsonPropertyName}",
                     property.Name,
                     property.PropertyType,
@@ -120,7 +124,7 @@ public abstract class LoadableViewModelBase : ViewModelBase, IJsonLoadableState
             // Check if property is in the JSON object
             if (!state.TryGetPropertyValue(name, out var value))
             {
-                Logger.Trace("Skipping {Property} - not in JSON object", property.Name);
+                Logger.ConditionalTrace("Skipping {Property} - not in JSON object", property.Name);
                 continue;
             }
 
@@ -133,7 +137,7 @@ public abstract class LoadableViewModelBase : ViewModelBase, IJsonLoadableState
             // For types that also implement IJsonLoadableState, defer to their load implementation
             if (typeof(IJsonLoadableState).IsAssignableFrom(property.PropertyType))
             {
-                Logger.Trace(
+                Logger.ConditionalTrace(
                     "Loading {Property} ({Type}) with IJsonLoadableState",
                     property.Name,
                     property.PropertyType
@@ -171,7 +175,11 @@ public abstract class LoadableViewModelBase : ViewModelBase, IJsonLoadableState
             }
             else
             {
-                Logger.Trace("Loading {Property} ({Type})", property.Name, property.PropertyType);
+                Logger.ConditionalTrace(
+                    "Loading {Property} ({Type})",
+                    property.Name,
+                    property.PropertyType
+                );
 
                 var propertyValue = value.Deserialize(property.PropertyType, SerializerOptions);
                 property.SetValue(this, propertyValue);
@@ -195,7 +203,11 @@ public abstract class LoadableViewModelBase : ViewModelBase, IJsonLoadableState
     {
         // Get all of our properties using reflection.
         var properties = GetType().GetProperties();
-        Logger.Trace("Serializing {Type} with {Count} properties", GetType(), properties.Length);
+        Logger.ConditionalTrace(
+            "Serializing {Type} with {Count} properties",
+            GetType(),
+            properties.Length
+        );
 
         // Create a JSON object to store the state.
         var state = new JsonObject();
@@ -218,7 +230,7 @@ public abstract class LoadableViewModelBase : ViewModelBase, IJsonLoadableState
                 is JsonPropertyNameAttribute jsonPropertyName
             )
             {
-                Logger.Trace(
+                Logger.ConditionalTrace(
                     "Serializing {Property} ({Type}) with JsonPropertyName {JsonPropertyName}",
                     property.Name,
                     property.PropertyType,
@@ -230,7 +242,7 @@ public abstract class LoadableViewModelBase : ViewModelBase, IJsonLoadableState
             // For types that also implement IJsonLoadableState, defer to their implementation.
             if (typeof(IJsonLoadableState).IsAssignableFrom(property.PropertyType))
             {
-                Logger.Trace(
+                Logger.ConditionalTrace(
                     "Serializing {Property} ({Type}) with IJsonLoadableState",
                     property.Name,
                     property.PropertyType
@@ -245,7 +257,7 @@ public abstract class LoadableViewModelBase : ViewModelBase, IJsonLoadableState
             }
             else
             {
-                Logger.Trace(
+                Logger.ConditionalTrace(
                     "Serializing {Property} ({Type})",
                     property.Name,
                     property.PropertyType

StabilityMatrix.Avalonia/ViewModels/Dialogs/ImageViewerViewModel.cs

@@ -69,13 +69,13 @@ public partial class ImageViewerViewModel : ContentDialogViewModelBase
     [RelayCommand]
     private void OnNavigateNext()
     {
-        NavigationRequested?.Invoke(this, DirectionalNavigationEventArgs.Up);
+        NavigationRequested?.Invoke(this, DirectionalNavigationEventArgs.Down);
     }
 
     [RelayCommand]
     private void OnNavigatePrevious()
     {
-        NavigationRequested?.Invoke(this, DirectionalNavigationEventArgs.Down);
+        NavigationRequested?.Invoke(this, DirectionalNavigationEventArgs.Up);
     }
 
     [RelayCommand]

StabilityMatrix.Avalonia/ViewModels/Inference/BatchSizeCardViewModel.cs

@@ -8,7 +8,15 @@ namespace StabilityMatrix.Avalonia.ViewModels.Inference;
 [View(typeof(BatchSizeCard))]
 public partial class BatchSizeCardViewModel : LoadableViewModelBase
 {
-    [ObservableProperty] private int batchSize = 1;
+    [ObservableProperty]
+    private int batchSize = 1;
 
-    [ObservableProperty] private int batchCount = 1;
+    [ObservableProperty]
+    private int batchCount = 1;
+
+    [ObservableProperty]
+    private bool isBatchIndexEnabled;
+
+    [ObservableProperty]
+    private int batchIndex = 1;
 }

StabilityMatrix.Avalonia/ViewModels/Inference/InferenceTextToImageViewModel.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Collections.Generic;
 using System.ComponentModel.DataAnnotations;
 using System.Linq;
 using System.Text.Json.Serialization;
@@ -173,12 +174,20 @@ public class InferenceTextToImageViewModel
         var builder = args.Builder;
         var nodes = builder.Nodes;
 
+        if (args.SeedOverride is { } seed)
+        {
+            builder.Connections.Seed = Convert.ToUInt64(seed);
+        }
+        else
+        {
+            builder.Connections.Seed = Convert.ToUInt64(SeedCardViewModel.Seed);
+        }
+
         // Setup empty latent
         builder.SetupLatentSource(BatchSizeCardViewModel, SamplerCardViewModel);
 
         // Setup base stage
         builder.SetupBaseSampler(
-            SeedCardViewModel,
             SamplerCardViewModel,
             PromptCardViewModel,
             ModelCardViewModel,
@@ -210,7 +219,6 @@ public class InferenceTextToImageViewModel
         )
         {
             builder.SetupRefinerSampler(
-                SeedCardViewModel,
                 SamplerCardViewModel,
                 PromptCardViewModel,
                 ModelCardViewModel,
@@ -285,7 +293,7 @@ public class InferenceTextToImageViewModel
                 ComfyNodeBuilder.KSampler(
                     "HiresSampler",
                     builder.Connections.GetRefinerOrBaseModel(),
-                    Convert.ToUInt64(SeedCardViewModel.Seed),
+                    builder.Connections.Seed,
                     HiresSamplerCardViewModel.Steps,
                     HiresSamplerCardViewModel.CfgScale,
                     // Use hires sampler name if not null, otherwise use the normal sampler
@@ -353,7 +361,19 @@ public class InferenceTextToImageViewModel
             seedCard.GenerateNewSeed();
         }
 
-        var buildPromptArgs = new BuildPromptEventArgs { Overrides = overrides };
+        var batches = BatchSizeCardViewModel.BatchCount;
+
+        var batchArgs = new List<ImageGenerationEventArgs>();
+
+        for (var i = 0; i < batches; i++)
+        {
+            var seed = seedCard.Seed + i;
+
+            var buildPromptArgs = new BuildPromptEventArgs
+            {
+                Overrides = overrides,
+                SeedOverride = seed
+            };
             BuildPrompt(buildPromptArgs);
 
             var generationArgs = new ImageGenerationEventArgs
@@ -363,7 +383,7 @@ public class InferenceTextToImageViewModel
                 OutputNodeNames = buildPromptArgs.Builder.Connections.OutputNodeNames.ToArray(),
                 Parameters = new GenerationParameters
                 {
-                Seed = (ulong)seedCard.Seed,
+                    Seed = (ulong)seed,
                     Steps = SamplerCardViewModel.Steps,
                     CfgScale = SamplerCardViewModel.CfgScale,
                     Sampler = SamplerCardViewModel.SelectedSampler?.Name,
@@ -377,10 +397,19 @@ public class InferenceTextToImageViewModel
                     PositivePrompt = PromptCardViewModel.PromptDocument.Text,
                     NegativePrompt = PromptCardViewModel.NegativePromptDocument.Text
                 },
-            Project = InferenceProjectDocument.FromLoadable(this)
+                Project = InferenceProjectDocument.FromLoadable(this),
+                // Only clear output images on the first batch
+                ClearOutputImages = i == 0
             };
 
-        await RunGeneration(generationArgs, cancellationToken);
+            batchArgs.Add(generationArgs);
+        }
+
+        // Run batches
+        foreach (var args in batchArgs)
+        {
+            await RunGeneration(args, cancellationToken);
+        }
     }
 
     /// <inheritdoc />

StabilityMatrix.Avalonia/ViewModels/Inference/SeedCardViewModel.cs

@@ -18,9 +18,8 @@ public partial class SeedCardViewModel : LoadableViewModelBase
     [ObservableProperty]
     private long seed;
 
-    public string RandomizeButtonToolTip => IsRandomizeEnabled 
-        ? "Randomizing Seed on each run"
-        : "Seed is locked";
+    public string RandomizeButtonToolTip =>
+        IsRandomizeEnabled ? "Randomizing Seed on each run" : "Seed is locked";
 
     [RelayCommand]
     public void GenerateNewSeed()
@@ -33,17 +32,15 @@ public partial class SeedCardViewModel : LoadableViewModelBase
     {
         var model = DeserializeModel<SeedCardModel>(state);
 
-        Seed = long.TryParse(model.Seed, out var result) ? result : 0;
+        Seed = model.Seed;
         IsRandomizeEnabled = model.IsRandomizeEnabled;
     }
 
     /// <inheritdoc />
     public override JsonObject SaveStateToJsonObject()
     {
-        return SerializeModel(new SeedCardModel
-        {
-            Seed = Seed.ToString(),
-            IsRandomizeEnabled = IsRandomizeEnabled
-        });
+        return SerializeModel(
+            new SeedCardModel { Seed = Seed, IsRandomizeEnabled = IsRandomizeEnabled }
+        );
     }
 }

StabilityMatrix.Avalonia/ViewModels/Inference/UpscalerCardViewModel.cs

@@ -28,7 +28,7 @@ public partial class UpscalerCardViewModel : LoadableViewModelBase
     private readonly ServiceManager<ViewModelBase> vmFactory;
 
     [ObservableProperty]
-    private double scale = 1;
+    private double scale = 2;
 
     [ObservableProperty]
     private ComfyUpscaler? selectedUpscaler = ComfyUpscaler.Defaults[0];

StabilityMatrix.Avalonia/Views/Dialogs/ImageViewerDialog.axaml

@@ -19,9 +19,9 @@
     mc:Ignorable="d">
 
     <controls:UserControlBase.KeyBindings>
-        <KeyBinding Command="{Binding NavigateNextCommand}" Gesture="Up" />
+        <KeyBinding Command="{Binding NavigateNextCommand}" Gesture="Down" />
         <KeyBinding Command="{Binding NavigateNextCommand}" Gesture="Right" />
-        <KeyBinding Command="{Binding NavigatePreviousCommand}" Gesture="Down" />
+        <KeyBinding Command="{Binding NavigatePreviousCommand}" Gesture="Up" />
         <KeyBinding Command="{Binding NavigatePreviousCommand}" Gesture="Left" />
     </controls:UserControlBase.KeyBindings>
 

StabilityMatrix.Avalonia/Views/MainWindow.axaml.cs

@@ -151,6 +151,8 @@ public partial class MainWindow : AppWindowBase
     }
 
     private void OnUpdateAvailable(object? sender, UpdateInfo? updateInfo)
+    {
+        Dispatcher.UIThread.Post(() =>
         {
             var vm = DataContext as MainWindowViewModel;
 
@@ -163,6 +165,7 @@ public partial class MainWindow : AppWindowBase
                 tip.Subtitle = $"{Compat.AppVersion} -> {updateInfo.Version}";
                 tip.IsOpen = true;
             }
+        });
     }
 
     public void SetDefaultFonts()

StabilityMatrix.Core/Helper/ImageMetadata.cs

@@ -13,8 +13,9 @@ public class ImageMetadata
 {
     private IReadOnlyList<Directory>? Directories { get; set; }
 
-    private static readonly byte[] Idat = { 0x49, 0x44, 0x41, 0x54 };
-    private static readonly byte[] Text = { 0x74, 0x45, 0x58, 0x74 };
+    private static readonly byte[] PngHeader = { 0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A };
+    private static readonly byte[] Idat = "IDAT"u8.ToArray();
+    private static readonly byte[] Text = "tEXt"u8.ToArray();
 
     public static ImageMetadata ParseFile(FilePath path)
     {
@@ -139,12 +140,19 @@ public class ImageMetadata
 
     public static string ReadTextChunk(BinaryReader byteStream, string key)
     {
-        // skip to end of png header stuff
-        byteStream.BaseStream.Position = 0x21;
+        byteStream.BaseStream.Position = 0;
+
+        // Read first 8 bytes and make sure they match the png header
+        if (!byteStream.ReadBytes(8).SequenceEqual(PngHeader))
+        {
+            return string.Empty;
+        }
+
         while (byteStream.BaseStream.Position < byteStream.BaseStream.Length - 4)
         {
             var chunkSize = BitConverter.ToInt32(byteStream.ReadBytes(4).Reverse().ToArray());
             var chunkType = Encoding.UTF8.GetString(byteStream.ReadBytes(4));
+
             if (chunkType == Encoding.UTF8.GetString(Idat))
             {
                 return string.Empty;
@@ -159,6 +167,11 @@ public class ImageMetadata
                     return text[(key.Length + 1)..];
                 }
             }
+            else
+            {
+                // skip chunk data
+                byteStream.BaseStream.Position += chunkSize;
+            }
 
             // skip crc
             byteStream.BaseStream.Position += 4;

StabilityMatrix.Core/Models/Api/Comfy/Nodes/ComfyNodeBuilder.cs

@@ -144,6 +144,25 @@ public class ComfyNodeBuilder
         };
     }
 
+    public static NamedComfyNode<LatentNodeConnection> LatentFromBatch(
+        string name,
+        LatentNodeConnection samples,
+        int batchIndex,
+        int length
+    )
+    {
+        return new NamedComfyNode<LatentNodeConnection>(name)
+        {
+            ClassType = "LatentFromBatch",
+            Inputs = new Dictionary<string, object?>
+            {
+                ["samples"] = samples.Data,
+                ["batch_index"] = batchIndex,
+                ["length"] = length,
+            }
+        };
+    }
+
     public static NamedComfyNode<ImageNodeConnection> ImageUpscaleWithModel(
         string name,
         UpscaleModelNodeConnection upscaleModel,
@@ -632,6 +651,8 @@ public class ComfyNodeBuilder
 
     public class NodeBuilderConnections
     {
+        public ulong Seed { get; set; }
+
         public ModelNodeConnection? BaseModel { get; set; }
         public VAENodeConnection? BaseVAE { get; set; }
         public ClipNodeConnection? BaseClip { get; set; }