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 3c6e03f
Show file tree
Hide file tree
Showing 2 changed files with 284 additions and 37 deletions.
148 changes: 123 additions & 25 deletions StabilityMatrix.Core/Models/GenerationParameters.cs
View file
Original file line number Diff line number Diff line change
@@ -1,14 +1,12 @@
using System.ComponentModel.DataAnnotations;
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 @@ -124,36 +122,139 @@ public static GenerationParameters Parse(string text)
/// fields are separated by commas and key-value pairs are separated by colons.
/// i.e. "key1: value1, key2: value2"
/// </summary>
internal static Dictionary<string, string> ParseLine(string fields)
internal static Dictionary<string, string> ParseLine(string line)
{
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 < line.Length; i++)
{
if (!match.Success)
continue;
var c = line[i];

var key = match.Groups[1].Value.Trim();
var value = UnquoteValue(match.Groups[2].Value.Trim());
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;

dict.Add(key, value);
}
case '[':
case '{':
case '(':
case '<':
quoteStack.Push(c);
break;

return dict;
}
case ']':
if (quoteStack.Count > 0 && quoteStack.Peek() == '[')
{
quoteStack.Pop();
}
break;
case '}':
if (quoteStack.Count > 0 && quoteStack.Peek() == '{')
{
quoteStack.Pop();
}
break;
case ')':
if (quoteStack.Count > 0 && quoteStack.Peek() == '(')
{
quoteStack.Pop();
}
break;
case '>':
if (quoteStack.Count > 0 && quoteStack.Peek() == '<')
{
quoteStack.Pop();
}
break;

/// <summary>
/// Unquotes a quoted value field if required
/// </summary>
private static string UnquoteValue(string quotedField)
{
if (!(quotedField.StartsWith('"') && quotedField.EndsWith('"')))
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)
{
break;
}

if (!currentKeyRange.HasValue)
{
// unexpected comma, reset and start from current position
currentStart = i + 1;
break;
}

try
{
// extract the key and value
var key = new string(line.AsSpan()[currentKeyRange!.Value].Trim());
var value = new string(line.AsSpan()[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
}

currentKeyRange = null;
currentStart = i + 1;
break;
default:
break;
} // end of switch
} // end of for

// if we have a key-value pair at the end of the string
if (currentKeyRange.HasValue)
{
return quotedField;
try
{
var key = new string(line.AsSpan()[currentKeyRange!.Value].Trim());
var value = new string(line.AsSpan()[currentStart..].Trim());

if (!string.IsNullOrWhiteSpace(key) && !dict.ContainsKey(key))
{
dict[key] = value;
}
}
catch (Exception)
{
// ignore individual key-value pair errors
}
}

return JsonNode.Parse(quotedField)?.GetValue<string>() ?? "";
return dict;
}

/// <summary>
Expand Down Expand Up @@ -213,7 +314,4 @@ public static GenerationParameters GetSample()
Sampler = "DPM++ 2M Karras"
};
}

[GeneratedRegex("""\s*([\w ]+):\s*("(?:\\.|[^\\"])+"|[^,]*)(?:,|$)""")]
private static partial Regex ParametersFieldsRegex();
}
173 changes: 161 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,169 @@ 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" }
)]
// extra ending brackets
[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 3c6e03f

Please sign in to comment.