Merge pull request #519 from ionite34/model-configs

Inference - Model config support

CHANGELOG.md

@@ -5,6 +5,11 @@ 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.9.0-dev.2
+### Added
+#### Inference
+- Added option to load a .yaml config file next to the model with the same name. Can be used with VPred and other models that require a config file.
+
 ## v2.9.0-dev.1
 ### Added
 - Added new package: [StableSwarmUI](https://github.com/Stability-AI/StableSwarmUI) by Stability AI

StabilityMatrix.Avalonia/Controls/Inference/ModelCard.axaml

@@ -1,14 +1,16 @@
 <Styles
     xmlns="https://github.com/avaloniaui"
     xmlns:x="http://schemas.microsoft.com/winfx/2006/xaml"
+    xmlns:avalonia="https://github.com/projektanker/icons.avalonia"
     xmlns:controls="using:StabilityMatrix.Avalonia.Controls"
+    xmlns:converters="clr-namespace:StabilityMatrix.Avalonia.Converters"
+    xmlns:fluent="clr-namespace:FluentIcons.Avalonia.Fluent;assembly=FluentIcons.Avalonia.Fluent"
     xmlns:inference="clr-namespace:StabilityMatrix.Avalonia.ViewModels.Inference"
     xmlns:lang="clr-namespace:StabilityMatrix.Avalonia.Languages"
     xmlns:mocks="clr-namespace:StabilityMatrix.Avalonia.DesignData"
     xmlns:models="clr-namespace:StabilityMatrix.Core.Models;assembly=StabilityMatrix.Core"
-    xmlns:ui="clr-namespace:FluentAvalonia.UI.Controls;assembly=FluentAvalonia"
-    xmlns:converters="clr-namespace:StabilityMatrix.Avalonia.Converters"
     xmlns:sg="clr-namespace:SpacedGridControl.Avalonia;assembly=SpacedGridControl.Avalonia"
+    xmlns:ui="clr-namespace:FluentAvalonia.UI.Controls;assembly=FluentAvalonia"
     x:DataType="inference:ModelCardViewModel">
     <Design.PreviewWith>
         <Panel Width="400" Height="200">
@@ -19,44 +21,73 @@
         </Panel>
     </Design.PreviewWith>
 
+    <Styles.Resources>
+        <fluent:SymbolIconSource
+            x:Key="IconCube"
+            IsFilled="False"
+            Symbol="Cube" />
+        <fluent:SymbolIconSource
+            x:Key="IconQuestionCircle"
+            IsFilled="True"
+            Symbol="QuestionCircle" />
+        <fluent:SymbolIconSource
+            x:Key="IconTableCellEdit"
+            IsFilled="True"
+            Symbol="TableCellEdit" />
+    </Styles.Resources>
+
     <Style Selector="controls|ModelCard">
         <Setter Property="HorizontalAlignment" Value="Stretch" />
         <Setter Property="Template">
             <ControlTemplate>
                 <controls:Card Padding="12">
                     <sg:SpacedGrid
+                        ColumnDefinitions="Auto,*,Auto"
                         ColumnSpacing="8"
-                        RowSpacing="0"
+                        RowDefinitions="*,*,*,*"
-                        ColumnDefinitions="Auto,*,Auto" 
+                        RowSpacing="0">
-                        RowDefinitions="*,*,*,*">
                         <!--  Model  -->
                         <TextBlock
                             Grid.Column="0"
                             VerticalAlignment="Center"
                             Text="{x:Static lang:Resources.Label_Model}"
                             TextAlignment="Left" />
-                        
+
                         <controls:BetterComboBox
                             Grid.Row="0"
                             Grid.Column="1"
                             Padding="8,6,4,6"
-                            Theme="{StaticResource  BetterComboBoxHybridModelTheme}"
                             HorizontalAlignment="Stretch"
                             ItemsSource="{Binding ClientManager.Models}"
-                            SelectedItem="{Binding SelectedModel}"/>
+                            SelectedItem="{Binding SelectedModel}"
+                            Theme="{StaticResource BetterComboBoxHybridModelTheme}" />
 
                         <Button
                             Grid.Row="0"
                             Grid.Column="2"
-                            VerticalAlignment="Stretch"
+                            Margin="0,0,0,0"
                             HorizontalAlignment="Stretch"
-                            Margin="0,0,0,0">
+                            VerticalAlignment="Stretch">
                             <ui:SymbolIcon FontSize="16" Symbol="Setting" />
                             <Button.Flyout>
                                 <ui:FAMenuFlyout Placement="BottomEdgeAlignedLeft">
-                                    <ui:ToggleMenuFlyoutItem IsChecked="{Binding IsRefinerSelectionEnabled}" Text="{x:Static lang:Resources.Label_Refiner}" />
+                                    <ui:ToggleMenuFlyoutItem
-                                    <ui:ToggleMenuFlyoutItem IsChecked="{Binding IsVaeSelectionEnabled}" Text="{x:Static lang:Resources.Label_VAE}" />
+                                        IconSource="{StaticResource IconCube}"
-                                    <ui:ToggleMenuFlyoutItem IsChecked="{Binding IsClipSkipEnabled}" Text="{x:Static lang:Resources.Label_CLIPSkip}" />
+                                        IsChecked="{Binding IsRefinerSelectionEnabled}"
+                                        Text="{x:Static lang:Resources.Label_Refiner}" />
+                                    <ui:ToggleMenuFlyoutItem
+                                        IconSource="{StaticResource IconCube}"
+                                        IsChecked="{Binding IsVaeSelectionEnabled}"
+                                        Text="{x:Static lang:Resources.Label_VAE}" />
+                                    <ui:ToggleMenuFlyoutItem
+                                        IconSource="{StaticResource IconTableCellEdit}"
+                                        IsChecked="{Binding IsClipSkipEnabled}"
+                                        Text="{x:Static lang:Resources.Label_CLIPSkip}" />
+                                    <ui:MenuFlyoutSeparator />
+                                    <ui:MenuFlyoutItem
+                                        Command="{Binding ConfigClickCommand}"
+                                        IconSource="{StaticResource IconQuestionCircle}"
+                                        Text="{x:Static lang:Resources.Label_Config}" />
                                 </ui:FAMenuFlyout>
                             </Button.Flyout>
                         </Button>
@@ -77,12 +108,12 @@
                             Grid.ColumnSpan="2"
                             Margin="0,8,0,0"
                             Padding="8,6,4,6"
-                            IsVisible="{Binding IsRefinerSelectionEnabled}"
-                            Theme="{StaticResource  BetterComboBoxHybridModelTheme}"
                             HorizontalAlignment="Stretch"
+                            IsVisible="{Binding IsRefinerSelectionEnabled}"
                             ItemsSource="{Binding ClientManager.Models}"
-                            SelectedItem="{Binding SelectedRefiner}"/>
+                            SelectedItem="{Binding SelectedRefiner}"
-                        
+                            Theme="{StaticResource BetterComboBoxHybridModelTheme}" />
+
                         <!--  VAE  -->
                         <TextBlock
                             Grid.Row="2"
@@ -104,7 +135,7 @@
                             IsVisible="{Binding IsVaeSelectionEnabled}"
                             ItemsSource="{Binding ClientManager.VaeModels}"
                             SelectedItem="{Binding SelectedVae}" />
-                        
+
                         <!--  CLIP Skip  -->
                         <TextBlock
                             Grid.Row="3"
@@ -119,16 +150,16 @@
                             Grid.Row="3"
                             Grid.Column="1"
                             Grid.ColumnSpan="2"
-                            IsVisible="{Binding IsClipSkipEnabled}"
-                            Watermark="1"
-                            HorizontalAlignment="Stretch"
                             Margin="0,8,0,0"
-                            Minimum="1"
+                            HorizontalAlignment="Stretch"
-                            Maximum="24"
+                            ClipValueToMinMax="True"
                             Increment="1"
+                            IsVisible="{Binding IsClipSkipEnabled}"
+                            Maximum="24"
+                            Minimum="1"
                             ParsingNumberStyle="Integer"
-                            Value="{Binding ClipSkip, Converter={x:Static converters:NullableDefaultNumericConverters.IntToDecimal}}"
+                            Watermark="1"
-                            ClipValueToMinMax="True"/>
+                            Value="{Binding ClipSkip, Converter={x:Static converters:NullableDefaultNumericConverters.IntToDecimal}}" />
 
                     </sg:SpacedGrid>
                 </controls:Card>

StabilityMatrix.Avalonia/Languages/Resources.Designer.cs

@@ -941,6 +941,15 @@ namespace StabilityMatrix.Avalonia.Languages {
             }
         }
         
