Merge pull request #278 from ionite34/fixes

StabilityMatrix.Avalonia/App.axaml

@@ -58,6 +58,7 @@
         <StyleInclude Source="Controls/SelectImageCard.axaml"/>
         <StyleInclude Source="Controls/ImageFolderCard.axaml"/>
         <StyleInclude Source="Controls/SharpenCard.axaml"/>
+        <StyleInclude Source="Controls/FreeUCard.axaml"/>
         <StyleInclude Source="Controls/Paginator.axaml"/>
         
         <Style Selector="DockControl">

StabilityMatrix.Avalonia/App.axaml.cs

@@ -338,6 +338,7 @@ public sealed class App : Application
         services.AddTransient<BatchSizeCardViewModel>();
         services.AddTransient<SelectImageCardViewModel>();
         services.AddTransient<SharpenCardViewModel>();
+        services.AddTransient<FreeUCardViewModel>();
 
         // Dialog factory
         services.AddSingleton<ServiceManager<ViewModelBase>>(
@@ -376,6 +377,7 @@ public sealed class App : Application
                     .Register(provider.GetRequiredService<InferenceConnectionHelpViewModel>)
                     .Register(provider.GetRequiredService<SharpenCardViewModel>)
                     .Register(provider.GetRequiredService<DownloadResourceViewModel>)
+                    .Register(provider.GetRequiredService<FreeUCardViewModel>)
         );
     }
 
@@ -390,6 +392,7 @@ public sealed class App : Application
         services.AddSingleton<CheckpointBrowserPage>();
         services.AddSingleton<ProgressManagerPage>();
         services.AddSingleton<InferencePage>();
+        services.AddSingleton<NewCheckpointsPage>();
 
         // Inference tabs
         services.AddTransient<InferenceTextToImageView>();
@@ -408,7 +411,7 @@ public sealed class App : Application
         services.AddTransient<BatchSizeCard>();
         services.AddTransient<SelectImageCard>();
         services.AddTransient<SharpenCard>();
-        services.AddSingleton<NewCheckpointsPage>();
+        services.AddTransient<FreeUCard>();
 
         // Dialogs
         services.AddTransient<SelectDataDirectoryDialog>();

StabilityMatrix.Avalonia/Controls/AdvancedImageBox.axaml.cs

@@ -487,9 +487,15 @@ public class AdvancedImageBox : TemplatedControl
     #endregion
 
     #region UI Controls
-    public ScrollBar HorizontalScrollBar { get; private set; }
-    public ScrollBar VerticalScrollBar { get; private set; }
-    public ContentPresenter ViewPort { get; private set; }
+
+    [NotNull]
+    public ScrollBar? HorizontalScrollBar { get; private set; }
+
+    [NotNull]
+    public ScrollBar? VerticalScrollBar { get; private set; }
+
+    [NotNull]
+    public ContentPresenter? ViewPort { get; private set; }
 
     // ReSharper disable once ConditionIsAlwaysTrueOrFalseAccordingToNullableAPIContract
     public bool IsViewPortLoaded => ViewPort is not null;

StabilityMatrix.Avalonia/Controls/CodeCompletion/CompletionList.cs

@@ -502,7 +502,11 @@ public class CompletionList : TemplatedControl
         if (string.IsNullOrEmpty(query))
             return;
 
-        var suggestedIndex = _listBox.SelectedIndex;
+        var suggestedIndex = _listBox?.SelectedIndex ?? -1;
+        if (suggestedIndex == -1)
+        {
+            return;
+        }
 
         var bestIndex = -1;
         var bestQuality = -1;

StabilityMatrix.Avalonia/Controls/DropTargetTemplatedControlBase.cs

@@ -1,10 +1,9 @@
-using Avalonia.Controls.Primitives;
-using Avalonia.Input;
+using Avalonia.Input;
 using StabilityMatrix.Avalonia.ViewModels;
 
 namespace StabilityMatrix.Avalonia.Controls;
 
