From 4897b2193023116afca215db4cd610bca798d485 Mon Sep 17 00:00:00 2001 From: genteure Date: Thu, 26 Dec 2024 20:36:17 +0800 Subject: [PATCH] Improve GenerationParameters fields parsing --- .../Models/GenerationParameters.cs | 126 +++++++++++-- .../Models/GenerationParametersTests.cs | 168 ++++++++++++++++-- 2 files changed, 270 insertions(+), 24 deletions(-) diff --git a/StabilityMatrix.Core/Models/GenerationParameters.cs b/StabilityMatrix.Core/Models/GenerationParameters.cs index 59fe7a95f..f17b2da23 100644 --- a/StabilityMatrix.Core/Models/GenerationParameters.cs +++ b/StabilityMatrix.Core/Models/GenerationParameters.cs @@ -2,13 +2,12 @@ using System.Diagnostics.CodeAnalysis; using System.Text.Json.Nodes; using System.Text.Json.Serialization; -using System.Text.RegularExpressions; using StabilityMatrix.Core.Models.Api.Comfy; namespace StabilityMatrix.Core.Models; [JsonSerializable(typeof(GenerationParameters))] -public partial record GenerationParameters +public record GenerationParameters { public string? PositivePrompt { get; set; } public string? NegativePrompt { get; set; } @@ -128,16 +127,122 @@ internal static Dictionary ParseLine(string fields) { var dict = new Dictionary(); - // Values main contain commas or colons - foreach (var match in ParametersFieldsRegex().Matches(fields).Cast()) + var quoteStack = new Stack(); + // the Range for the key + Range? currentKeyRange = null; + // the start of the key or value + Index currentStart = 0; + + for (var i = 0; i < fields.Length; i++) { - if (!match.Success) - continue; + var c = fields[i]; + + switch (c) + { + case '"': + // if we are in a " quote, pop the stack + if (quoteStack.Count > 0 && quoteStack.Peek() == '"') + { + quoteStack.Pop(); + } + else + { + // start of a new quoted section + quoteStack.Push(c); + } + break; + + case '[': + case '{': + case '(': + case '<': + quoteStack.Push(c); + break; + + case ']': + case '}': + case ')': + case '>': + // check if we have a matching pair and pop the stack + var peek = quoteStack.Peek(); + if ( + (peek == '[' && c == ']') + || (peek == '{' && c == '}') + || (peek == '(' && c == ')') + || (peek == '<' && c == '>') + ) + { + quoteStack.Pop(); + } + break; + + case ':': + // : marks the end of the key + + // if we already have a key, ignore this colon as it is part of the value + // if we are not in a quote, we have a key + if (!currentKeyRange.HasValue && quoteStack.Count == 0) + { + currentKeyRange = new Range(currentStart, i); + currentStart = i + 1; + } + break; + case ',': + // , marks the end of a key-value pair + // if we are not in a quote, we have a value + if (quoteStack.Count == 0) + { + if (!currentKeyRange.HasValue) + { + // unexpected comma, reset and start from current position + currentStart = i + 1; + break; + } + + try + { + // extract the key and value + var key = fields[currentKeyRange!.Value].Trim(); + var value = UnquoteValue(fields[currentStart..i].Trim()); + + // check duplicates and prefer the first occurrence + if (!string.IsNullOrWhiteSpace(key) && !dict.ContainsKey(key)) + { + dict[key] = value; + } + } + catch (Exception) + { + // ignore individual key-value pair errors + // JsonNode.Parse() in UnquoteValue() can throw exceptions + } - var key = match.Groups[1].Value.Trim(); - var value = UnquoteValue(match.Groups[2].Value.Trim()); + currentKeyRange = null; + currentStart = i + 1; + } + break; + default: + break; + } // end of switch + } // end of for - dict.Add(key, value); + // if we have a key-value pair at the end of the string + if (currentKeyRange.HasValue) + { + try + { + var key = fields[currentKeyRange!.Value].Trim(); + var value = UnquoteValue(fields[currentStart..].Trim()); + + if (!string.IsNullOrWhiteSpace(key) && !dict.ContainsKey(key)) + { + dict[key] = value; + } + } + catch (Exception) + { + // ignore individual key-value pair errors + } } return dict; @@ -213,7 +318,4 @@ public static GenerationParameters GetSample() Sampler = "DPM++ 2M Karras" }; } - - [GeneratedRegex("""\s*([\w ]+):\s*("(?:\\.|[^\\"])+"|[^,]*)(?:,|$)""")] - private static partial Regex ParametersFieldsRegex(); } diff --git a/StabilityMatrix.Tests/Models/GenerationParametersTests.cs b/StabilityMatrix.Tests/Models/GenerationParametersTests.cs index d22caf8c0..815cd655c 100644 --- a/StabilityMatrix.Tests/Models/GenerationParametersTests.cs +++ b/StabilityMatrix.Tests/Models/GenerationParametersTests.cs @@ -51,20 +51,164 @@ public void TestParse_NoNegative() } [TestMethod] - public void TestParseLineFields() + // basic data + [DataRow( + """Steps: 30, Sampler: DPM++ 2M Karras, CFG scale: 7, Seed: 2216407431, Size: 640x896, Model hash: eb2h052f91, Model: anime_v1""", + 7, + "30", + "DPM++ 2M Karras", + "7", + "2216407431", + "640x896", + "eb2h052f91", + "anime_v1", + new string[] { "Steps", "Sampler", "CFG scale", "Seed", "Size", "Model hash", "Model" } + )] + // duplicated keys + [DataRow( + """Steps: 30, Sampler: DPM++ 2M Karras, CFG scale: 7, Seed: 2216407431, Size: 640x896, Model hash: eb2h052f91, Model: anime_v1, Steps: 40, Sampler: Whatever, CFG scale: 1, Seed: 1234567890, Size: 1024x1024, Model hash: 1234567890, Model: anime_v2""", + 7, + "30", + "DPM++ 2M Karras", + "7", + "2216407431", + "640x896", + "eb2h052f91", + "anime_v1", + new string[] { "Steps", "Sampler", "CFG scale", "Seed", "Size", "Model hash", "Model" } + )] + public void TestParseLineFields( + string line, + int totalFields, + string? expectedSteps, + string? expectedSampler, + string? expectedCfgScale, + string? expectedSeed, + string? expectedSize, + string? expectedModelHash, + string? expectedModel, + string[] expectedKeys + ) { - const string lastLine = - @"Steps: 30, Sampler: DPM++ 2M Karras, CFG scale: 7, Seed: 2216407431, Size: 640x896, Model hash: eb2h052f91, Model: anime_v1"; + var fields = GenerationParameters.ParseLine(line); - var fields = GenerationParameters.ParseLine(lastLine); + Assert.AreEqual(totalFields, fields.Count); + Assert.AreEqual(expectedSteps, fields["Steps"]); + Assert.AreEqual(expectedSampler, fields["Sampler"]); + Assert.AreEqual(expectedCfgScale, fields["CFG scale"]); + Assert.AreEqual(expectedSeed, fields["Seed"]); + Assert.AreEqual(expectedSize, fields["Size"]); + Assert.AreEqual(expectedModelHash, fields["Model hash"]); + Assert.AreEqual(expectedModel, fields["Model"]); + CollectionAssert.AreEqual(expectedKeys, fields.Keys); + } + + [TestMethod] + // empty line + [DataRow("", new string[] { })] + [DataRow(" ", new string[] { })] + // basic data + [DataRow( + "Steps: 30, Sampler: DPM++ 2M Karras, CFG scale: 7, Seed: 2216407431, Size: 640x896, Model hash: eb2h052f91, Model: anime_v1", + new string[] { "Steps", "Sampler", "CFG scale", "Seed", "Size", "Model hash", "Model" } + )] + // no spaces + [DataRow( + "Steps:30,Sampler:DPM++2MKarras,CFGscale:7,Seed:2216407431,Size:640x896,Modelhash:eb2h052f91,Model:anime_v1", + new string[] { "Steps", "Sampler", "CFGscale", "Seed", "Size", "Modelhash", "Model" } + )] + // extra commas + [DataRow( + "Steps: 30, Sampler: DPM++ 2M Karras, CFG scale: 7, Seed: 2216407431, Size: 640x896,,,,,, Model hash: eb2h052f91, Model: anime_v1,,,,,,,", + new string[] { "Steps", "Sampler", "CFG scale", "Seed", "Size", "Model hash", "Model" } + )] + // quoted string + [DataRow( + """Name: "John, Doe", Json: {"key:with:colon": "value, with, comma"}, It still: should work""", + new string[] { "Name", "Json", "It still" } + )] + // civitai + [DataRow( + """Steps: 8, Sampler: Euler, CFG scale: 1, Seed: 12346789098, Size: 832x1216, Clip skip: 2, Created Date: 2024-12-22T01:01:01.0222111Z, Civitai resources: [{"type":"checkpoint","modelVersionId":123456,"modelName":"Some model name here [Pony XL] which hopefully doesnt contains half pair of quotes and brackets","modelVersionName":"v2.0"},{"type":"lycoris","weight":0.7,"modelVersionId":11111111,"modelName":"some style","modelVersionName":"v1.0 pony"},{"type":"lora","weight":1,"modelVersionId":222222222,"modelName":"another name","modelVersionName":"v1.0"},{"type":"lora","modelVersionId":3333333,"modelName":"name for 33333333333","modelVersionName":"version name here"}], Civitai metadata: {"remixOfId":11111100000}""", + new string[] + { + "Steps", + "Sampler", + "CFG scale", + "Seed", + "Size", + "Clip skip", + "Created Date", + "Civitai resources", + "Civitai metadata" + } + )] + // github.com/nkchocoai/ComfyUI-SaveImageWithMetaData + [DataRow( + """Steps: 20, Sampler: DPM++ SDE Karras, CFG scale: 6.0, Seed: 1111111111111, Clip skip: 2, Size: 1024x1024, Model: the_main_model.safetensors, Model hash: ababababab, Lora_0 Model name: name_of_the_first_lora.safetensors, Lora_0 Model hash: ababababab, Lora_0 Strength model: -1.1, Lora_0 Strength clip: -1.1, Lora_1 Model name: name_of_the_second_lora.safetensors, Lora_1 Model hash: ababababab, Lora_1 Strength model: 1, Lora_1 Strength clip: 1, Lora_2 Model name: name_of_the_third_lora.safetensors, Lora_2 Model hash: ababababab, Lora_2 Strength model: 0.9, Lora_2 Strength clip: 0.9, Hashes: {"model": "ababababab", "lora:name_of_the_first_lora": "ababababab", "lora:name_of_the_second_lora": "ababababab", "lora:name_of_the_third_lora": "ababababab"}""", + new string[] + { + "Steps", + "Sampler", + "CFG scale", + "Seed", + "Clip skip", + "Size", + "Model", + "Model hash", + "Lora_0 Model name", + "Lora_0 Model hash", + "Lora_0 Strength model", + "Lora_0 Strength clip", + "Lora_1 Model name", + "Lora_1 Model hash", + "Lora_1 Strength model", + "Lora_1 Strength clip", + "Lora_2 Model name", + "Lora_2 Model hash", + "Lora_2 Strength model", + "Lora_2 Strength clip", + "Hashes" + } + )] + // asymmetrical bracket + [DataRow( + """Steps: 20, Missing closing bracket: {"name": "Someone did not close [this bracket"}, But: the parser, should: still return, the: fields before it""", + new string[] { "Steps", "Missing closing bracket" } + )] + public void TestParseLineEdgeCases(string line, string[] expectedKeys) + { + var fields = GenerationParameters.ParseLine(line); + + Assert.AreEqual(expectedKeys.Length, fields.Count); + CollectionAssert.AreEqual(expectedKeys, fields.Keys); + } + + [TestMethod] + public void TestParseLine() + { + var fields = GenerationParameters.ParseLine( + """Steps: 8, Sampler: Euler, CFG scale: 1, Seed: 12346789098, Size: 832x1216, Clip skip: 2, """ + + """Created Date: 2024-12-22T01:01:01.0222111Z, Civitai resources: [{"type":"checkpoint","modelVersionId":123456,"modelName":"Some model name here [Pony XL] which hopefully doesnt contains half pair of quotes and brackets","modelVersionName":"v2.0"},{"type":"lycoris","weight":0.7,"modelVersionId":11111111,"modelName":"some style","modelVersionName":"v1.0 pony"},{"type":"lora","weight":1,"modelVersionId":222222222,"modelName":"another name","modelVersionName":"v1.0"},{"type":"lora","modelVersionId":3333333,"modelName":"name for 33333333333","modelVersionName":"version name here"}], Civitai metadata: {"remixOfId":11111100000},""" + + """Hashes: {"model": "1234455678", "lora:aaaaaaa": "1234455678", "lora:bbbbbb": "1234455678", "lora:cccccccc": "1234455678"}""" + ); - Assert.AreEqual(7, fields.Count); - Assert.AreEqual("30", fields["Steps"]); - Assert.AreEqual("DPM++ 2M Karras", fields["Sampler"]); - Assert.AreEqual("7", fields["CFG scale"]); - Assert.AreEqual("2216407431", fields["Seed"]); - Assert.AreEqual("640x896", fields["Size"]); - Assert.AreEqual("eb2h052f91", fields["Model hash"]); - Assert.AreEqual("anime_v1", fields["Model"]); + Assert.AreEqual(10, fields.Count); + Assert.AreEqual("8", fields["Steps"]); + Assert.AreEqual("Euler", fields["Sampler"]); + Assert.AreEqual("1", fields["CFG scale"]); + Assert.AreEqual("12346789098", fields["Seed"]); + Assert.AreEqual("832x1216", fields["Size"]); + Assert.AreEqual("2", fields["Clip skip"]); + Assert.AreEqual("2024-12-22T01:01:01.0222111Z", fields["Created Date"]); + Assert.AreEqual( + """[{"type":"checkpoint","modelVersionId":123456,"modelName":"Some model name here [Pony XL] which hopefully doesnt contains half pair of quotes and brackets","modelVersionName":"v2.0"},{"type":"lycoris","weight":0.7,"modelVersionId":11111111,"modelName":"some style","modelVersionName":"v1.0 pony"},{"type":"lora","weight":1,"modelVersionId":222222222,"modelName":"another name","modelVersionName":"v1.0"},{"type":"lora","modelVersionId":3333333,"modelName":"name for 33333333333","modelVersionName":"version name here"}]""", + fields["Civitai resources"] + ); + Assert.AreEqual("""{"remixOfId":11111100000}""", fields["Civitai metadata"]); + Assert.AreEqual( + """{"model": "1234455678", "lora:aaaaaaa": "1234455678", "lora:bbbbbb": "1234455678", "lora:cccccccc": "1234455678"}""", + fields["Hashes"] + ); } }