Add safetensor metadata parsing and display

LykosAI · Dec 30, 2024 · f2f471b · f2f471b
1 parent d34603c
commit f2f471b
Show file tree

Hide file tree

Showing 12 changed files with 521 additions and 0 deletions.
diff --git a/Directory.Packages.props b/Directory.Packages.props
@@ -92,6 +92,7 @@
     <PackageVersion Include="Avalonia.Diagnostics" Version="$(AvaloniaVersion)" />
     <PackageVersion Include="Avalonia.HtmlRenderer" Version="11.0.0" />
     <PackageVersion Include="Avalonia.Labs.Controls" Version="11.2.0" />
+    <PackageVersion Include="Avalonia.Labs.Panels" Version="11.2.0" />
     <PackageVersion Include="Avalonia.Xaml.Behaviors" Version="11.2.0" />
     <PackageVersion Include="AvaloniaEdit.TextMate" Version="11.0.6" />
     <PackageVersion Include="TextMateSharp.Grammars" Version="1.0.56" />

diff --git a/StabilityMatrix.Avalonia/DesignData/DesignData.cs b/StabilityMatrix.Avalonia/DesignData/DesignData.cs
@@ -1167,6 +1167,50 @@ public static CompletionList SampleCompletionList
             }
         };
 
+    public static SafetensorMetadataViewModel SafetensorMetadataViewModel =>
+        DialogFactory.Get<SafetensorMetadataViewModel>(vm =>
+        {
+            vm.Metadata = new SafetensorMetadata
+            {
+                TagFrequency =
+                [
+                    ("tag1", 36),
+                    ("tag2", 33),
+                    ("tag3", 31),
+                    ("tag4", 30),
+                    ("tag5", 29),
+                    ("tag6", 28),
+                    ("tag7", 27),
+                    ("tag8", 26),
+                    ("tag9", 25),
+                    ("tag10", 24),
+                    ("tag11", 23),
+                    ("tag12", 22),
+                    ("tag13", 21),
+                    ("tag14", 20),
+                    ("tag15", 19),
+                    ("tag16", 18),
+                    ("tag17", 17),
+                    ("tag18", 16),
+                    ("tag19", 15),
+                    ("tag20", 14),
+                    ("tag21", 13),
+                    ("tag22", 12),
+                    ("tag23", 11),
+                    ("tag24", 10),
+                    ("tag25", 9),
+                    ("tag26", 8),
+                    ("tag27", 7),
+                    ("tag28", 6),
+                    ("tag29", 5),
+                    ("tag30", 4),
+                    ("tag31", 3),
+                    ("tag32", 2),
+                    ("tag33", 1),
+                ]
+            };
+        });
+
     public static ModelMetadataEditorDialogViewModel MetadataEditorDialogViewModel =>
         DialogFactory.Get<ModelMetadataEditorDialogViewModel>(vm =>
         {

diff --git a/StabilityMatrix.Avalonia/StabilityMatrix.Avalonia.csproj b/StabilityMatrix.Avalonia/StabilityMatrix.Avalonia.csproj
@@ -86,6 +86,7 @@
         <PackageReference Condition="'$(Configuration)' == 'Debug'" Include="Avalonia.Diagnostics" />
         <PackageReference Include="Avalonia.HtmlRenderer" />
         <PackageReference Include="Avalonia.Labs.Controls" />
+        <PackageReference Include="Avalonia.Labs.Panels" />
         <PackageReference Include="Avalonia.Xaml.Behaviors" />
         <PackageReference Include="AvaloniaEdit.TextMate" />
         <PackageReference Include="bodong.Avalonia.PropertyGrid" />

diff --git a/StabilityMatrix.Avalonia/ViewModels/CheckpointManager/CheckpointFileViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/CheckpointManager/CheckpointFileViewModel.cs
@@ -307,6 +307,24 @@ private async Task RenameAsync()
         }
     }
 