+        /// <summary>
+        ///   Looks up a localized string similar to Config.
+        /// </summary>
+        public static string Label_Config {
+            get {
+                return ResourceManager.GetString("Label_Config", resourceCulture);
+            }
+        }
+        
         /// <summary>
         ///   Looks up a localized string similar to Confirm Delete.
         /// </summary>

StabilityMatrix.Avalonia/Languages/Resources.resx

@@ -972,4 +972,7 @@
   <data name="Label_SelectDownloadLocation" xml:space="preserve">
     <value>Select Download Location:</value>
   </data>
+  <data name="Label_Config" xml:space="preserve">
+    <value>Config</value>
+  </data>
 </root>

StabilityMatrix.Avalonia/Models/Inference/ModuleApplyStepEventArgs.cs

@@ -1,7 +1,10 @@
 using System;
 using System.Collections.Generic;
-using StabilityMatrix.Core.Models.Api.Comfy.NodeTypes;
+using System.IO;
+using System.IO.Hashing;
+using System.Text;
 using StabilityMatrix.Core.Models.Api.Comfy.Nodes;
+using StabilityMatrix.Core.Models.Api.Comfy.NodeTypes;
 
 namespace StabilityMatrix.Avalonia.Models.Inference;
 
@@ -17,19 +20,40 @@ public class ModuleApplyStepEventArgs : EventArgs
     public ModuleApplyStepTemporaryArgs Temp { get; } = new();
 
     /// <summary>
-    /// Index of the step in the pipeline.
+    /// Generation overrides (like hires fix generate, current seed generate, etc.)
     /// </summary>
