Skip to content

Commit e63443a

Browse files
mohnjilesmohnjiles
authored
Merge pull request LykosAI#1074 from Genteure/feat/model-metadata
Add safetensor metadata parsing and display
2 parents f7c5f1e + 8dad8db commit e63443a

File tree

11 files changed

+593
-2
lines changed

11 files changed

+593
-2
lines changed

‎StabilityMatrix.Avalonia/DesignData/DesignData.cs‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1167,6 +1167,24 @@ public static CompletionList SampleCompletionList
11671167
}
11681168
};
11691169

1170+
public static SafetensorMetadataViewModel SafetensorMetadataViewModel =>
1171+
DialogFactory.Get<SafetensorMetadataViewModel>(vm =>
1172+
{
1173+
vm.Metadata = new SafetensorMetadata
1174+
{
1175+
TagFrequency = Enumerable
1176+
.Range(1, 100)
1177+
.Select(i => new SafetensorMetadata.Tag("tag" + i, i))
1178+
.ToList(),
1179+
OtherMetadata = new List<SafetensorMetadata.Metadata>
1180+
{
1181+
new("Name1", "Value1"),
1182+
new("Name2", "Value2"),
1183+
new("Name3", "Value3"),
1184+
}
1185+
};
1186+
});
1187+
11701188
public static ModelMetadataEditorDialogViewModel MetadataEditorDialogViewModel =>
11711189
DialogFactory.Get<ModelMetadataEditorDialogViewModel>(vm =>
11721190
{

‎StabilityMatrix.Avalonia/ViewModels/CheckpointManager/CheckpointFileViewModel.cs‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -307,6 +307,55 @@ private async Task RenameAsync()
307307
}
308308
}
309309

