Merge pull request #296 from LykosAI/main

ionite34 · web-flow · commit 24cb710ada18 · 2023-11-27T16:00:29.000-05:00
diff --git a/CHANGELOG.md b/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.6.5
+### Fixed
+- Fixed error when receiving unknown model format values from the Model Browser
+- Fixed process errors when installing or updating Pip packages using the Python packages dialog
+
 ## v2.6.4
 ### Fixed
 - Fixed errors preventing Model Browser from finding results with certain search queries
diff --git a/StabilityMatrix.Avalonia/App.axaml.cs b/StabilityMatrix.Avalonia/App.axaml.cs
@@ -424,6 +424,7 @@ internal static IServiceCollection ConfigureServices()
         jsonSerializerOptions.Converters.Add(new ObjectToInferredTypesConverter());
         jsonSerializerOptions.Converters.Add(new DefaultUnknownEnumConverter<CivitFileType>());
         jsonSerializerOptions.Converters.Add(new DefaultUnknownEnumConverter<CivitModelType>());
+        jsonSerializerOptions.Converters.Add(new DefaultUnknownEnumConverter<CivitModelFormat>());
         jsonSerializerOptions.Converters.Add(
             new JsonStringEnumConverter(JsonNamingPolicy.CamelCase)
         );
diff --git a/StabilityMatrix.Core/Converters/Json/DefaultUnknownEnumConverter.cs b/StabilityMatrix.Core/Converters/Json/DefaultUnknownEnumConverter.cs
@@ -1,38 +1,110 @@
-﻿using System.Text.Json;
+﻿using System.Diagnostics;
+using System.Diagnostics.CodeAnalysis;
+using System.Reflection;
+using System.Runtime.Serialization;
+using System.Text.Json;
 using System.Text.Json.Serialization;
-using StabilityMatrix.Core.Extensions;
 
 namespace StabilityMatrix.Core.Converters.Json;
 
-public class DefaultUnknownEnumConverter<T> : JsonConverter<T>
+public class DefaultUnknownEnumConverter<
+    [DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicFields)] T
+> : JsonConverter<T>
     where T : Enum
 {
+    /// <summary>
+    /// Lazy initialization for <see cref="EnumMemberValues"/>.
+    /// </summary>
+    private readonly Lazy<Dictionary<string, T>> _enumMemberValuesLazy =
+        new(
+            () =>
+                typeof(T)
+                    .GetFields()
+                    .Where(field => field.IsStatic)
+                    .Select(
+                        field =>
+                            new
+                            {
+                                FieldName = field.Name,
+                                FieldValue = (T)field.GetValue(null)!,
+                                EnumMemberValue = field
+                                    .GetCustomAttributes<EnumMemberAttribute>(false)
+                                    .FirstOrDefault()
+                                    ?.Value?.ToString()
+                            }
+                    )
+                    .ToDictionary(x => x.EnumMemberValue ?? x.FieldName, x => x.FieldValue)
+        );
+
+    /// <summary>
+    /// Gets a dictionary of enum member values, keyed by the EnumMember attribute value, or the field name if no EnumMember attribute is present.
+    /// </summary>
+    private Dictionary<string, T> EnumMemberValues => _enumMemberValuesLazy.Value;
+
+    /// <summary>
+    /// Lazy initialization for <see cref="EnumMemberNames"/>.
+    /// </summary>
+    private readonly Lazy<Dictionary<T, string>> _enumMemberNamesLazy;
+
+    /// <summary>
+    /// Gets a dictionary of enum member names, keyed by the enum member value.
+    /// </summary>
+    private Dictionary<T, string> EnumMemberNames => _enumMemberNamesLazy.Value;
+
+    /// <summary>
+    /// Gets the value of the "Unknown" enum member, or the 0 value if no "Unknown" member is present.
+    /// </summary>
+    private T UnknownValue =>
+        EnumMemberValues.TryGetValue("Unknown", out var res) ? res : (T)Enum.ToObject(typeof(T), 0);
+
+    /// <inheritdoc />
+    public override bool HandleNull => true;
+
+    public DefaultUnknownEnumConverter()
+    {
+        _enumMemberNamesLazy = new Lazy<Dictionary<T, string>>(
+            () => EnumMemberValues.ToDictionary(x => x.Value, x => x.Key)
+        );
+    }
+
+    /// <inheritdoc />
     public override T Read(
         ref Utf8JsonReader reader,
         Type typeToConvert,
         JsonSerializerOptions options
     )
     {
-        if (reader.TokenType != JsonTokenType.String)
+        if (reader.TokenType is not (JsonTokenType.String or JsonTokenType.PropertyName))
         {
-            throw new JsonException();
+            throw new JsonException("Expected String or PropertyName token");
         }
 
-        var enumText = reader.GetString()?.Replace(" ", "_");
-        if (Enum.TryParse(typeof(T), enumText, true, out var result))
+        if (reader.GetString() is { } readerString)
         {
-            return (T)result!;
-        }
+            // First try get exact match
+            if (EnumMemberValues.TryGetValue(readerString, out var enumMemberValue))
+            {
+                return enumMemberValue;
+            }
 
-        // Unknown value handling
-        if (Enum.TryParse(typeof(T), "Unknown", true, out var unknownResult))
-        {
-            return (T)unknownResult!;
+            // Otherwise try get case-insensitive match
+            if (
+                EnumMemberValues.Keys.FirstOrDefault(
+                    key => key.Equals(readerString, StringComparison.OrdinalIgnoreCase)
+                ) is
+                { } enumMemberName
+            )
+            {
+                return EnumMemberValues[enumMemberName];
+            }
+
+            Debug.WriteLine($"Unknown enum member value for {typeToConvert}: {readerString}");
         }
 
-        throw new JsonException($"Unable to parse '{enumText}' to enum '{typeof(T)}'.");
+        return UnknownValue;
     }
 
+    /// <inheritdoc />
     public override void Write(Utf8JsonWriter writer, T? value, JsonSerializerOptions options)
     {
         if (value == null)
@@ -41,49 +113,20 @@ public override void Write(Utf8JsonWriter writer, T? value, JsonSerializerOption
             return;
         }
 
-        writer.WriteStringValue(value.GetStringValue().Replace("_", " "));
+        writer.WriteStringValue(EnumMemberNames[value]);
     }
 
     /// <inheritdoc />
     public override T ReadAsPropertyName(
         ref Utf8JsonReader reader,
         Type typeToConvert,
         JsonSerializerOptions options
-    )
-    {
-        if (reader.TokenType != JsonTokenType.PropertyName)
-        {
-            throw new JsonException();
-        }
-
-        var enumText = reader.GetString()?.Replace(" ", "_");
-        if (Enum.TryParse(typeof(T), enumText, true, out var result))
-        {
-            return (T)result!;
-        }
-
-        // Unknown value handling
-        if (Enum.TryParse(typeof(T), "Unknown", true, out var unknownResult))
-        {
-            return (T)unknownResult!;
-        }
-
-        throw new JsonException($"Unable to parse '{enumText}' to enum '{typeof(T)}'.");
-    }
+    ) => Read(ref reader, typeToConvert, options);
 
     /// <inheritdoc />
     public override void WriteAsPropertyName(
         Utf8JsonWriter writer,
         T? value,
         JsonSerializerOptions options
-    )
-    {
-        if (value == null)
-        {
-            writer.WriteNullValue();
-            return;
-        }
-
-        writer.WritePropertyName(value.GetStringValue().Replace("_", " "));
-    }
+    ) => Write(writer, value, options);
 }
