Skip to content

Commit

Permalink
Improve GenerationParameters fields parsing
Browse files Browse the repository at this point in the history
  • Loading branch information
@Genteure
Genteure committed Dec 26, 2024
1 parent d34603c commit 4897b21
Show file tree
Hide file tree
Showing 2 changed files with 270 additions and 24 deletions.
126 changes: 114 additions & 12 deletions StabilityMatrix.Core/Models/GenerationParameters.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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; }
Expand Down Expand Up @@ -128,16 +127,122 @@ internal static Dictionary<string, string> ParseLine(string fields)
{
var dict = new Dictionary<string, string>();

// Values main contain commas or colons
foreach (var match in ParametersFieldsRegex().Matches(fields).Cast<Match>())
var quoteStack = new Stack<char>();
// 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;
Expand Down Expand Up @@ -213,7 +318,4 @@ public static GenerationParameters GetSample()
Sampler = "DPM++ 2M Karras"
};
}

[GeneratedRegex("""\s*([\w ]+):\s*("(?:\\.|[^\\"])+"|[^,]*)(?:,|$)""")]
private static partial Regex ParametersFieldsRegex();
}
168 changes: 156 additions & 12 deletions StabilityMatrix.Tests/Models/GenerationParametersTests.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
);
}
}

0 comments on commit 4897b21

Please sign in to comment.