Skip to content

Commit 149a109

Browse files
committed
Add safetensor metadata parsing and display
1 parent a353386 commit 149a109

12 files changed

Lines changed: 521 additions & 0 deletions

File tree

Directory.Packages.props

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -92,6 +92,7 @@
9292
<PackageVersion Include="Avalonia.Diagnostics" Version="$(AvaloniaVersion)" />
9393
<PackageVersion Include="Avalonia.HtmlRenderer" Version="11.0.0" />
9494
<PackageVersion Include="Avalonia.Labs.Controls" Version="11.2.0" />
95+
<PackageVersion Include="Avalonia.Labs.Panels" Version="11.2.0" />
9596
<PackageVersion Include="Avalonia.Xaml.Behaviors" Version="11.2.0" />
9697
<PackageVersion Include="AvaloniaEdit.TextMate" Version="11.0.6" />
9798
<PackageVersion Include="TextMateSharp.Grammars" Version="1.0.56" />

StabilityMatrix.Avalonia/DesignData/DesignData.cs

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1167,6 +1167,50 @@ 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 =
1176+
[
1177+
("tag1", 36),
1178+
("tag2", 33),
1179+
("tag3", 31),
1180+
("tag4", 30),
1181+
("tag5", 29),
1182+
("tag6", 28),
1183+
("tag7", 27),
1184+
("tag8", 26),
1185+
("tag9", 25),
1186+
("tag10", 24),
1187+
("tag11", 23),
1188+
("tag12", 22),
1189+
("tag13", 21),
1190+
("tag14", 20),
1191+
("tag15", 19),
1192+
("tag16", 18),
1193+
("tag17", 17),
1194+
("tag18", 16),
1195+
("tag19", 15),
1196+
("tag20", 14),
1197+
("tag21", 13),
1198+
("tag22", 12),
1199+
("tag23", 11),
1200+
("tag24", 10),
1201+
("tag25", 9),
1202+
("tag26", 8),
1203+
("tag27", 7),
1204+
("tag28", 6),
1205+
("tag29", 5),
1206+
("tag30", 4),
1207+
("tag31", 3),
1208+
("tag32", 2),
1209+
("tag33", 1),
1210+
]
1211+
};
1212+
});
1213+
11701214
public static ModelMetadataEditorDialogViewModel MetadataEditorDialogViewModel =>
11711215
DialogFactory.Get<ModelMetadataEditorDialogViewModel>(vm =>
11721216
{

StabilityMatrix.Avalonia/StabilityMatrix.Avalonia.csproj

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,7 @@
8686
<PackageReference Condition="'$(Configuration)' == 'Debug'" Include="Avalonia.Diagnostics" />
8787
<PackageReference Include="Avalonia.HtmlRenderer" />
8888
<PackageReference Include="Avalonia.Labs.Controls" />
89+
<PackageReference Include="Avalonia.Labs.Panels" />
8990
<PackageReference Include="Avalonia.Xaml.Behaviors" />
9091
<PackageReference Include="AvaloniaEdit.TextMate" />
9192
<PackageReference Include="bodong.Avalonia.PropertyGrid" />

StabilityMatrix.Avalonia/ViewModels/CheckpointManager/CheckpointFileViewModel.cs

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

310+
[RelayCommand]
311+
private async Task OpenSafetensorMetadataViewer()
312+
{
313+
var vm = vmFactory.Get<SafetensorMetadataViewModel>(vm =>
314+
{
315+
vm.ModelName = CheckpointFile.DisplayModelName;
316+
vm.Metadata = CheckpointFile.SafetensorMetadata ?? default;
317+
});
318+
319+
var dialog = vm.GetDialog();
320+
dialog.MinDialogHeight = 800;
321+
dialog.MinDialogWidth = 700;
322+
dialog.CloseButtonText = "Close";
323+
dialog.DefaultButton = ContentDialogButton.Close;
324+
325+
await dialog.ShowAsync();
326+
}
327+
310328
[RelayCommand]
311329
private async Task OpenMetadataEditor()
312330
{
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.HasSafetensorMetadata}"
496+
Text="View Safetensor Metadata" />
492497
<ui:MenuFlyoutItem
493498
Command="{Binding OpenMetadataEditorCommand}"
494499
IconSource="Edit"
Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
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:panels="using:Avalonia.Labs.Panels"
13+
xmlns:ui="using:FluentAvalonia.UI.Controls"
14+
d:DataContext="{x:Static mocks:DesignData.SafetensorMetadataViewModel}"
15+
d:DesignHeight="550"
16+
d:DesignWidth="700"
17+
x:DataType="dialogs:SafetensorMetadataViewModel"
18+
mc:Ignorable="d">
19+
<Grid RowDefinitions="Auto,Auto,Auto,Auto,Auto,Auto">
20+
<TextBlock
21+
Margin="8"
22+
HorizontalAlignment="Center"
23+
FontSize="24"
24+
FontWeight="SemiBold"
25+
Text="Safetensor Metadata" />
26+
<TextBlock
27+
Grid.Row="1"
28+
Margin="8"
29+
HorizontalAlignment="Center"
30+
FontSize="16"
31+
FontWeight="SemiBold"
32+
Text="{Binding ModelName}" />
33+
34+
<!-- List of tags -->
35+
<TextBlock
36+
Grid.Row="2"
37+
Margin="8"
38+
HorizontalAlignment="Left"
39+
FontSize="16"
40+
FontWeight="SemiBold"
41+
IsVisible="{Binding Metadata.TagFrequency, Converter={x:Static ObjectConverters.IsNotNull}}"
42+
Text="Trained Tags" />
43+
<ItemsControl
44+
Grid.Row="3"
45+
Margin="8"
46+
IsVisible="{Binding Metadata.TagFrequency, Converter={x:Static ObjectConverters.IsNotNull}}"
47+
ItemsSource="{Binding Metadata.TagFrequency}">
48+
<ItemsControl.ItemTemplate>
49+
<DataTemplate>
50+
<Button
51+
Command="{Binding $parent[ItemsControl].((dialogs:SafetensorMetadataViewModel)DataContext).CopyTagToClipboardCommand}"
52+
CommandParameter="{Binding Name}"
53+
Cursor="Hand">
54+
<StackPanel Orientation="Horizontal">
55+
<TextBlock FontWeight="SemiBold" Text="{Binding Name}" />
56+
<TextBlock Margin="5,0,0,0" Text="{Binding Frequency}" />
57+
</StackPanel>
58+
</Button>
59+
</DataTemplate>
60+
</ItemsControl.ItemTemplate>
61+
<ItemsControl.ItemsPanel>
62+
<ItemsPanelTemplate>
63+
<panels:FlexPanel
64+
ColumnSpacing="5"
65+
RowSpacing="5"
66+
Wrap="Wrap" />
67+
</ItemsPanelTemplate>
68+
</ItemsControl.ItemsPanel>
69+
</ItemsControl>
70+
71+
<!-- All other metadata -->
72+
<TextBlock
73+
Grid.Row="4"
74+
Margin="8"
75+
HorizontalAlignment="Left"
76+
FontSize="16"
77+
FontWeight="SemiBold"
78+
Text="Other Metadata" />
79+
80+
<TextBlock
81+
Grid.Row="5"
82+
Margin="8"
83+
FontSize="16"
84+
FontStyle="Italic"
85+
IsVisible="{Binding !Metadata.OtherMetadata.Count}"
86+
Text="No Metadata" />
87+
<ItemsControl
88+
Grid.Row="5"
89+
Margin="8"
90+
ItemsSource="{Binding Metadata.OtherMetadata}">
91+
<ItemsControl.ItemTemplate>
92+
<DataTemplate>
93+
<StackPanel Margin="5" Orientation="Vertical">
94+
<TextBlock FontWeight="SemiBold" Text="{Binding Name}" />
95+
<TextBlock Text="{Binding Value}" />
96+
</StackPanel>
97+
</DataTemplate>
98+
</ItemsControl.ItemTemplate>
99+
</ItemsControl>
100+
</Grid>
101+
</controls:UserControlBase>
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/Models/Database/LocalModelFile.cs

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,12 @@ public override int GetHashCode()
6868
/// </summary>
6969
public ConnectedModelInfo? ConnectedModelInfo { get; set; }
7070

71+
/// <summary>
72+
/// Optional Safetensor metadata.
73+
/// </summary>
74+
[BsonIgnore]
75+
public SafetensorMetadata? SafetensorMetadata { get; set; }
76+
7177
/// <summary>
7278
/// Optional preview image relative path.
7379
/// </summary>
@@ -151,6 +157,10 @@ public override int GetHashCode()
151157
[MemberNotNullWhen(true, nameof(ConnectedModelInfo))]
152158
public bool HasCivitMetadata => HasConnectedModel && ConnectedModelInfo.ModelId != null;
153159

160+
[BsonIgnore]
161+
[MemberNotNullWhen(true, nameof(SafetensorMetadata))]
162+
public bool HasSafetensorMetadata => SafetensorMetadata != null;
163+
154164
public string GetFullPath(string rootModelDirectory)
155165
{
156166
return Path.Combine(rootModelDirectory, RelativePath);

0 commit comments

Comments
 (0)