-    public int StepIndex { get; init; }
+    public IReadOnlyDictionary<Type, bool> IsEnabledOverrides { get; init; } = new Dictionary<Type, bool>();
 
-    /// <summary>
+    public List<(string SourcePath, string DestinationRelativePath)> FilesToTransfer { get; init; } = [];
-    /// Index
+
-    /// </summary>
+    public void AddFileTransfer(string sourcePath, string destinationRelativePath)
-    public int StepTypeIndex { get; init; }
+    {
+        FilesToTransfer.Add((sourcePath, destinationRelativePath));
+    }
 
     /// <summary>
-    /// Generation overrides (like hires fix generate, current seed generate, etc.)
+    /// Adds a file transfer to `models/configs`
     /// </summary>
-    public IReadOnlyDictionary<Type, bool> IsEnabledOverrides { get; init; } = new Dictionary<Type, bool>();
+    /// <returns>The destination relative path</returns>
+    public string AddFileTransferToConfigs(string sourcePath)
+    {
+        // To avoid conflicts, we'll add the file name's crc32 before the extension
+        var sourceNameWithoutExtension = Path.GetFileNameWithoutExtension(sourcePath);
+        var sourceExtension = Path.GetExtension(sourcePath);
+        var sourceNameCrc = Crc32.Hash(Encoding.UTF8.GetBytes(sourceNameWithoutExtension));
+        var sourceNameCrcShort = BitConverter
+            .ToString(sourceNameCrc)
+            .ToLowerInvariant()
+            .Replace("-", string.Empty)[..4];
+
+        var destFileName = $"{sourceNameWithoutExtension}_{sourceNameCrcShort}{sourceExtension}";
+
+        var destPath = Path.Combine("models", "configs", destFileName);
+
+        FilesToTransfer.Add((sourcePath, destPath));
+
+        return destPath;
+    }
 
     public class ModuleApplyStepTemporaryArgs
     {

StabilityMatrix.Avalonia/Services/InferenceClientManager.cs

@@ -369,6 +369,7 @@ public partial class InferenceClientManager : ObservableObject, IInferenceClient
         inferenceInputs.Create();
     }
 
+    [MemberNotNull(nameof(Client))]
     private async Task ConnectAsyncImpl(Uri uri, CancellationToken cancellationToken = default)
     {
         if (IsConnected)
@@ -480,11 +481,8 @@ public partial class InferenceClientManager : ObservableObject, IInferenceClient
 
         await ConnectAsyncImpl(uri, cancellationToken);
 
-        var packageDir = new DirectoryPath(packagePair.InstalledPackage.FullPath);
+        // Set package path as server path
-
+        Client.LocalServerPath = packagePair.InstalledPackage.FullPath!;
-        // Set package paths
-        Client!.OutputImagesDir = packageDir.JoinDir("output");
-        Client!.InputImagesDir = packageDir.JoinDir("input");
     }
 
     public async Task CloseAsync()

StabilityMatrix.Avalonia/StabilityMatrix.Avalonia.csproj

@@ -94,6 +94,7 @@
         <PackageReference Include="Sylvan.Common" Version="0.4.3" />
         <PackageReference Include="Sylvan.Data" Version="0.2.12" />
         <PackageReference Include="Sylvan.Data.Csv" Version="1.3.5" />
+        <PackageReference Include="System.IO.Hashing" Version="8.0.0" />
         <PackageReference Include="TextMateSharp.Grammars" Version="1.0.56" />
         <PackageReference Include="URISchemeTools" Version="1.0.2" />
     </ItemGroup>

StabilityMatrix.Avalonia/Styles/ControlThemes/BetterComboBoxStyles.axaml

@@ -10,7 +10,10 @@
 
     <Design.PreviewWith>
         <Panel Width="450" Height="600">
-            <StackPanel Margin="8" Spacing="4" Width="250">
+            <StackPanel
+                Width="250"
+                Margin="8"
+                Spacing="4">
                 <controls:BetterComboBox
                     HorizontalAlignment="Stretch"
                     ItemsSource="{x:Static mocks:DesignData.SampleHybridModels}"
@@ -83,7 +86,7 @@
             </Template>
         </Setter>
     </ControlTheme>
-    
+
     <!--  ReSharper disable once Xaml.StaticResourceNotResolved  -->
     <ControlTheme
         x:Key="{x:Type controls:BetterComboBox}"
@@ -139,11 +142,16 @@
                         <sg:SpacedGrid
                             Grid.Row="1"
                             Grid.Column="1"
+                            ColumnDefinitions="Auto,*"
                             RowDefinitions="Auto,Auto,Auto"
                             RowSpacing="1">
-                            <TextBlock Text="{Binding Local.DisplayModelName}" TextTrimming="CharacterEllipsis" />
+                            <TextBlock
+                                Grid.ColumnSpan="2"
+                                Text="{Binding Local.DisplayModelName}"
+                                TextTrimming="CharacterEllipsis" />
                             <TextBlock
                                 Grid.Row="1"
+                                Grid.Column="0"
                                 FontSize="12"
                                 FontWeight="Regular"
                                 Foreground="{DynamicResource TextFillColorSecondaryBrush}"