310+
[RelayCommand]
311+
private async Task OpenSafetensorMetadataViewer()
312+
{
313+
if (!CheckpointFile.SafetensorMetadataParsed)
314+
{
315+
if (
316+
!settingsManager.IsLibraryDirSet
317+
|| new DirectoryPath(settingsManager.ModelsDirectory) is not { Exists: true } modelsDir
318+
)
319+
{
320+
return;
321+
}
322+
323+
try
324+
{
325+
var safetensorPath = CheckpointFile.GetFullPath(modelsDir);
326+
327+
var metadata = await SafetensorMetadata.ParseAsync(safetensorPath);
328+
329+
CheckpointFile.SafetensorMetadataParsed = true;
330+
CheckpointFile.SafetensorMetadata = metadata;
331+
}
332+
catch (Exception ex)
333+
{
334+
logger.LogWarning(ex, "Failed to parse safetensor metadata");
335+
return;
336+
}
337+
}
338+
339+
if (!CheckpointFile.SafetensorMetadataParsed)
340+
{
341+
return;
342+
}
343+
344+
var vm = vmFactory.Get<SafetensorMetadataViewModel>(vm =>
345+
{
346+
vm.ModelName = CheckpointFile.DisplayModelName;
347+
vm.Metadata = CheckpointFile.SafetensorMetadata;
348+
});
349+
350+
var dialog = vm.GetDialog();
351+
dialog.MinDialogHeight = 800;
352+
dialog.MinDialogWidth = 700;
353+
dialog.CloseButtonText = "Close";
354+
dialog.DefaultButton = ContentDialogButton.Close;
355+
356+
await dialog.ShowAsync();
357+
}
358+
310359
[RelayCommand]
311360
private async Task OpenMetadataEditor()
312361
{

‎StabilityMatrix.Avalonia/ViewModels/Dialogs/SafetensorMetadataViewModel.cs‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
using CommunityToolkit.Mvvm.ComponentModel;
2+
using CommunityToolkit.Mvvm.Input;
3+
using Injectio.Attributes;
4+
using StabilityMatrix.Avalonia.ViewModels.Base;
5+
using StabilityMatrix.Avalonia.Views.Dialogs;
6+
using StabilityMatrix.Core.Attributes;
7+
using StabilityMatrix.Core.Models;
8+
9+
namespace StabilityMatrix.Avalonia.ViewModels.Dialogs;
10+
11+
[View(typeof(SafetensorMetadataDialog))]
12+
[ManagedService]
13+
[RegisterSingleton<SafetensorMetadataViewModel>]
14+
public partial class SafetensorMetadataViewModel : ContentDialogViewModelBase
15+
{
16+
[ObservableProperty]
17+
private string? modelName;
18+
19+
[ObservableProperty]
20+
private SafetensorMetadata? metadata;
21+
22+
[RelayCommand]
23+
public void CopyTagToClipboard(string tag)
24+
{
25+
App.Clipboard?.SetTextAsync(tag);
26+
}
27+
}

‎StabilityMatrix.Avalonia/Views/CheckpointsPage.axaml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -489,6 +489,11 @@
489489
</ui:MenuFlyoutItem.IsVisible>
490490
</ui:MenuFlyoutItem>
491491

492+
<ui:MenuFlyoutItem
493+
Command="{Binding OpenSafetensorMetadataViewerCommand}"
494+
IconSource="Tag"
495+
IsVisible="{Binding CheckpointFile.IsSafetensorFile}"
496+
Text="View Safetensor Metadata" />
492497
<ui:MenuFlyoutItem
493498
Command="{Binding OpenMetadataEditorCommand}"
494499
IconSource="Edit"

‎StabilityMatrix.Avalonia/Views/Dialogs/SafetensorMetadataDialog.axaml‎

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,113 @@
1+
<controls:UserControlBase
2+
x:Class="StabilityMatrix.Avalonia.Views.Dialogs.SafetensorMetadataDialog"
3+
xmlns="https://github.com/avaloniaui"
4+
xmlns:x="http://schemas.microsoft.com/winfx/2006/xaml"
5+
xmlns:controls="clr-namespace:StabilityMatrix.Avalonia.Controls"
6+
xmlns:d="http://schemas.microsoft.com/expression/blend/2008"
7+
xmlns:dialogs="clr-namespace:StabilityMatrix.Avalonia.ViewModels.Dialogs"
8+
xmlns:lang="clr-namespace:StabilityMatrix.Avalonia.Languages"
9+
xmlns:mc="http://schemas.openxmlformats.org/markup-compatibility/2006"
10+
xmlns:mocks="clr-namespace:StabilityMatrix.Avalonia.DesignData"
11+
xmlns:models="clr-namespace:StabilityMatrix.Core.Models;assembly=StabilityMatrix.Core"
12+
xmlns:ui="using:FluentAvalonia.UI.Controls"
13+
d:DataContext="{x:Static mocks:DesignData.SafetensorMetadataViewModel}"
14+
d:DesignHeight="550"
15+
d:DesignWidth="700"
16+
x:DataType="dialogs:SafetensorMetadataViewModel"
17+
mc:Ignorable="d">
18+
<Grid RowDefinitions="Auto,Auto,Auto">
19+
<TextBlock
20+
Margin="8"
21+
HorizontalAlignment="Center"
22+
FontSize="24"
23+
FontWeight="SemiBold"
24+
Text="Safetensor Metadata" />
25+
<TextBlock
26+
Grid.Row="1"
27+
Margin="8"
28+
HorizontalAlignment="Center"
29+
FontSize="16"
30+
FontWeight="SemiBold"
31+
Text="{Binding ModelName}" />
32+
33+
<TextBlock
34+
Grid.Row="2"
35+
Margin="10,20"
36+
HorizontalAlignment="Center"
37+
FontSize="16"
38+
FontStyle="Italic"
39+
IsVisible="{Binding Metadata, Converter={x:Static ObjectConverters.IsNull}}"
40+
Text="No Metadata" />
41+
42+
<Grid
43+
Grid.Row="2"
44+
IsVisible="{Binding Metadata, Converter={x:Static ObjectConverters.IsNotNull}}"
45+
RowDefinitions="Auto,Auto,Auto,Auto">
46+
47+
<!-- List of tags -->
48+
<TextBlock
49+
Grid.Row="0"
50+
Margin="8"
51+
HorizontalAlignment="Left"
52+
FontSize="16"
53+
FontWeight="SemiBold"
54+
IsVisible="{Binding Metadata.TagFrequency, Converter={x:Static ObjectConverters.IsNotNull}}"
55+
Text="Trained Tags" />
56+
57+
<ui:ItemsRepeater
58+
Grid.Row="1"
59+
Margin="8"
60+
IsVisible="{Binding Metadata.TagFrequency, Converter={x:Static ObjectConverters.IsNotNull}}"
61+
ItemsSource="{Binding Metadata.TagFrequency}">
62+
<ui:ItemsRepeater.Layout>
63+
<ui:FlowLayout
64+
MinColumnSpacing="4"
65+
MinRowSpacing="4"
66+
Orientation="Horizontal" />
67+
</ui:ItemsRepeater.Layout>
68+
<ui:ItemsRepeater.ItemTemplate>
69+
<DataTemplate>
70+
<Button
71+
Command="{Binding $parent[ui:ItemsRepeater].((dialogs:SafetensorMetadataViewModel)DataContext).CopyTagToClipboardCommand}"
72+
CommandParameter="{Binding Name}"
73+
Cursor="Hand">
74+
<StackPanel Orientation="Horizontal">
75+
<TextBlock FontWeight="SemiBold" Text="{Binding Name}" />
76+
<TextBlock Margin="5,0,0,0" Text="{Binding Frequency}" />
77+
</StackPanel>
78+
</Button>
79+
</DataTemplate>
80+
</ui:ItemsRepeater.ItemTemplate>
81+
</ui:ItemsRepeater>
82+
83+
<!-- All other metadata -->
84+
<TextBlock
85+
Grid.Row="2"
86+
Margin="8"
87+
FontSize="16"
88+
FontWeight="SemiBold"
89+
Text="Other Metadata" />
90+
<TextBlock
91+
Grid.Row="3"
92+
Margin="8"
93+
FontSize="16"
94+
FontStyle="Italic"
95+
IsVisible="{Binding !Metadata.OtherMetadata.Count}"
96+
Text="No Other Metadata" />
97+
<ui:ItemsRepeater
98+
Grid.Row="3"
99+
Margin="8"
100+
IsVisible="{Binding !!Metadata.OtherMetadata.Count}"
101+
ItemsSource="{Binding Metadata.OtherMetadata}">
102+
<ui:ItemsRepeater.ItemTemplate>
103+
<DataTemplate>
104+
<StackPanel Margin="5" Orientation="Vertical">
105+
<TextBlock FontWeight="SemiBold" Text="{Binding Name}" />
106+
<TextBlock Text="{Binding Value}" />
107+
</StackPanel>
108+
</DataTemplate>
109+
</ui:ItemsRepeater.ItemTemplate>
110+
</ui:ItemsRepeater>
111+
</Grid>
112+
</Grid>
113+
</controls:UserControlBase>

‎StabilityMatrix.Avalonia/Views/Dialogs/SafetensorMetadataDialog.axaml.cs‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
using Avalonia.Controls;
2+
using Avalonia.Markup.Xaml;
3+
using Injectio.Attributes;
4+
5+
namespace StabilityMatrix.Avalonia.Views.Dialogs;
6+
7+
[RegisterTransient<SafetensorMetadataDialog>]
8+
public partial class SafetensorMetadataDialog : UserControl
9+
{
10+
public SafetensorMetadataDialog()
11+
{
12+
InitializeComponent();
13+
}
14+
15+
private void InitializeComponent()
16+
{
17+
AvaloniaXamlLoader.Load(this);
18+
}
19+
}

‎StabilityMatrix.Core/Helper/Compat.cs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -189,7 +189,11 @@ public static string GetExecutableName()
189189
var appImage = Environment.GetEnvironmentVariable("APPIMAGE");
190190
if (string.IsNullOrEmpty(appImage))
191191
{
192+
#if DEBUG
193+
return "DEBUG_NOT_RUNNING_IN_APPIMAGE";
194+
#else
192195
throw new Exception("Could not find APPIMAGE environment variable");
196+
#endif
193197
}
194198
return Path.GetFileName(appImage);
195199
}