-public abstract class DropTargetTemplatedControlBase : TemplatedControl
+public abstract class DropTargetTemplatedControlBase : TemplatedControlBase
 {
     protected DropTargetTemplatedControlBase()
     {

StabilityMatrix.Avalonia/Controls/FreeUCard.axaml

@@ -0,0 +1,75 @@
+<Styles xmlns="https://github.com/avaloniaui"
+        xmlns:x="http://schemas.microsoft.com/winfx/2006/xaml"
+        xmlns:controls="using:StabilityMatrix.Avalonia.Controls"
+        xmlns:ui="clr-namespace:FluentAvalonia.UI.Controls;assembly=FluentAvalonia"
+        xmlns:mocks="clr-namespace:StabilityMatrix.Avalonia.DesignData"
+        xmlns:vmInference="clr-namespace:StabilityMatrix.Avalonia.ViewModels.Inference"
+        x:DataType="vmInference:FreeUCardViewModel">
+    <Design.PreviewWith>
+        <controls:FreeUCard DataContext="{x:Static mocks:DesignData.FreeUCardViewModel}" />
+    </Design.PreviewWith>
+
+    <Style Selector="controls|FreeUCard">
+        <Setter Property="Template">
+            <ControlTemplate>
+                <controls:Card x:Name="PART_Card">
+                    <controls:Card.Styles>
+                        <Style Selector="ui|NumberBox">
+                            <Setter Property="Margin" Value="12,4,0,4"/>
+                            <Setter Property="MinWidth" Value="70"/>
+                            <Setter Property="HorizontalAlignment" Value="Stretch"/>
+                            <Setter Property="ValidationMode" Value="InvalidInputOverwritten"/>
+                            <Setter Property="SmallChange" Value="0.01"/>
+                            <Setter Property="LargeChange" Value="0.1"/>
+                            <Setter Property="SimpleNumberFormat" Value="F2"/>
+                            <Setter Property="SpinButtonPlacementMode" Value="Inline"/>
+                        </Style>
+                    </controls:Card.Styles>
+                    <Grid RowDefinitions="*,*,*,*" ColumnDefinitions="Auto,*">
+                        <TextBlock
+                            Grid.Row="0"
+                            Grid.Column="0"
+                            Text="B₁"
+                            VerticalAlignment="Center" />
+                        <ui:NumberBox
+                            Grid.Row="0"
+                            Grid.Column="1"
+                            Margin="12,0,0,4"
+                            Value="{Binding B1}"/>
+                        
+                        <TextBlock
+                            Grid.Row="1"
+                            Grid.Column="0"
+                            Text="B₂"
+                            VerticalAlignment="Center" />
+                        <ui:NumberBox
+                            Grid.Row="1"
+                            Grid.Column="1"
+                            Value="{Binding B2}"/>
+                        
+                        <TextBlock
+                            Grid.Row="2"
+                            Grid.Column="0"
+                            Text="S₁"
+                            VerticalAlignment="Center" />
+                        <ui:NumberBox
+                            Grid.Row="2"
+                            Grid.Column="1"
+                            Value="{Binding S1}"/>
+                        
+                        <TextBlock
+                            Grid.Row="3"
+                            Grid.Column="0"
+                            Text="S₂"
+                            VerticalAlignment="Center" />
+                        <ui:NumberBox
+                            Grid.Row="3"
+                            Grid.Column="1"
+                            Margin="12,4,0,0"
+                            Value="{Binding S2}"/>
+                    </Grid>
+                </controls:Card>
+            </ControlTemplate>
+        </Setter>
+    </Style>
+</Styles>

StabilityMatrix.Avalonia/Controls/FreeUCard.axaml.cs

@@ -0,0 +1,5 @@
+using Avalonia.Controls.Primitives;
+
+namespace StabilityMatrix.Avalonia.Controls;
+
+public class FreeUCard : TemplatedControl { }

StabilityMatrix.Avalonia/Controls/ImageFolderCard.axaml.cs

@@ -1,31 +1,9 @@
-using AsyncAwaitBestPractices;
-using Avalonia.Controls;
-using Avalonia.Input;
-using Avalonia.Interactivity;
-using Avalonia.Threading;
-using StabilityMatrix.Avalonia.ViewModels.Base;
+using Avalonia.Input;
 
 namespace StabilityMatrix.Avalonia.Controls;
 
 public class ImageFolderCard : DropTargetTemplatedControlBase
 {
-    /// <inheritdoc />
-    protected override void OnLoaded(RoutedEventArgs e)
-    {
-        base.OnLoaded(e);
-
-        if (DataContext is ViewModelBase vm)
-        {
-            vm.OnLoaded();
-            Dispatcher.UIThread
-                .InvokeAsync(async () =>
-                {
-                    await vm.OnLoadedAsync();
-                })
-                .SafeFireAndForget();
-        }
-    }
-
     /// <inheritdoc />
     protected override void DropHandler(object? sender, DragEventArgs e)
     {

StabilityMatrix.Avalonia/Controls/TemplatedControlBase.cs

@@ -0,0 +1,40 @@
+using AsyncAwaitBestPractices;
+using Avalonia.Controls.Primitives;
+using Avalonia.Interactivity;
+using Avalonia.Threading;
+using StabilityMatrix.Avalonia.ViewModels.Base;
+
+namespace StabilityMatrix.Avalonia.Controls;
+
+public abstract class TemplatedControlBase : TemplatedControl
+{
+    /// <inheritdoc />
+    protected override void OnLoaded(RoutedEventArgs e)
+    {
+        base.OnLoaded(e);
+
+        if (DataContext is not ViewModelBase viewModel)
+            return;
+
+        // Run synchronous load then async load
+        viewModel.OnLoaded();
+
+        // Can't block here so we'll run as async on UI thread
+        Dispatcher.UIThread.InvokeAsync(viewModel.OnLoadedAsync).SafeFireAndForget();
+    }
+
+    /// <inheritdoc />
+    protected override void OnUnloaded(RoutedEventArgs e)
+    {
+        base.OnUnloaded(e);
+
+        if (DataContext is not ViewModelBase viewModel)
+            return;
+
+        // Run synchronous load then async load
+        viewModel.OnUnloaded();
+
+        // Can't block here so we'll run as async on UI thread
+        Dispatcher.UIThread.InvokeAsync(viewModel.OnUnloadedAsync).SafeFireAndForget();
+    }
+}

StabilityMatrix.Avalonia/Controls/UserControlBase.cs

@@ -12,40 +12,34 @@ public class UserControlBase : UserControl
 {
     static UserControlBase()
     {
-        LoadedEvent.AddClassHandler<UserControlBase>(
-            (cls, args) => cls.OnLoadedEvent(args));
+        LoadedEvent.AddClassHandler<UserControlBase>((cls, args) => cls.OnLoadedEvent(args));
 
-        UnloadedEvent.AddClassHandler<UserControlBase>(
-            (cls, args) => cls.OnUnloadedEvent(args));
+        UnloadedEvent.AddClassHandler<UserControlBase>((cls, args) => cls.OnUnloadedEvent(args));
     }
 
     // ReSharper disable once UnusedParameter.Global
     protected virtual void OnLoadedEvent(RoutedEventArgs? e)
     {
-        if (DataContext is not ViewModelBase viewModel) return;
-        
+        if (DataContext is not ViewModelBase viewModel)
+            return;
+
         // Run synchronous load then async load
         viewModel.OnLoaded();
-        
+
         // Can't block here so we'll run as async on UI thread
-        Dispatcher.UIThread.InvokeAsync(async () =>
-        {
-            await viewModel.OnLoadedAsync();
-        }).SafeFireAndForget();
+        Dispatcher.UIThread.InvokeAsync(viewModel.OnLoadedAsync).SafeFireAndForget();
     }
-    
+
     // ReSharper disable once UnusedParameter.Global
     protected virtual void OnUnloadedEvent(RoutedEventArgs? e)
     {
-        if (DataContext is not ViewModelBase viewModel) return;
-        
+        if (DataContext is not ViewModelBase viewModel)
+            return;
+
         // Run synchronous load then async load
         viewModel.OnUnloaded();
-        
+
         // Can't block here so we'll run as async on UI thread
-        Dispatcher.UIThread.InvokeAsync(async () =>
-        {
-            await viewModel.OnUnloadedAsync();
-        }).SafeFireAndForget();
+        Dispatcher.UIThread.InvokeAsync(viewModel.OnUnloadedAsync).SafeFireAndForget();
     }
 }

StabilityMatrix.Avalonia/DesignData/DesignData.cs

@@ -8,6 +8,7 @@ using System.Net.Http;
 using System.Text;
 using AvaloniaEdit.Utils;
 using Microsoft.Extensions.DependencyInjection;
+using StabilityMatrix.Avalonia.Controls;
 using StabilityMatrix.Avalonia.Controls.CodeCompletion;
 using StabilityMatrix.Avalonia.Models;
 using StabilityMatrix.Avalonia.Models.TagCompletion;
@@ -539,6 +540,8 @@ The gallery images are often inpainted, but you will get something very similar
     public static ImageFolderCardViewModel ImageFolderCardViewModel =>
         DialogFactory.Get<ImageFolderCardViewModel>();
 
+    public static FreeUCardViewModel FreeUCardViewModel => DialogFactory.Get<FreeUCardViewModel>();
+
     public static PromptCardViewModel PromptCardViewModel =>
         DialogFactory.Get<PromptCardViewModel>(vm =>
         {

StabilityMatrix.Avalonia/DesignData/MockLiteDbContext.cs

@@ -47,7 +47,7 @@ public class MockLiteDbContext : ILiteDbContext
 
     public Task<GithubCacheEntry?> GetGithubCacheEntry(string cacheKey)
     {
-        return Task.FromResult<GithubCacheEntry>(null);
+        return Task.FromResult<GithubCacheEntry?>(null);
     }
 
     public Task<bool> UpsertGithubCacheEntry(GithubCacheEntry cacheEntry)

StabilityMatrix.Avalonia/Extensions/ComfyNodeBuilderExtensions.cs

@@ -39,7 +39,8 @@ public static class ComfyNodeBuilderExtensions
         SamplerCardViewModel samplerCardViewModel,
         PromptCardViewModel promptCardViewModel,
         ModelCardViewModel modelCardViewModel,
-        IModelIndexService modelIndexService
+        IModelIndexService modelIndexService,
+        Action<ComfyNodeBuilder>? postModelLoad = null
     )
     {
         // Load base checkpoint
@@ -51,11 +52,12 @@ public static class ComfyNodeBuilderExtensions
             )
         );
 
+        builder.Connections.BaseModel = checkpointLoader.GetOutput<ModelNodeConnection>(0);
+        builder.Connections.BaseClip = checkpointLoader.GetOutput<ClipNodeConnection>(1);
         builder.Connections.BaseVAE = checkpointLoader.GetOutput<VAENodeConnection>(2);
 
-        // Define model and clip for connections for chaining
-        var modelSource = checkpointLoader.GetOutput<ModelNodeConnection>(0);
-        var clipSource = checkpointLoader.GetOutput<ClipNodeConnection>(1);
+        // Run post model load action
+        postModelLoad?.Invoke(builder);
 
         // Load prompts
         var prompt = promptCardViewModel.GetPrompt();
@@ -69,25 +71,28 @@ public static class ComfyNodeBuilderExtensions
             // Convert to local file names
             var lorasGroup = builder.Group_LoraLoadMany(
                 "Loras",
-                modelSource,
-                clipSource,
+                builder.Connections.BaseModel,
+                builder.Connections.BaseClip,
                 prompt.GetExtraNetworksAsLocalModels(modelIndexService)
             );
 
             // Set as source
-            modelSource = lorasGroup.Output1;
-            clipSource = lorasGroup.Output2;
+            builder.Connections.BaseModel = lorasGroup.Output1;
+            builder.Connections.BaseClip = lorasGroup.Output2;
         }
-        builder.Connections.BaseModel = modelSource;
 
         // Clips
         var positiveClip = builder.Nodes.AddNamedNode(
-            ComfyNodeBuilder.ClipTextEncode("PositiveCLIP", clipSource, prompt.ProcessedText)
+            ComfyNodeBuilder.ClipTextEncode(
+                "PositiveCLIP",
+                builder.Connections.BaseClip,
+                prompt.ProcessedText
+            )
         );
         var negativeClip = builder.Nodes.AddNamedNode(
             ComfyNodeBuilder.ClipTextEncode(
                 "NegativeCLIP",
-                clipSource,
+                builder.Connections.BaseClip,
                 negativePrompt.ProcessedText
             )
         );
@@ -103,7 +108,7 @@ public static class ComfyNodeBuilderExtensions
             var sampler = builder.Nodes.AddNamedNode(
                 ComfyNodeBuilder.KSampler(
                     "Sampler",
-                    modelSource,
+                    builder.Connections.BaseModel,
                     Convert.ToUInt64(seedCardViewModel.Seed),
                     samplerCardViewModel.Steps,
                     samplerCardViewModel.CfgScale,
@@ -129,7 +134,7 @@ public static class ComfyNodeBuilderExtensions
             var sampler = builder.Nodes.AddNamedNode(
                 ComfyNodeBuilder.KSamplerAdvanced(
                     "Sampler",
-                    modelSource,
+                    builder.Connections.BaseModel,
                     true,
                     Convert.ToUInt64(seedCardViewModel.Seed),
                     totalSteps,
@@ -157,7 +162,8 @@ public static class ComfyNodeBuilderExtensions
         SamplerCardViewModel samplerCardViewModel,
         PromptCardViewModel promptCardViewModel,
         ModelCardViewModel modelCardViewModel,
-        IModelIndexService modelIndexService
+        IModelIndexService modelIndexService,
+        Action<ComfyNodeBuilder>? postModelLoad = null
     )
     {
         // Load refiner checkpoint
@@ -169,11 +175,12 @@ public static class ComfyNodeBuilderExtensions
             )
         );
 
+        builder.Connections.RefinerModel = checkpointLoader.GetOutput<ModelNodeConnection>(0);
+        builder.Connections.RefinerClip = checkpointLoader.GetOutput<ClipNodeConnection>(1);
         builder.Connections.RefinerVAE = checkpointLoader.GetOutput<VAENodeConnection>(2);
 
-        // Define model and clip for connections for chaining
-        var modelSource = checkpointLoader.GetOutput<ModelNodeConnection>(0);
-        var clipSource = checkpointLoader.GetOutput<ClipNodeConnection>(1);
+        // Run post model load action
+        postModelLoad?.Invoke(builder);
 
         // Load prompts
         var prompt = promptCardViewModel.GetPrompt();
@@ -187,29 +194,28 @@ public static class ComfyNodeBuilderExtensions
             // Convert to local file names
             var lorasGroup = builder.Group_LoraLoadMany(
                 "Refiner_Loras",
-                modelSource,
-                clipSource,
+                builder.Connections.RefinerModel,
+                builder.Connections.RefinerClip,
                 prompt.GetExtraNetworksAsLocalModels(modelIndexService)
             );
 
             // Set as source
-            modelSource = lorasGroup.Output1;
-            clipSource = lorasGroup.Output2;
+            builder.Connections.RefinerModel = lorasGroup.Output1;
+            builder.Connections.RefinerClip = lorasGroup.Output2;
         }
-        builder.Connections.RefinerModel = modelSource;
 
         // Clips
         var positiveClip = builder.Nodes.AddNamedNode(
             ComfyNodeBuilder.ClipTextEncode(
                 "Refiner_PositiveCLIP",
-                clipSource,
+                builder.Connections.RefinerClip,
                 prompt.ProcessedText
             )
         );
         var negativeClip = builder.Nodes.AddNamedNode(
             ComfyNodeBuilder.ClipTextEncode(
                 "Refiner_NegativeCLIP",
-                clipSource,
+                builder.Connections.RefinerClip,
                 negativePrompt.ProcessedText
             )
         );
@@ -224,7 +230,7 @@ public static class ComfyNodeBuilderExtensions
         var sampler = builder.Nodes.AddNamedNode(
             ComfyNodeBuilder.KSamplerAdvanced(
                 "Refiner_Sampler",
-                modelSource,
+                builder.Connections.RefinerModel,
                 false,
                 Convert.ToUInt64(seedCardViewModel.Seed),
                 totalSteps,

StabilityMatrix.Avalonia/Models/ILoadableState.cs

@@ -1,26 +0,0 @@
-using System;
-using System.Text.Json;
-using System.Text.Json.Nodes;
-
-namespace StabilityMatrix.Avalonia.Models;
-
-public interface ILoadableState<T> : IJsonLoadableState
-{
-    new Type LoadableStateType => typeof(T);
-    
-    void LoadState(T state);
-
-    new void LoadStateFromJsonObject(JsonObject state)
-    {
-        state.Deserialize(LoadableStateType);
-    }
-    
-    T SaveState();
-    
-    new JsonObject SaveStateToJsonObject()
-    {
-        var node = JsonSerializer.SerializeToNode(SaveState());
-        return node?.AsObject() ?? throw new 
-            InvalidOperationException("Failed to serialize state to JSON object.");
-    }
-}

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-pre.2</Version>
+        <Version>2.5.0-pre.3</Version>
         <InformationalVersion>$(Version)</InformationalVersion>
         <EnableWindowsTargeting>true</EnableWindowsTargeting>
         <AllowUnsafeBlocks>true</AllowUnsafeBlocks>

StabilityMatrix.Avalonia/ViewModels/Inference/FreeUCardViewModel.cs

@@ -0,0 +1,31 @@
+using System.ComponentModel.DataAnnotations;
+using CommunityToolkit.Mvvm.ComponentModel;
+using StabilityMatrix.Avalonia.Controls;
+using StabilityMatrix.Avalonia.ViewModels.Base;
+using StabilityMatrix.Core.Attributes;
+
+namespace StabilityMatrix.Avalonia.ViewModels.Inference;
+
+[View(typeof(FreeUCard))]
+public partial class FreeUCardViewModel : LoadableViewModelBase
+{
+    [ObservableProperty]
+    [Required]
+    [Range(0D, 10D)]
+    private double b1 = 1.1;
+
+    [ObservableProperty]
+    [Required]
+    [Range(0D, 10D)]
+    private double b2 = 1.2;
+
+    [ObservableProperty]
+    [Required]
+    [Range(0D, 10D)]
+    private double s1 = 0.9;
+
+    [ObservableProperty]
+    [Required]
+    [Range(0D, 10D)]
+    private double s2 = 0.2;
+}

StabilityMatrix.Avalonia/ViewModels/Inference/ImageFolderCardViewModel.cs

@@ -6,9 +6,6 @@ using AsyncAwaitBestPractices;
 using AsyncImageLoader;
 using Avalonia;
 using Avalonia.Controls.Notifications;
-using Avalonia.Controls.Primitives;
-using Avalonia.Input;
-using Avalonia.Layout;
 using Avalonia.Platform.Storage;
 using Avalonia.Threading;
 using CommunityToolkit.Mvvm.ComponentModel;
@@ -17,7 +14,6 @@ using DynamicData;
 using DynamicData.Binding;
 using FuzzySharp;
 using FuzzySharp.PreProcess;
-using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Logging;
 using SkiaSharp;
 using StabilityMatrix.Avalonia.Controls;
@@ -27,8 +23,6 @@ using StabilityMatrix.Avalonia.Models;
 using StabilityMatrix.Avalonia.Services;
 using StabilityMatrix.Avalonia.ViewModels.Base;
 using StabilityMatrix.Avalonia.ViewModels.Dialogs;
-using StabilityMatrix.Avalonia.Views;
-using StabilityMatrix.Avalonia.Views.Dialogs;
 using StabilityMatrix.Core.Attributes;
 using StabilityMatrix.Core.Models.Database;
 using StabilityMatrix.Core.Models.FileInterfaces;

StabilityMatrix.Avalonia/ViewModels/Inference/InferenceTextToImageViewModel.cs

@@ -55,24 +55,33 @@ public class InferenceTextToImageViewModel
     [JsonPropertyName("HiresUpscaler")]
     public UpscalerCardViewModel HiresUpscalerCardViewModel { get; }
 
+    [JsonPropertyName("FreeU")]
+    public FreeUCardViewModel FreeUCardViewModel { get; }
+
     [JsonPropertyName("BatchSize")]
     public BatchSizeCardViewModel BatchSizeCardViewModel { get; }
 
     [JsonPropertyName("Seed")]
     public SeedCardViewModel SeedCardViewModel { get; }
 
-    public bool IsHiresFixEnabled
+    public bool IsFreeUEnabled
     {
         get => StackCardViewModel.GetCard<StackExpanderViewModel>().IsEnabled;
         set => StackCardViewModel.GetCard<StackExpanderViewModel>().IsEnabled = value;
     }
 
-    public bool IsUpscaleEnabled
+    public bool IsHiresFixEnabled
     {
         get => StackCardViewModel.GetCard<StackExpanderViewModel>(1).IsEnabled;
         set => StackCardViewModel.GetCard<StackExpanderViewModel>(1).IsEnabled = value;
     }
 
+    public bool IsUpscaleEnabled
+    {
+        get => StackCardViewModel.GetCard<StackExpanderViewModel>(2).IsEnabled;
+        set => StackCardViewModel.GetCard<StackExpanderViewModel>(2).IsEnabled = value;
+    }
+
     public InferenceTextToImageViewModel(
         INotificationService notificationService,
         IInferenceClientManager inferenceClientManager,
@@ -106,6 +115,7 @@ public class InferenceTextToImageViewModel
         });
         HiresUpscalerCardViewModel = vmFactory.Get<UpscalerCardViewModel>();
         UpscalerCardViewModel = vmFactory.Get<UpscalerCardViewModel>();
+        FreeUCardViewModel = vmFactory.Get<FreeUCardViewModel>();
         BatchSizeCardViewModel = vmFactory.Get<BatchSizeCardViewModel>();
 
         StackCardViewModel = vmFactory.Get<StackCardViewModel>();
@@ -115,6 +125,12 @@ public class InferenceTextToImageViewModel
             {
                 ModelCardViewModel,
                 SamplerCardViewModel,
+                // Free U
+                vmFactory.Get<StackExpanderViewModel>(stackExpander =>
+                {
+                    stackExpander.Title = "FreeU";
+                    stackExpander.AddCards(new LoadableViewModelBase[] { FreeUCardViewModel });
+                }),
                 // Hires Fix
                 vmFactory.Get<StackExpanderViewModel>(stackExpander =>
                 {
@@ -166,7 +182,25 @@ public class InferenceTextToImageViewModel
             SamplerCardViewModel,
             PromptCardViewModel,
             ModelCardViewModel,
-            modelIndexService
+            modelIndexService,
+            postModelLoad: x =>
+            {
+                if (IsFreeUEnabled)
+                {
+                    builder.Connections.BaseModel = nodes
+                        .AddNamedNode(
+                            ComfyNodeBuilder.FreeU(
+                                "FreeU",
+                                x.Connections.BaseModel!,
+                                FreeUCardViewModel.B1,
+                                FreeUCardViewModel.B2,
+                                FreeUCardViewModel.S1,
+                                FreeUCardViewModel.S2
+                            )
+                        )
+                        .Output;
+                }
+            }
         );
 
         // Setup refiner stage if enabled
@@ -180,7 +214,25 @@ public class InferenceTextToImageViewModel
                 SamplerCardViewModel,
                 PromptCardViewModel,
                 ModelCardViewModel,
-                modelIndexService
+                modelIndexService,
+                postModelLoad: x =>
+                {
+                    if (IsFreeUEnabled)
+                    {
+                        builder.Connections.RefinerModel = nodes
+                            .AddNamedNode(
+                                ComfyNodeBuilder.FreeU(
+                                    "Refiner_FreeU",
+                                    x.Connections.RefinerModel!,
+                                    FreeUCardViewModel.B1,
+                                    FreeUCardViewModel.B2,
+                                    FreeUCardViewModel.S1,
+                                    FreeUCardViewModel.S2
+                                )
+                            )
+                            .Output;
+                    }
+                }
             );
         }
 

StabilityMatrix.Avalonia/ViewModels/SettingsViewModel.cs

@@ -238,7 +238,7 @@ public partial class SettingsViewModel : PageViewModelBase
 
     partial void OnSelectedLanguageChanged(CultureInfo? oldValue, CultureInfo newValue)
     {
-        if (oldValue is null || newValue.Name == Cultures.Current.Name)
+        if (oldValue is null || newValue.Name == Cultures.Current?.Name)
             return;
 
         // Set locale

StabilityMatrix.Avalonia/Views/Inference/InferenceTextToImageView.axaml

@@ -210,7 +210,7 @@
                             
                             <Grid
                                 x:CompileBindings="False"
-                                DataContext="{Binding RelativeSource={RelativeSource Mode=FindAncestor, AncestorType=DockControl}, Path=DataContext}">
+                                DataContext="{Binding ElementName=Dock, Path=DataContext}">
                                 <controls:ImageFolderCard
                                     DataContext="{Binding ImageFolderCardViewModel}" />
                             </Grid>

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

@@ -240,6 +240,29 @@ public class ComfyNodeBuilder
         };
     }
 
+    public static NamedComfyNode<ModelNodeConnection> FreeU(
+        string name,
+        ModelNodeConnection model,
+        double b1,
+        double b2,
+        double s1,
+        double s2
+    )
+    {
+        return new NamedComfyNode<ModelNodeConnection>(name)
+        {
+            ClassType = "FreeU",
+            Inputs = new Dictionary<string, object?>
+            {
+                ["model"] = model.Data,
+                ["b1"] = b1,
+                ["b2"] = b2,
+                ["s1"] = s1,
+                ["s2"] = s2
+            }
+        };
+    }
+
     public static NamedComfyNode<ConditioningNodeConnection> ClipTextEncode(
         string name,
         ClipNodeConnection clip,
@@ -611,12 +634,14 @@ public class ComfyNodeBuilder
     {
         public ModelNodeConnection? BaseModel { get; set; }
         public VAENodeConnection? BaseVAE { get; set; }
+        public ClipNodeConnection? BaseClip { get; set; }
 
         public ConditioningNodeConnection? BaseConditioning { get; set; }
         public ConditioningNodeConnection? BaseNegativeConditioning { get; set; }
 
         public ModelNodeConnection? RefinerModel { get; set; }
         public VAENodeConnection? RefinerVAE { get; set; }
+        public ClipNodeConnection? RefinerClip { get; set; }
         public ConditioningNodeConnection? RefinerConditioning { get; set; }
         public ConditioningNodeConnection? RefinerNegativeConditioning { get; set; }
 

StabilityMatrix.Core/Models/Packages/BaseGitPackage.cs

@@ -239,8 +239,13 @@ public abstract class BaseGitPackage : BasePackage
     public override async Task<bool> CheckForUpdates(InstalledPackage package)
     {
         var currentVersion = package.Version;
-        if (currentVersion == null)
+        if (currentVersion is null or { InstalledReleaseVersion: null, InstalledBranch: null })
         {
+            Logger.Warn(
+                "Could not check updates for package {Name}, version is invalid: {@currentVersion}",
+                Name,
+                currentVersion
+            );
             return false;
         }
 
@@ -254,7 +259,7 @@ public abstract class BaseGitPackage : BasePackage
             }
 
             var allCommits = (
-                await GetAllCommits(currentVersion.InstalledBranch).ConfigureAwait(false)
+                await GetAllCommits(currentVersion.InstalledBranch!).ConfigureAwait(false)
             )?.ToList();
             if (allCommits == null || !allCommits.Any())
             {

StabilityMatrix.Core/Models/Packages/ComfyUI.cs

@@ -455,24 +455,4 @@ public class ComfyUI : BaseGitPackage
             })
             .ConfigureAwait(false);
     }
-
-    public class ComfyModelPathsYaml
-    {
-        public class SmData
-        {
-            public string Checkpoints { get; set; }
-            public string Vae { get; set; }
-            public string Loras { get; set; }
-            public string UpscaleModels { get; set; }
-            public string Embeddings { get; set; }
-            public string Hypernetworks { get; set; }
-            public string Controlnet { get; set; }
-            public string Clip { get; set; }
-            public string Diffusers { get; set; }
-            public string Gligen { get; set; }
-            public string VaeApprox { get; set; }
-        }
-
-        public SmData? StabilityMatrix { get; set; }
-    }
 }

StabilityMatrix.Core/Models/Packages/PackageVersionOptions.cs

@@ -2,6 +2,8 @@
 
 public class PackageVersionOptions
 {
-    public IEnumerable<PackageVersion>? AvailableVersions { get; set; }
-    public IEnumerable<PackageVersion>? AvailableBranches { get; set; }
+    public IEnumerable<PackageVersion> AvailableVersions { get; set; } =
+        Enumerable.Empty<PackageVersion>();
+    public IEnumerable<PackageVersion> AvailableBranches { get; set; } =
+        Enumerable.Empty<PackageVersion>();
 }