@@ -151,10 +159,27 @@
                                 TextTrimming="CharacterEllipsis" />
                             <TextBlock
                                 Grid.Row="2"
+                                Grid.Column="0"
                                 FontSize="11"
                                 FontWeight="Normal"
                                 Foreground="{DynamicResource TextFillColorTertiaryBrush}"
                                 Text="{Binding Local.DisplayModelFileName}" />
+
+                            <!--  Badges  -->
+                            <StackPanel
+                                Grid.Row="1"
+                                Grid.RowSpan="2"
+                                Grid.Column="1"
+                                HorizontalAlignment="Right"
+                                VerticalAlignment="Bottom"
+                                Orientation="Horizontal">
+                                <fluentIcons:SymbolIcon
+                                    Margin="4"
+                                    Foreground="{DynamicResource TextFillColorSecondaryBrush}"
+                                    IsVisible="{Binding Local.ConfigFullPath, Converter={x:Static StringConverters.IsNotNullOrEmpty}, FallbackValue=False}"
+                                    Symbol="BeakerSettings"
+                                    ToolTip.Tip="{Binding Local.DisplayConfigFileName}" />
+                            </StackPanel>
                         </sg:SpacedGrid>
                     </sg:SpacedGrid>
                 </DataTemplate>
@@ -236,7 +261,7 @@
         <Setter Property="SelectionBoxItemTemplate" Value="{StaticResource HybridModelSelectionBoxTemplateSelector}" />
         <Setter Property="ItemTemplate" Value="{StaticResource HybridModelTemplateSelector}" />
         <Setter Property="ItemContainerTheme" Value="{StaticResource BetterComboBoxItemHybridModelTheme}" />
-        
+
         <Style Selector="^ /template/ Popup#PART_Popup">
             <Setter Property="Width" Value="400" />
             <Setter Property="Placement" Value="Bottom" />

StabilityMatrix.Avalonia/ViewModels/Base/InferenceGenerationViewModelBase.cs

@@ -199,6 +199,26 @@ public abstract partial class InferenceGenerationViewModelBase
     /// </summary>
     protected virtual void BuildPrompt(BuildPromptEventArgs args) { }
 
+    /// <summary>
+    /// Uploads files required for the prompt
+    /// </summary>
+    protected virtual async Task UploadPromptFiles(
+        IEnumerable<(string SourcePath, string DestinationRelativePath)> files,
+        ComfyClient client
+    )
+    {
+        foreach (var (sourcePath, destinationRelativePath) in files)
+        {
+            Logger.Debug(
+                "Uploading prompt file {SourcePath} to relative path {DestinationPath}",
+                sourcePath,
+                destinationRelativePath
+            );
+
+            await client.UploadFileAsync(sourcePath, destinationRelativePath);
+        }
+    }
+
     /// <summary>
     /// Gets ImageSources that need to be uploaded as inputs
     /// </summary>
@@ -257,6 +277,9 @@ public abstract partial class InferenceGenerationViewModelBase
         // Upload input images
         await UploadInputImages(client);
 
+        // Upload required files
+        await UploadPromptFiles(args.FilesToTransfer, client);
+
         // Connect preview image handler
         client.PreviewImageReceived += OnPreviewImageReceived;
 
@@ -637,6 +660,7 @@ public abstract partial class InferenceGenerationViewModelBase
         public GenerationParameters? Parameters { get; init; }
         public InferenceProjectDocument? Project { get; init; }
         public bool ClearOutputImages { get; init; } = true;
+        public List<(string SourcePath, string DestinationRelativePath)> FilesToTransfer { get; init; } = [];
     }
 
     public class BuildPromptEventArgs : EventArgs
@@ -644,17 +668,28 @@ public abstract partial class InferenceGenerationViewModelBase
         public ComfyNodeBuilder Builder { get; } = new();
         public GenerateOverrides Overrides { get; init; } = new();
         public long? SeedOverride { get; init; }
+        public List<(string SourcePath, string DestinationRelativePath)> FilesToTransfer { get; init; } = [];
 
-        public static implicit operator ModuleApplyStepEventArgs(BuildPromptEventArgs args)
+        public ModuleApplyStepEventArgs ToModuleApplyStepEventArgs()
         {
             var overrides = new Dictionary<Type, bool>();
 
-            if (args.Overrides.IsHiresFixEnabled.HasValue)
+            if (Overrides.IsHiresFixEnabled.HasValue)
             {
-                overrides[typeof(HiresFixModule)] = args.Overrides.IsHiresFixEnabled.Value;
+                overrides[typeof(HiresFixModule)] = Overrides.IsHiresFixEnabled.Value;
             }
 
-            return new ModuleApplyStepEventArgs { Builder = args.Builder, IsEnabledOverrides = overrides };
+            return new ModuleApplyStepEventArgs
+            {
+                Builder = Builder,
+                IsEnabledOverrides = overrides,
+                FilesToTransfer = FilesToTransfer
+            };
+        }
+
+        public static implicit operator ModuleApplyStepEventArgs(BuildPromptEventArgs args)
+        {
+            return args.ToModuleApplyStepEventArgs();
         }
     }
 }

