Browse Source

Merge pull request #519 from ionite34/model-configs

Inference - Model config support
pull/495/head
Ionite 9 months ago committed by GitHub
parent
commit
188be7bb45
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
  1. 5
      CHANGELOG.md
  2. 75
      StabilityMatrix.Avalonia/Controls/Inference/ModelCard.axaml
  3. 9
      StabilityMatrix.Avalonia/Languages/Resources.Designer.cs
  4. 3
      StabilityMatrix.Avalonia/Languages/Resources.resx
  5. 42
      StabilityMatrix.Avalonia/Models/Inference/ModuleApplyStepEventArgs.cs
  6. 8
      StabilityMatrix.Avalonia/Services/InferenceClientManager.cs
  7. 1
      StabilityMatrix.Avalonia/StabilityMatrix.Avalonia.csproj
  8. 29
      StabilityMatrix.Avalonia/Styles/ControlThemes/BetterComboBoxStyles.axaml
  9. 43
      StabilityMatrix.Avalonia/ViewModels/Base/InferenceGenerationViewModelBase.cs
  10. 1
      StabilityMatrix.Avalonia/ViewModels/Inference/InferenceImageToVideoViewModel.cs
  11. 1
      StabilityMatrix.Avalonia/ViewModels/Inference/InferenceTextToImageViewModel.cs
  12. 71
      StabilityMatrix.Avalonia/ViewModels/Inference/ModelCardViewModel.cs
  13. 79
      StabilityMatrix.Core/Inference/ComfyClient.cs
  14. 7
      StabilityMatrix.Core/Models/Api/Comfy/Nodes/ComfyNodeBuilder.cs
  15. 31
      StabilityMatrix.Core/Models/Database/LocalModelFile.cs
  16. 15
      StabilityMatrix.Core/Models/FileInterfaces/DirectoryPath.cs
  17. 14
      StabilityMatrix.Core/Models/FileInterfaces/FilePath.cs
  18. 4
      StabilityMatrix.Core/Models/FileInterfaces/FileSystemPath.cs
  19. 12
      StabilityMatrix.Core/Models/HybridModelFile.cs
  20. 6
      StabilityMatrix.Core/Services/ModelIndexService.cs

5
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

75
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,16 +21,31 @@
</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
ColumnSpacing="8"
RowSpacing="0"
ColumnDefinitions="Auto,*,Auto"
RowDefinitions="*,*,*,*">
ColumnSpacing="8"
RowDefinitions="*,*,*,*"
RowSpacing="0">
<!-- Model -->
<TextBlock
Grid.Column="0"
@ -40,23 +57,37 @@
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 IsChecked="{Binding IsVaeSelectionEnabled}" Text="{x:Static lang:Resources.Label_VAE}" />
<ui:ToggleMenuFlyoutItem IsChecked="{Binding IsClipSkipEnabled}" Text="{x:Static lang:Resources.Label_CLIPSkip}" />
<ui:ToggleMenuFlyoutItem
IconSource="{StaticResource IconCube}"
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,11 +108,11 @@
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
@ -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"
Maximum="24"
HorizontalAlignment="Stretch"
ClipValueToMinMax="True"
Increment="1"
IsVisible="{Binding IsClipSkipEnabled}"
Maximum="24"
Minimum="1"
ParsingNumberStyle="Integer"
Value="{Binding ClipSkip, Converter={x:Static converters:NullableDefaultNumericConverters.IntToDecimal}}"
ClipValueToMinMax="True"/>
Watermark="1"
Value="{Binding ClipSkip, Converter={x:Static converters:NullableDefaultNumericConverters.IntToDecimal}}" />
</sg:SpacedGrid>
</controls:Card>

9
StabilityMatrix.Avalonia/Languages/Resources.Designer.cs generated

@ -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>

3
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>

42
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>
/// Index
/// </summary>
public int StepTypeIndex { get; init; }
public List<(string SourcePath, string DestinationRelativePath)> FilesToTransfer { get; init; } = [];
public void AddFileTransfer(string sourcePath, string destinationRelativePath)
{
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
{

8
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 paths
Client!.OutputImagesDir = packageDir.JoinDir("output");
Client!.InputImagesDir = packageDir.JoinDir("input");
// Set package path as server path
Client.LocalServerPath = packagePair.InstalledPackage.FullPath!;
}
public async Task CloseAsync()

1
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>

29
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}"
@ -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>

43
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();
}
}
}

1
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
};

1
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
};

71
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
{
Name = "CheckpointLoader_Base",
CkptName = SelectedModel?.RelativePath ?? throw new ValidationException("Model not selected")
}
GetModelLoader(
e,
"CheckpointLoader_Base",
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
{
Name = "CheckpointLoader_Refiner",
CkptName =
SelectedRefiner?.RelativePath
?? throw new ValidationException("Refiner Model enabled but not selected")
}
GetModelLoader(
e,
"CheckpointLoader_Refiner",
SelectedRefiner ?? throw new ValidationException("Refiner Model enabled but not selected")
)
);
e.Builder.Connections.Refiner.Model = refinerLoader.Output1;

79
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>(
jsonSerializerOptions
);
var executingData = json.GetDataAsType<ComfyWebSocketExecutingData>(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(
$"Could not find task for prompt {executingData.PromptId}, skipping"
);
Logger.Warn($"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>(
jsonSerializerOptions
);
var progressData = json.GetDataAsType<ComfyWebSocketProgressData>(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
{
ErrorData = errorData,
JsonData = json.Data.ToString()
}
new ComfyNodeException { 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
.Stop(WebSocketCloseStatus.NormalClosure, string.Empty)
.ConfigureAwait(false);
await webSocketClient.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(
streamPart,
"true",
"input",
"Inference",
cancellationToken
);
return comfyApi.PostUploadImage(streamPart, "true", "input", "Inference", cancellationToken);
}
/// <summary>
/// Upload a file to the server at the given relative path from server's root
/// </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
.GetObjectInfo(nodeName, cancellationToken)
.ConfigureAwait(false);
var response = await comfyApi.GetObjectInfo(nodeName, cancellationToken).ConfigureAwait(false);
var info = response[nodeName];
return info.Input.GetRequiredValueAsNestedList(optionName);

7
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>
{

31
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",

15
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;

14
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);
}

4
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; }

12
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)
&& x.IsNone == y.IsNone
&& x.IsDefault == y.IsDefault;
if (!Equals(x.RelativePath, y.RelativePath))
return false;
// 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)

6
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);

Loading…
Cancel
Save