Skip to content

Commit 6a42aa6

Browse files
GenteureGenteure
committed
Parse and load safetensor metadata in background
1 parent 149a109 commit 6a42aa6

10 files changed

Lines changed: 294 additions & 226 deletions

File tree

‎StabilityMatrix.Avalonia/DesignData/DesignData.cs‎

Lines changed: 10 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -1172,42 +1172,16 @@ public static CompletionList SampleCompletionList
11721172
{
11731173
vm.Metadata = new SafetensorMetadata
11741174
{
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-
]
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+
}
12111185
};
12121186
});
12131187

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

Lines changed: 32 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -310,10 +310,41 @@ private async Task RenameAsync()
310310
[RelayCommand]
311311
private async Task OpenSafetensorMetadataViewer()
312312
{
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+
313344
var vm = vmFactory.Get<SafetensorMetadataViewModel>(vm =>
314345
{
315346
vm.ModelName = CheckpointFile.DisplayModelName;
316-
vm.Metadata = CheckpointFile.SafetensorMetadata ?? default;
347+
vm.Metadata = CheckpointFile.SafetensorMetadata;
317348
});
318349

319350
var dialog = vm.GetDialog();

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ public partial class SafetensorMetadataViewModel : ContentDialogViewModelBase
1717
private string? modelName;
1818

1919
[ObservableProperty]
20-
private SafetensorMetadata metadata;
20+
private SafetensorMetadata? metadata;
2121

2222
[RelayCommand]
2323
public void CopyTagToClipboard(string tag)

‎StabilityMatrix.Avalonia/Views/CheckpointsPage.axaml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -492,7 +492,7 @@
492492
<ui:MenuFlyoutItem
493493
Command="{Binding OpenSafetensorMetadataViewerCommand}"
494494
IconSource="Tag"
495-
IsVisible="{Binding CheckpointFile.HasSafetensorMetadata}"
495+
IsVisible="{Binding CheckpointFile.IsSafetensorFile}"
496496
Text="View Safetensor Metadata" />
497497
<ui:MenuFlyoutItem
498498
Command="{Binding OpenMetadataEditorCommand}"

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

Lines changed: 70 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
d:DesignWidth="700"
1717
x:DataType="dialogs:SafetensorMetadataViewModel"
1818
mc:Ignorable="d">
19-
<Grid RowDefinitions="Auto,Auto,Auto,Auto,Auto,Auto">
19+
<Grid RowDefinitions="Auto,Auto,Auto">
2020
<TextBlock
2121
Margin="8"
2222
HorizontalAlignment="Center"
@@ -31,71 +31,79 @@
3131
FontWeight="SemiBold"
3232
Text="{Binding ModelName}" />
3333

34-
<!-- List of tags -->
3534
<TextBlock
3635
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"
36+
Margin="10,20"
37+
HorizontalAlignment="Center"
8338
FontSize="16"
8439
FontStyle="Italic"
85-
IsVisible="{Binding !Metadata.OtherMetadata.Count}"
40+
IsVisible="{Binding Metadata, Converter={x:Static ObjectConverters.IsNull}}"
8641
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>
42+
43+
<Grid
44+
Grid.Row="2"
45+
IsVisible="{Binding Metadata, Converter={x:Static ObjectConverters.IsNotNull}}"
46+
RowDefinitions="Auto,Auto,Auto,Auto">
47+
48+
<!-- List of tags -->
49+
<TextBlock
50+
Grid.Row="0"
51+
Margin="8"
52+
HorizontalAlignment="Left"
53+
FontSize="16"
54+
FontWeight="SemiBold"
55+
IsVisible="{Binding Metadata.TagFrequency, Converter={x:Static ObjectConverters.IsNotNull}}"
56+
Text="Trained Tags" />
57+
<ItemsControl
58+
Grid.Row="1"
59+
Margin="8"
60+
IsVisible="{Binding Metadata.TagFrequency, Converter={x:Static ObjectConverters.IsNotNull}}"
61+
ItemsSource="{Binding Metadata.TagFrequency}">
62+
<ItemsControl.ItemTemplate>
63+
<DataTemplate>
64+
<Button
65+
Command="{Binding $parent[ItemsControl].((dialogs:SafetensorMetadataViewModel)DataContext).CopyTagToClipboardCommand}"
66+
CommandParameter="{Binding Name}"
67+
Cursor="Hand">
68+
<StackPanel Orientation="Horizontal">
69+
<TextBlock FontWeight="SemiBold" Text="{Binding Name}" />
70+
<TextBlock Margin="5,0,0,0" Text="{Binding Frequency}" />
71+
</StackPanel>
72+
</Button>
73+
</DataTemplate>
74+
</ItemsControl.ItemTemplate>
75+
<ItemsControl.ItemsPanel>
76+
<ItemsPanelTemplate>
77+
<panels:FlexPanel
78+
ColumnSpacing="5"
79+
RowSpacing="5"
80+
Wrap="Wrap" />
81+
</ItemsPanelTemplate>
82+
</ItemsControl.ItemsPanel>
83+
</ItemsControl>
84+
85+
<!-- All other metadata -->
86+
<TextBlock
87+
Grid.Row="2"
88+
Margin="8"
89+
FontSize="16"
90+
FontWeight="SemiBold"
91+
IsVisible="{Binding !!Metadata.OtherMetadata.Count}"
92+
Text="Other Metadata" />
93+
<ItemsControl
94+
Grid.Row="3"
95+
Margin="8"
96+
IsVisible="{Binding !!Metadata.OtherMetadata.Count}"
97+
ItemsSource="{Binding Metadata.OtherMetadata}">
98+
<ItemsControl.ItemTemplate>
99+
<DataTemplate>
100+
<StackPanel Margin="5" Orientation="Vertical">
101+
<TextBlock FontWeight="SemiBold" Text="{Binding Name}" />
102+
<TextBlock Text="{Binding Value}" />
103+
</StackPanel>
104+
</DataTemplate>
105+
</ItemsControl.ItemTemplate>
106+
</ItemsControl>
107+
</Grid>
100108
</Grid>
101109
</controls:UserControlBase>

‎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: 8 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -68,12 +68,6 @@ 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-
7771
/// <summary>
7872
/// Optional preview image relative path.
7973
/// </summary>
@@ -127,7 +121,10 @@ public override int GetHashCode()
127121
/// <summary>
128122
/// Blake3 hash of the file.
129123
/// </summary>
130-
public string? HashBlake3 => ConnectedModelInfo?.Hashes.BLAKE3;
124+
public string? HashBlake3 => ConnectedModelInfo?.Hashes?.BLAKE3;
125+
126+
[BsonIgnore]
127+
public bool IsSafetensorFile => Path.GetExtension(RelativePath) == ".safetensors";
131128

132129
[BsonIgnore]
133130
public string? PreviewImageFullPathGlobal =>
@@ -158,8 +155,10 @@ public override int GetHashCode()
158155
public bool HasCivitMetadata => HasConnectedModel && ConnectedModelInfo.ModelId != null;
159156

160157
[BsonIgnore]
161-
[MemberNotNullWhen(true, nameof(SafetensorMetadata))]
162-
public bool HasSafetensorMetadata => SafetensorMetadata != null;
158+
public SafetensorMetadata? SafetensorMetadata { get; set; }
159+
160+
[BsonIgnore]
161+
public bool SafetensorMetadataParsed { get; set; }
163162

164163
public string GetFullPath(string rootModelDirectory)
165164
{

0 commit comments

Comments
 (0)