StabilityMatrix.Avalonia/ViewModels/Inference/InferenceImageToVideoViewModel.cs

@@ -199,6 +199,7 @@ public partial class InferenceImageToVideoViewModel
                 OutputNodeNames = buildPromptArgs.Builder.Connections.OutputNodeNames.ToArray(),
                 Parameters = SaveStateToParameters(new GenerationParameters()),
                 Project = InferenceProjectDocument.FromLoadable(this),
+                FilesToTransfer = buildPromptArgs.FilesToTransfer,
                 // Only clear output images on the first batch
                 ClearOutputImages = i == 0
             };

StabilityMatrix.Avalonia/ViewModels/Inference/InferenceTextToImageViewModel.cs

@@ -213,6 +213,7 @@ public class InferenceTextToImageViewModel : InferenceGenerationViewModelBase, I
                 OutputNodeNames = buildPromptArgs.Builder.Connections.OutputNodeNames.ToArray(),
                 Parameters = SaveStateToParameters(new GenerationParameters()),
                 Project = InferenceProjectDocument.FromLoadable(this),
+                FilesToTransfer = buildPromptArgs.FilesToTransfer,
                 // Only clear output images on the first batch
                 ClearOutputImages = i == 0
             };

StabilityMatrix.Avalonia/ViewModels/Inference/ModelCardViewModel.cs

@@ -1,9 +1,11 @@
 using System;
 using System.ComponentModel.DataAnnotations;
+using System.IO;
 using System.Linq;
 using System.Text.Json.Nodes;
 using System.Threading.Tasks;
 using CommunityToolkit.Mvvm.ComponentModel;
+using CommunityToolkit.Mvvm.Input;
 using StabilityMatrix.Avalonia.Controls;
 using StabilityMatrix.Avalonia.Languages;
 using StabilityMatrix.Avalonia.Models;
@@ -53,6 +55,28 @@ public partial class ModelCardViewModel(IInferenceClientManager clientManager)
 
     public IInferenceClientManager ClientManager { get; } = clientManager;
 
+    [RelayCommand]
+    private static async Task OnConfigClickAsync()
+    {
+        await DialogHelper
+            .CreateMarkdownDialog(
+                """
+                                                 You can use a config (.yaml) file to load a model with specific settings.
+                                                 
+                                                 Place the config file next to the model file with the same name:
+                                                 ```md
+                                                 Models/
+                                                     StableDiffusion/
+                                                         my_model.safetensors
+                                                         my_model.yaml <-
+                                                 ```
+                                                 """,
+                "Using Model Configs",
+                TextEditorPreset.Console
+            )
+            .ShowAsync();
+    }
+
     public async Task<bool> ValidateModel()
     {
         if (SelectedModel != null)
@@ -66,16 +90,41 @@ public partial class ModelCardViewModel(IInferenceClientManager clientManager)
         return false;
     }
 