‎StabilityMatrix.Core/Models/Database/LocalModelFile.cs‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -121,7 +121,10 @@ public override int GetHashCode()
121121
/// <summary>
122122
/// Blake3 hash of the file.
123123
/// </summary>
124-
public string? HashBlake3 => ConnectedModelInfo?.Hashes.BLAKE3;
124+
public string? HashBlake3 => ConnectedModelInfo?.Hashes?.BLAKE3;
125+
126+
[BsonIgnore]
127+
public bool IsSafetensorFile => Path.GetExtension(RelativePath) == ".safetensors";
125128

126129
[BsonIgnore]
127130
public string? PreviewImageFullPathGlobal =>
@@ -151,6 +154,12 @@ public override int GetHashCode()
151154
[MemberNotNullWhen(true, nameof(ConnectedModelInfo))]
152155
public bool HasCivitMetadata => HasConnectedModel && ConnectedModelInfo.ModelId != null;
153156

157+
[BsonIgnore]
158+
public SafetensorMetadata? SafetensorMetadata { get; set; }
159+
160+
[BsonIgnore]
161+
public bool SafetensorMetadataParsed { get; set; }
162+
154163
public string GetFullPath(string rootModelDirectory)
155164
{
156165
return Path.Combine(rootModelDirectory, RelativePath);

0 commit comments

Comments
 (0)