diff --git a/StabilityMatrix.Core/Helper/ArchiveHelper.cs b/StabilityMatrix.Core/Helper/ArchiveHelper.cs
@@ -86,22 +86,16 @@ public static async Task AddToArchive7Z(string archivePath, string sourceDirecto
 
     public static async Task<ArchiveInfo> Extract7Z(string archivePath, string extractDirectory)
     {
-        var args =
-            $"x {ProcessRunner.Quote(archivePath)} -o{ProcessRunner.Quote(extractDirectory)} -y";
+        var result = await ProcessRunner
+            .GetProcessResultAsync(
+                SevenZipPath,
+                new[] { "x", archivePath, "-o" + ProcessRunner.Quote(extractDirectory), "-y" }
+            )
+            .ConfigureAwait(false);
 
-        Logger.Debug($"Starting process '{SevenZipPath}' with arguments '{args}'");
+        result.EnsureSuccessExitCode();
 
-        using var process = new Process();
-        process.StartInfo = new ProcessStartInfo(SevenZipPath, args)
-        {
-            RedirectStandardOutput = true,
-            RedirectStandardError = true,
-            UseShellExecute = false,
-            CreateNoWindow = true
-        };
-        process.Start();
-        await ProcessRunner.WaitForExitConditionAsync(process);
-        var output = await process.StandardOutput.ReadToEndAsync();
+        var output = result.StandardOutput ?? "";
 
         try
         {
@@ -153,8 +147,12 @@ IProgress<ProgressReport> progress
             $"x {ProcessRunner.Quote(archivePath)} -o{ProcessRunner.Quote(extractDirectory)} -y -bsp1";
         Logger.Debug($"Starting process '{SevenZipPath}' with arguments '{args}'");
 
-        var process = ProcessRunner.StartProcess(SevenZipPath, args, outputDataReceived: onOutput);
-        await ProcessRunner.WaitForExitConditionAsync(process);
+        using var process = ProcessRunner.StartProcess(
+            SevenZipPath,
+            args,
+            outputDataReceived: onOutput
+        );
+        await ProcessRunner.WaitForExitConditionAsync(process).ConfigureAwait(false);
 
         progress.Report(new ProgressReport(1f, "Finished extracting", type: ProgressType.Extract));
 
diff --git a/StabilityMatrix.Core/Models/Api/CivitFileType.cs b/StabilityMatrix.Core/Models/Api/CivitFileType.cs
@@ -1,13 +1,16 @@
-﻿using System.Text.Json.Serialization;
+﻿using System.Runtime.Serialization;
+using System.Text.Json.Serialization;
 using StabilityMatrix.Core.Converters.Json;
 
 namespace StabilityMatrix.Core.Models.Api;
 
 [JsonConverter(typeof(DefaultUnknownEnumConverter<CivitFileType>))]
 public enum CivitFileType
 {
+    Unknown,
     Model,
     VAE,
-    Training_Data,
-    Unknown,
+
+    [EnumMember(Value = "Training Data")]
+    TrainingData
 }
diff --git a/StabilityMatrix.Core/Models/Api/CivitModelFormat.cs b/StabilityMatrix.Core/Models/Api/CivitModelFormat.cs
@@ -1,12 +1,14 @@
 ﻿using System.Text.Json.Serialization;
+using StabilityMatrix.Core.Converters.Json;
 
 namespace StabilityMatrix.Core.Models.Api;
 
-
-[JsonConverter(typeof(JsonStringEnumConverter))]
+[JsonConverter(typeof(DefaultUnknownEnumConverter<CivitModelFormat>))]
 public enum CivitModelFormat
 {
+    Unknown,
     SafeTensor,
     PickleTensor,
+    Diffusers,
     Other
 }
diff --git a/StabilityMatrix.Core/Models/Api/CivitModelType.cs b/StabilityMatrix.Core/Models/Api/CivitModelType.cs
@@ -9,6 +9,8 @@ namespace StabilityMatrix.Core.Models.Api;
 [SuppressMessage("ReSharper", "InconsistentNaming")]
 public enum CivitModelType
 {
+    Unknown,
+
     [ConvertTo<SharedFolderType>(SharedFolderType.StableDiffusion)]
     Checkpoint,
 
@@ -39,6 +41,5 @@ public enum CivitModelType
     Wildcards,
     Workflows,
     Other,
-    All,
-    Unknown
+    All
 }
diff --git a/StabilityMatrix.Core/Processes/ProcessRunner.cs b/StabilityMatrix.Core/Processes/ProcessRunner.cs
@@ -401,35 +401,6 @@ public static string Quote(string argument)
         return inner.Contains(' ') ? $"\"{inner}\"" : argument;
     }
 
-    /// <summary>
-    /// Check if the process exited with the expected exit code.
-    /// </summary>
-    /// <param name="process">Process to check.</param>
-    /// <param name="expectedExitCode">Expected exit code.</param>
-    /// <param name="stdout">Process stdout.</param>
-    /// <param name="stderr">Process stderr.</param>
-    /// <exception cref="ProcessException">Thrown if exit code does not match expected value.</exception>
-    // ReSharper disable once MemberCanBePrivate.Global
-    public static Task ValidateExitConditionAsync(
-        Process process,
-        int expectedExitCode = 0,
-        string? stdout = null,
-        string? stderr = null
-    )
-    {
-        var exitCode = process.ExitCode;
-        if (exitCode == expectedExitCode)
-        {
-            return Task.CompletedTask;
-        }
-
-        var pName = process.StartInfo.FileName;
-        var msg =
-            $"Process {pName} failed with exit-code {exitCode}. stdout: '{stdout}', stderr: '{stderr}'";
-        Logger.Error(msg);
-        throw new ProcessException(msg);
-    }
-
     /// <summary>
     /// Waits for process to exit, then validates exit code.
     /// </summary>
@@ -443,25 +414,28 @@ public static async Task WaitForExitConditionAsync(
         CancellationToken cancelToken = default
     )
     {
-        if (process is AnsiProcess)
+        if (!process.HasExited)
         {
-            throw new ArgumentException(
-                $"{nameof(WaitForExitConditionAsync)} does not support AnsiProcess, which uses custom async data reading",
-                nameof(process)
-            );
+            await process.WaitForExitAsync(cancelToken).ConfigureAwait(false);
         }
 
-        var stdout = new StringBuilder();
-        var stderr = new StringBuilder();
-        process.OutputDataReceived += (_, args) => stdout.Append(args.Data);
-        process.ErrorDataReceived += (_, args) => stderr.Append(args.Data);
-        await process.WaitForExitAsync(cancelToken).ConfigureAwait(false);
-        await ValidateExitConditionAsync(
-                process,
-                expectedExitCode,
-                stdout.ToString(),
-                stderr.ToString()
-            )
-            .ConfigureAwait(false);
+        if (process.ExitCode == expectedExitCode)
+        {
+            return;
+        }
+
+        // Accessing ProcessName may error on some platforms
+        string? processName = null;
+        try
+        {
+            processName = process.ProcessName;
+        }
+        catch (SystemException) { }
+
+        throw new ProcessException(
+            "Process "
+                + (processName == null ? "" : processName + " ")
+                + $"failed with exit-code {process.ExitCode}."
+        );
     }
 }
diff --git a/StabilityMatrix.Tests/Core/DefaultUnknownEnumConverterTests.cs b/StabilityMatrix.Tests/Core/DefaultUnknownEnumConverterTests.cs