+    private static ComfyTypedNodeBase<
+        ModelNodeConnection,
+        ClipNodeConnection,
+        VAENodeConnection
+    > GetModelLoader(ModuleApplyStepEventArgs e, string nodeName, HybridModelFile model)
+    {
+        // Check if config
+        if (model.Local?.ConfigFullPath is { } configPath)
+        {
+            // We'll need to upload the config file to `models/configs` later
+            var uploadConfigPath = e.AddFileTransferToConfigs(configPath);
+
+            return new ComfyNodeBuilder.CheckpointLoader
+            {
+                Name = nodeName,
+                // Only the file name is needed
+                ConfigName = Path.GetFileName(uploadConfigPath),
+                CkptName = model.RelativePath
+            };
+        }
+
+        // Simple loader if no config
+        return new ComfyNodeBuilder.CheckpointLoaderSimple { Name = nodeName, CkptName = model.RelativePath };
+    }
+
     /// <inheritdoc />
     public virtual void ApplyStep(ModuleApplyStepEventArgs e)
     {
         // Load base checkpoint
         var baseLoader = e.Nodes.AddTypedNode(
-            new ComfyNodeBuilder.CheckpointLoaderSimple
+            GetModelLoader(
-            {
+                e,
-                Name = "CheckpointLoader_Base",
+                "CheckpointLoader_Base",
-                CkptName = SelectedModel?.RelativePath ?? throw new ValidationException("Model not selected")
+                SelectedModel ?? throw new ValidationException("Model not selected")
-            }
+            )
         );
 
         e.Builder.Connections.Base.Model = baseLoader.Output1;
@@ -86,13 +135,11 @@ public partial class ModelCardViewModel(IInferenceClientManager clientManager)
         if (IsRefinerSelectionEnabled && SelectedRefiner is { IsNone: false })
         {
             var refinerLoader = e.Nodes.AddTypedNode(
-                new ComfyNodeBuilder.CheckpointLoaderSimple
+                GetModelLoader(
-                {
+                    e,
-                    Name = "CheckpointLoader_Refiner",
+                    "CheckpointLoader_Refiner",
-                    CkptName =
+                    SelectedRefiner ?? throw new ValidationException("Refiner Model enabled but not selected")
-                        SelectedRefiner?.RelativePath
+                )
-                        ?? throw new ValidationException("Refiner Model enabled but not selected")
-                }
             );
 
             e.Builder.Connections.Refiner.Model = refinerLoader.Output1;

StabilityMatrix.Core/Inference/ComfyClient.cs

@@ -35,14 +35,19 @@ public class ComfyClient : InferenceClientBase
     public Uri BaseAddress { get; }
 
     /// <summary>
-    /// Optional local path to output images.
+    /// If available, the local path to the server root directory.
     /// </summary>
-    public DirectoryPath? OutputImagesDir { get; set; }
+    public DirectoryPath? LocalServerPath { get; set; }
 
     /// <summary>
-    /// Optional local path to input images.
+    /// Path to the "output" folder from LocalServerPath
     /// </summary>
-    public DirectoryPath? InputImagesDir { get; set; }
+    public DirectoryPath? OutputImagesDir => LocalServerPath?.JoinDir("output");
+
+    /// <summary>
+    /// Path to the "input" folder from LocalServerPath
+    /// </summary>
+    public DirectoryPath? InputImagesDir => LocalServerPath?.JoinDir("input");
 
     /// <summary>
     /// Dictionary of ongoing prompt execution tasks
@@ -144,9 +149,7 @@ public class ComfyClient : InferenceClientBase
 
         if (json.Type == ComfyWebSocketResponseType.Executing)
         {
-            var executingData = json.GetDataAsType<ComfyWebSocketExecutingData>(
+            var executingData = json.GetDataAsType<ComfyWebSocketExecutingData>(jsonSerializerOptions);
-                jsonSerializerOptions
-            );
             if (executingData?.PromptId is null)
             {
                 Logger.Warn($"Could not parse executing data {json.Data}, skipping");
@@ -165,9 +168,7 @@ public class ComfyClient : InferenceClientBase
                 }
                 else
                 {
-                    Logger.Warn(
+                    Logger.Warn($"Could not find task for prompt {executingData.PromptId}, skipping");
-                        $"Could not find task for prompt {executingData.PromptId}, skipping"
-                    );
                 }
             }
             // Otherwise set the task's active node to the one received
@@ -194,9 +195,7 @@ public class ComfyClient : InferenceClientBase
         }
         else if (json.Type == ComfyWebSocketResponseType.Progress)
         {
-            var progressData = json.GetDataAsType<ComfyWebSocketProgressData>(
+            var progressData = json.GetDataAsType<ComfyWebSocketProgressData>(jsonSerializerOptions);
-                jsonSerializerOptions
-            );
             if (progressData is null)
             {
                 Logger.Warn($"Could not parse progress data {json.Data}, skipping");
@@ -224,11 +223,7 @@ public class ComfyClient : InferenceClientBase
             {
                 task.RunningNode = null;
                 task.SetException(
-                    new ComfyNodeException
+                    new ComfyNodeException { ErrorData = errorData, JsonData = json.Data.ToString() }
-                    {
-                        ErrorData = errorData,
-                        JsonData = json.Data.ToString()
-                    }
                 );
                 currentPromptTask = null;
             }
@@ -306,9 +301,7 @@ public class ComfyClient : InferenceClientBase
 
     public override async Task CloseAsync(CancellationToken cancellationToken = default)
     {
-        await webSocketClient
+        await webSocketClient.Stop(WebSocketCloseStatus.NormalClosure, string.Empty).ConfigureAwait(false);
-            .Stop(WebSocketCloseStatus.NormalClosure, string.Empty)
-            .ConfigureAwait(false);
     }
 
     public async Task<ComfyTask> QueuePromptAsync(
@@ -349,13 +342,39 @@ public class ComfyClient : InferenceClientBase
     )
     {
         var streamPart = new StreamPart(image, fileName);
-        return comfyApi.PostUploadImage(
+        return comfyApi.PostUploadImage(streamPart, "true", "input", "Inference", cancellationToken);
-            streamPart,
+    }
-            "true",
+
-            "input",
+    /// <summary>
-            "Inference",
+    /// Upload a file to the server at the given relative path from server's root
-            cancellationToken
+    /// </summary>
-        );
+    public async Task UploadFileAsync(
+        string sourcePath,
+        string destinationRelativePath,
+        CancellationToken cancellationToken = default
+    )
+    {
+        cancellationToken.ThrowIfCancellationRequested();
+
+        // Currently there is no api, so we do a local file copy
+        if (LocalServerPath is null)
+        {
+            throw new InvalidOperationException("LocalServerPath is not set");
+        }
+
+        var sourceFile = new FilePath(sourcePath);
+        var destFile = LocalServerPath.JoinFile(destinationRelativePath);
+
+        Logger.Info("Copying file from {Source} to {Dest}", sourcePath, destFile);
+
+        if (!sourceFile.Exists)
+        {
+            throw new FileNotFoundException("Source file does not exist", sourcePath);
+        }
+
+        destFile.Directory?.Create();
+
+        await sourceFile.CopyToAsync(destFile, true).ConfigureAwait(false);
     }
 
     public async Task<Dictionary<string, List<ComfyImage>?>> GetImagesForExecutedPromptAsync(
@@ -413,9 +432,7 @@ public class ComfyClient : InferenceClientBase
         CancellationToken cancellationToken = default
     )
     {
-        var response = await comfyApi
+        var response = await comfyApi.GetObjectInfo(nodeName, cancellationToken).ConfigureAwait(false);
-            .GetObjectInfo(nodeName, cancellationToken)
-            .ConfigureAwait(false);
 
         var info = response[nodeName];
         return info.Input.GetRequiredValueAsNestedList(optionName);

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

@@ -226,6 +226,13 @@ public class ComfyNodeBuilder
         };
     }
 
+    public record CheckpointLoader
+        : ComfyTypedNodeBase<ModelNodeConnection, ClipNodeConnection, VAENodeConnection>
+    {
+        public required string ConfigName { get; init; }
+        public required string CkptName { get; init; }
+    }
+
     public record CheckpointLoaderSimple
         : ComfyTypedNodeBase<ModelNodeConnection, ClipNodeConnection, VAENodeConnection>
     {

StabilityMatrix.Core/Models/Database/LocalModelFile.cs

@@ -34,6 +34,11 @@ public record LocalModelFile
     /// </summary>
     public string? PreviewImageFullPath { get; set; }
 
+    /// <summary>
+    /// Optional full path to the model's configuration (.yaml) file.
+    /// </summary>
+    public string? ConfigFullPath { get; set; }
+
     /// <summary>
     /// Whether or not an update is available for this model
     /// </summary>
@@ -88,6 +93,9 @@ public record LocalModelFile
     [BsonIgnore]
     public string DisplayModelFileName => FileName;
 
+    [BsonIgnore]
+    public string DisplayConfigFileName => Path.GetFileName(ConfigFullPath) ?? string.Empty;
+
     public string GetFullPath(string rootModelDirectory)
     {
         return Path.Combine(rootModelDirectory, RelativePath);
@@ -103,29 +111,6 @@ public record LocalModelFile
             : Path.Combine(rootModelDirectory, PreviewImageRelativePath);
     }
 
-    /*protected bool Equals(LocalModelFile other)
-    {
-        return RelativePath == other.RelativePath;
-    }
-
-    /// <inheritdoc />
-    public override bool Equals(object? obj)
-    {
-        if (ReferenceEquals(null, obj))
-            return false;
-        if (ReferenceEquals(this, obj))
-            return true;
-        if (obj.GetType() != this.GetType())
-            return false;
-        return Equals((LocalModelFile)obj);
-    }
-
-    /// <inheritdoc />
-    public override int GetHashCode()
-    {
-        return RelativePath.GetHashCode();
-    }*/
-
     public static readonly HashSet<string> SupportedCheckpointExtensions =
     [
         ".safetensors",

StabilityMatrix.Core/Models/FileInterfaces/DirectoryPath.cs

@@ -1,4 +1,5 @@
 using System.Collections;
+using System.ComponentModel;
 using System.Text.Json.Serialization;
 using JetBrains.Annotations;
 using StabilityMatrix.Core.Converters.Json;
@@ -6,6 +7,7 @@ using StabilityMatrix.Core.Converters.Json;
 namespace StabilityMatrix.Core.Models.FileInterfaces;
 
 [PublicAPI]
+[Localizable(false)]
 [JsonConverter(typeof(StringJsonConverter<DirectoryPath>))]
 public class DirectoryPath : FileSystemPath, IPathObject, IEnumerable<FileSystemPath>
 {
@@ -43,7 +45,7 @@ public class DirectoryPath : FileSystemPath, IPathObject, IEnumerable<FileSystem
     [JsonIgnore]
     public DirectoryPath? Parent => Info.Parent == null ? null : new DirectoryPath(Info.Parent);
 
-    public DirectoryPath(string path)
+    public DirectoryPath([Localizable(false)] string path)
         : base(path) { }
 
     public DirectoryPath(FileSystemPath path)
@@ -56,7 +58,7 @@ public class DirectoryPath : FileSystemPath, IPathObject, IEnumerable<FileSystem
         this.info = info;
     }
 
-    public DirectoryPath(params string[] paths)
+    public DirectoryPath([Localizable(false)] params string[] paths)
         : base(paths) { }
 
     public DirectoryPath RelativeTo(DirectoryPath path)
@@ -226,13 +228,13 @@ public class DirectoryPath : FileSystemPath, IPathObject, IEnumerable<FileSystem
     /// <summary>
     /// Join with other paths to form a new directory path.
     /// </summary>
-    public DirectoryPath JoinDir(params DirectoryPath[] paths) =>
+    public DirectoryPath JoinDir([Localizable(false)] params DirectoryPath[] paths) =>
         new(Path.Combine(FullPath, Path.Combine(paths.Select(path => path.FullPath).ToArray())));
 
     /// <summary>
     /// Join with other paths to form a new file path.
     /// </summary>
-    public FilePath JoinFile(params FilePath[] paths) =>
+    public FilePath JoinFile([Localizable(false)] params FilePath[] paths) =>
         new(Path.Combine(FullPath, Path.Combine(paths.Select(path => path.FullPath).ToArray())));
 
     /// <summary>
@@ -305,12 +307,13 @@ public class DirectoryPath : FileSystemPath, IPathObject, IEnumerable<FileSystem
         new(Path.Combine(path, other.FullName));
 
     // DirectoryPath + string = string
-    public static string operator +(DirectoryPath path, string other) => Path.Combine(path, other);
+    public static string operator +(DirectoryPath path, [Localizable(false)] string other) =>
+        Path.Combine(path, other);
 
     // Implicit conversions to and from string
     public static implicit operator string(DirectoryPath path) => path.FullPath;
 
-    public static implicit operator DirectoryPath(string path) => new(path);
+    public static implicit operator DirectoryPath([Localizable(false)] string path) => new(path);
 
     // Implicit conversions to and from DirectoryInfo
     public static implicit operator DirectoryInfo(DirectoryPath path) => path.Info;

StabilityMatrix.Core/Models/FileInterfaces/FilePath.cs

@@ -1,11 +1,13 @@
-using System.Diagnostics.CodeAnalysis;
+using System.ComponentModel;
 using System.Text;
 using System.Text.Json.Serialization;
+using JetBrains.Annotations;
 using StabilityMatrix.Core.Converters.Json;
 
 namespace StabilityMatrix.Core.Models.FileInterfaces;
 
-[SuppressMessage("ReSharper", "MemberCanBePrivate.Global")]
+[PublicAPI]
+[Localizable(false)]
 [JsonConverter(typeof(StringJsonConverter<FilePath>))]
 public partial class FilePath : FileSystemPath, IPathObject
 {
@@ -59,7 +61,7 @@ public partial class FilePath : FileSystemPath, IPathObject
         }
     }
 
-    public FilePath(string path)
+    public FilePath([Localizable(false)] string path)
         : base(path) { }
 
     public FilePath(FileInfo fileInfo)
@@ -71,7 +73,7 @@ public partial class FilePath : FileSystemPath, IPathObject
     public FilePath(FileSystemPath path)
         : base(path) { }
 
-    public FilePath(params string[] paths)
+    public FilePath([Localizable(false)] params string[] paths)
         : base(paths) { }
 
     public FilePath RelativeTo(DirectoryPath path)
@@ -151,7 +153,7 @@ public partial class FilePath : FileSystemPath, IPathObject
     /// <summary>
     /// Rename the file.
     /// </summary>
-    public FilePath Rename(string fileName)
+    public FilePath Rename([Localizable(false)] string fileName)
     {
         if (Path.GetDirectoryName(FullPath) is { } directory && !string.IsNullOrWhiteSpace(directory))
         {
@@ -245,5 +247,5 @@ public partial class FilePath : FileSystemPath, IPathObject
     // Implicit conversions to and from string
     public static implicit operator string(FilePath path) => path.FullPath;
 
-    public static implicit operator FilePath(string path) => new(path);
+    public static implicit operator FilePath([Localizable(false)] string path) => new(path);
 }

StabilityMatrix.Core/Models/FileInterfaces/FileSystemPath.cs

@@ -1,9 +1,11 @@
-using System.Diagnostics.CodeAnalysis;
+using System.ComponentModel;
+using System.Diagnostics.CodeAnalysis;
 using JetBrains.Annotations;
 
 namespace StabilityMatrix.Core.Models.FileInterfaces;
 
 [PublicAPI]
+[Localizable(false)]
 public class FileSystemPath : IEquatable<FileSystemPath>, IFormattable
 {
     public string FullPath { get; }

StabilityMatrix.Core/Models/HybridModelFile.cs

@@ -1,7 +1,6 @@
 using System.Diagnostics.CodeAnalysis;
 using System.Text.Json.Serialization;
 using StabilityMatrix.Core.Models.Database;
-using StabilityMatrix.Core.Models.FileInterfaces;
 
 namespace StabilityMatrix.Core.Models;
 
@@ -125,9 +124,14 @@ public record HybridModelFile
             if (x.GetType() != y.GetType())
                 return false;
 
-            return Equals(x.RelativePath, y.RelativePath)
+            if (!Equals(x.RelativePath, y.RelativePath))
-                && x.IsNone == y.IsNone
+                return false;
-                && x.IsDefault == y.IsDefault;
+
+            // This equality affects replacements of remote over local models
+            // We want local and remote models to be considered equal if they have the same relative path
+            // But 2 local models with the same path but different config paths should be considered different
+
+            return !(x.Type == y.Type && x.Local?.ConfigFullPath != y.Local?.ConfigFullPath);
         }
 
         public int GetHashCode(HybridModelFile obj)

StabilityMatrix.Core/Services/ModelIndexService.cs

@@ -158,6 +158,12 @@ public partial class ModelIndexService : IModelIndexService
                 localModel.PreviewImageRelativePath = Path.GetRelativePath(modelsDir, previewImagePath);
             }
 
+            // Try to find a config file (same name as model file but with .yaml extension)
+            if (file.WithName($"{file.NameWithoutExtension}.yaml") is { Exists: true } configFile)
+            {
+                localModel.ConfigFullPath = configFile;
+            }
+
             // Insert into database
             await localModelFiles.InsertAsync(localModel).ConfigureAwait(false);