+    [RelayCommand]
+    private async Task OpenSafetensorMetadataViewer()
+    {
+        var vm = vmFactory.Get<SafetensorMetadataViewModel>(vm =>
+        {
+            vm.ModelName = CheckpointFile.DisplayModelName;
+            vm.Metadata = CheckpointFile.SafetensorMetadata ?? default;
+        });
+
+        var dialog = vm.GetDialog();
+        dialog.MinDialogHeight = 800;
+        dialog.MinDialogWidth = 700;
+        dialog.CloseButtonText = "Close";
+        dialog.DefaultButton = ContentDialogButton.Close;
+
+        await dialog.ShowAsync();
+    }
+
     [RelayCommand]
     private async Task OpenMetadataEditor()
     {

diff --git a/StabilityMatrix.Avalonia/ViewModels/Dialogs/SafetensorMetadataViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/Dialogs/SafetensorMetadataViewModel.cs
@@ -0,0 +1,27 @@
+using CommunityToolkit.Mvvm.ComponentModel;
+using CommunityToolkit.Mvvm.Input;
+using Injectio.Attributes;
+using StabilityMatrix.Avalonia.ViewModels.Base;
+using StabilityMatrix.Avalonia.Views.Dialogs;
+using StabilityMatrix.Core.Attributes;
+using StabilityMatrix.Core.Models;
+
+namespace StabilityMatrix.Avalonia.ViewModels.Dialogs;
+
+[View(typeof(SafetensorMetadataDialog))]
+[ManagedService]
+[RegisterSingleton<SafetensorMetadataViewModel>]
+public partial class SafetensorMetadataViewModel : ContentDialogViewModelBase
+{
+    [ObservableProperty]
+    private string? modelName;
+
+    [ObservableProperty]
+    private SafetensorMetadata metadata;
+
+    [RelayCommand]
+    public void CopyTagToClipboard(string tag)
+    {
+        App.Clipboard?.SetTextAsync(tag);
+    }
+}
diff --git a/StabilityMatrix.Avalonia/Views/CheckpointsPage.axaml b/StabilityMatrix.Avalonia/Views/CheckpointsPage.axaml
@@ -489,6 +489,11 @@
                                         </ui:MenuFlyoutItem.IsVisible>
                                     </ui:MenuFlyoutItem>
 
+                                    <ui:MenuFlyoutItem
+                                        Command="{Binding OpenSafetensorMetadataViewerCommand}"
+                                        IconSource="Tag"
+                                        IsVisible="{Binding CheckpointFile.HasSafetensorMetadata}"
+                                        Text="View Safetensor Metadata" />
                                     <ui:MenuFlyoutItem
                                         Command="{Binding OpenMetadataEditorCommand}"
                                         IconSource="Edit"

diff --git a/StabilityMatrix.Avalonia/Views/Dialogs/SafetensorMetadataDialog.axaml b/StabilityMatrix.Avalonia/Views/Dialogs/SafetensorMetadataDialog.axaml
@@ -0,0 +1,101 @@
+<controls:UserControlBase
+    x:Class="StabilityMatrix.Avalonia.Views.Dialogs.SafetensorMetadataDialog"
+    xmlns="https://github.com/avaloniaui"
+    xmlns:x="http://schemas.microsoft.com/winfx/2006/xaml"
+    xmlns:controls="clr-namespace:StabilityMatrix.Avalonia.Controls"
+    xmlns:d="http://schemas.microsoft.com/expression/blend/2008"
+    xmlns:dialogs="clr-namespace:StabilityMatrix.Avalonia.ViewModels.Dialogs"
+    xmlns:lang="clr-namespace:StabilityMatrix.Avalonia.Languages"
+    xmlns:mc="http://schemas.openxmlformats.org/markup-compatibility/2006"
+    xmlns:mocks="clr-namespace:StabilityMatrix.Avalonia.DesignData"
+    xmlns:models="clr-namespace:StabilityMatrix.Core.Models;assembly=StabilityMatrix.Core"
+    xmlns:panels="using:Avalonia.Labs.Panels"
+    xmlns:ui="using:FluentAvalonia.UI.Controls"
+    d:DataContext="{x:Static mocks:DesignData.SafetensorMetadataViewModel}"
+    d:DesignHeight="550"
+    d:DesignWidth="700"
+    x:DataType="dialogs:SafetensorMetadataViewModel"
+    mc:Ignorable="d">
+    <Grid RowDefinitions="Auto,Auto,Auto,Auto,Auto,Auto">
+        <TextBlock
+            Margin="8"
+            HorizontalAlignment="Center"
+            FontSize="24"
+            FontWeight="SemiBold"
+            Text="Safetensor Metadata" />
+        <TextBlock
+            Grid.Row="1"
+            Margin="8"
+            HorizontalAlignment="Center"
+            FontSize="16"
+            FontWeight="SemiBold"
+            Text="{Binding ModelName}" />
+
+        <!--  List of tags  -->
+        <TextBlock
+            Grid.Row="2"
+            Margin="8"
+            HorizontalAlignment="Left"
+            FontSize="16"
+            FontWeight="SemiBold"
+            IsVisible="{Binding Metadata.TagFrequency, Converter={x:Static ObjectConverters.IsNotNull}}"
+            Text="Trained Tags" />
+        <ItemsControl
+            Grid.Row="3"
+            Margin="8"
+            IsVisible="{Binding Metadata.TagFrequency, Converter={x:Static ObjectConverters.IsNotNull}}"
+            ItemsSource="{Binding Metadata.TagFrequency}">
+            <ItemsControl.ItemTemplate>
+                <DataTemplate>
+                    <Button
+                        Command="{Binding $parent[ItemsControl].((dialogs:SafetensorMetadataViewModel)DataContext).CopyTagToClipboardCommand}"
+                        CommandParameter="{Binding Name}"
+                        Cursor="Hand">
+                        <StackPanel Orientation="Horizontal">
+                            <TextBlock FontWeight="SemiBold" Text="{Binding Name}" />
+                            <TextBlock Margin="5,0,0,0" Text="{Binding Frequency}" />
+                        </StackPanel>
+                    </Button>
+                </DataTemplate>
+            </ItemsControl.ItemTemplate>
+            <ItemsControl.ItemsPanel>
+                <ItemsPanelTemplate>
+                    <panels:FlexPanel
+                        ColumnSpacing="5"
+                        RowSpacing="5"
+                        Wrap="Wrap" />
+                </ItemsPanelTemplate>
+            </ItemsControl.ItemsPanel>
+        </ItemsControl>
+
+        <!--  All other metadata  -->
+        <TextBlock
+            Grid.Row="4"
+            Margin="8"
+            HorizontalAlignment="Left"
+            FontSize="16"
+            FontWeight="SemiBold"
+            Text="Other Metadata" />
+
+        <TextBlock
+            Grid.Row="5"
+            Margin="8"
+            FontSize="16"
+            FontStyle="Italic"
+            IsVisible="{Binding !Metadata.OtherMetadata.Count}"
+            Text="No Metadata" />
+        <ItemsControl
+            Grid.Row="5"
+            Margin="8"
+            ItemsSource="{Binding Metadata.OtherMetadata}">
+            <ItemsControl.ItemTemplate>
+                <DataTemplate>
+                    <StackPanel Margin="5" Orientation="Vertical">
+                        <TextBlock FontWeight="SemiBold" Text="{Binding Name}" />
+                        <TextBlock Text="{Binding Value}" />
+                    </StackPanel>
+                </DataTemplate>
+            </ItemsControl.ItemTemplate>
+        </ItemsControl>
+    </Grid>
+</controls:UserControlBase>
diff --git a/StabilityMatrix.Avalonia/Views/Dialogs/SafetensorMetadataDialog.axaml.cs b/StabilityMatrix.Avalonia/Views/Dialogs/SafetensorMetadataDialog.axaml.cs
@@ -0,0 +1,19 @@
+using Avalonia.Controls;
+using Avalonia.Markup.Xaml;
+using Injectio.Attributes;
+
+namespace StabilityMatrix.Avalonia.Views.Dialogs;
+
+[RegisterTransient<SafetensorMetadataDialog>]
+public partial class SafetensorMetadataDialog : UserControl
+{
+    public SafetensorMetadataDialog()
+    {
+        InitializeComponent();
+    }
+
+    private void InitializeComponent()
+    {
+        AvaloniaXamlLoader.Load(this);
+    }
+}
diff --git a/StabilityMatrix.Core/Models/Database/LocalModelFile.cs b/StabilityMatrix.Core/Models/Database/LocalModelFile.cs
@@ -68,6 +68,12 @@ public override int GetHashCode()
     /// </summary>
     public ConnectedModelInfo? ConnectedModelInfo { get; set; }
 
+    /// <summary>
+    /// Optional Safetensor metadata.
+    /// </summary>
+    [BsonIgnore]
+    public SafetensorMetadata? SafetensorMetadata { get; set; }
+
     /// <summary>
     /// Optional preview image relative path.
     /// </summary>
@@ -151,6 +157,10 @@ public override int GetHashCode()
     [MemberNotNullWhen(true, nameof(ConnectedModelInfo))]
     public bool HasCivitMetadata => HasConnectedModel && ConnectedModelInfo.ModelId != null;
 
+    [BsonIgnore]
+    [MemberNotNullWhen(true, nameof(SafetensorMetadata))]
+    public bool HasSafetensorMetadata => SafetensorMetadata != null;
+
     public string GetFullPath(string rootModelDirectory)
     {
         return Path.Combine(rootModelDirectory, RelativePath);