Skip to content

Commit 3c6e03f

Browse files
GenteureGenteure
committed
Improve GenerationParameters fields parsing
1 parent d34603c commit 3c6e03f

File tree

2 files changed

+284
-37
lines changed

2 files changed

+284
-37
lines changed

‎StabilityMatrix.Core/Models/GenerationParameters.cs‎

Lines changed: 123 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,12 @@
11
using System.ComponentModel.DataAnnotations;
22
using System.Diagnostics.CodeAnalysis;
3-
using System.Text.Json.Nodes;
43
using System.Text.Json.Serialization;
5-
using System.Text.RegularExpressions;
64
using StabilityMatrix.Core.Models.Api.Comfy;
75

86
namespace StabilityMatrix.Core.Models;
97

108
[JsonSerializable(typeof(GenerationParameters))]
11-
public partial record GenerationParameters
9+
public record GenerationParameters
1210
{
1311
public string? PositivePrompt { get; set; }
1412
public string? NegativePrompt { get; set; }
@@ -124,36 +122,139 @@ public static GenerationParameters Parse(string text)
124122
/// fields are separated by commas and key-value pairs are separated by colons.
125123
/// i.e. "key1: value1, key2: value2"
126124
/// </summary>
127-
internal static Dictionary<string, string> ParseLine(string fields)
125+
internal static Dictionary<string, string> ParseLine(string line)
128126
{
129127
var dict = new Dictionary<string, string>();
130128

131-
// Values main contain commas or colons
132-
foreach (var match in ParametersFieldsRegex().Matches(fields).Cast<Match>())
129+
var quoteStack = new Stack<char>();
130+
// the Range for the key
131+
Range? currentKeyRange = null;
132+
// the start of the key or value
133+
Index currentStart = 0;
134+
135+
for (var i = 0; i < line.Length; i++)
133136
{
134-
if (!match.Success)
135-
continue;
137+
var c = line[i];
136138

137-
var key = match.Groups[1].Value.Trim();
138-
var value = UnquoteValue(match.Groups[2].Value.Trim());
139+
switch (c)
140+
{
141+
case '"':
142+
// if we are in a " quote, pop the stack
143+
if (quoteStack.Count > 0 && quoteStack.Peek() == '"')
144+
{
145+
quoteStack.Pop();
146+
}
147+
else
148+
{
149+
// start of a new quoted section
150+
quoteStack.Push(c);
151+
}
152+
break;
139153

140-
dict.Add(key, value);
141-
}
154+
case '[':
155+
case '{':
156+
case '(':
157+
case '<':
158+
quoteStack.Push(c);
159+
break;
142160

143-
return dict;
144-
}
161+
case ']':
162+
if (quoteStack.Count > 0 && quoteStack.Peek() == '[')
163+
{
164+
quoteStack.Pop();
165+
}
166+
break;
167+
case '}':
168+
if (quoteStack.Count > 0 && quoteStack.Peek() == '{')
169+
{
170+
quoteStack.Pop();
171+
}
172+
break;
173+
case ')':
174+
if (quoteStack.Count > 0 && quoteStack.Peek() == '(')
175+
{
176+
quoteStack.Pop();
177+
}
178+
break;
179+
case '>':
180+
if (quoteStack.Count > 0 && quoteStack.Peek() == '<')
181+
{
182+
quoteStack.Pop();
183+
}
184+
break;
145185

146-
/// <summary>
147-
/// Unquotes a quoted value field if required
148-
/// </summary>
149-
private static string UnquoteValue(string quotedField)
150-
{
151-
if (!(quotedField.StartsWith('"') && quotedField.EndsWith('"')))
186+
case ':':
187+
// : marks the end of the key
188+
189+
// if we already have a key, ignore this colon as it is part of the value
190+
// if we are not in a quote, we have a key
191+
if (!currentKeyRange.HasValue && quoteStack.Count == 0)
192+
{
193+
currentKeyRange = new Range(currentStart, i);
194+
currentStart = i + 1;
195+
}
196+
break;
197+
198+
case ',':
199+
// , marks the end of a key-value pair
200+
// if we are not in a quote, we have a value
201+
if (quoteStack.Count != 0)
202+
{
203+
break;
204+
}
205+
206+
if (!currentKeyRange.HasValue)
207+
{
208+
// unexpected comma, reset and start from current position
209+
currentStart = i + 1;
210+
break;
211+
}
212+
213+
try
214+
{
215+
// extract the key and value
216+
var key = new string(line.AsSpan()[currentKeyRange!.Value].Trim());
217+
var value = new string(line.AsSpan()[currentStart..i].Trim());
218+
219+
// check duplicates and prefer the first occurrence
220+
if (!string.IsNullOrWhiteSpace(key) && !dict.ContainsKey(key))
221+
{
222+
dict[key] = value;
223+
}
224+
}
225+
catch (Exception)
226+
{
227+
// ignore individual key-value pair errors
228+
}
229+
230+
currentKeyRange = null;
231+
currentStart = i + 1;
232+
break;
233+
default:
234+
break;
235+
} // end of switch
236+
} // end of for
237+
238+
// if we have a key-value pair at the end of the string
239+
if (currentKeyRange.HasValue)
152240
{
153-
return quotedField;
241+
try
242+
{
243+
var key = new string(line.AsSpan()[currentKeyRange!.Value].Trim());
244+
var value = new string(line.AsSpan()[currentStart..].Trim());
245+
246+
if (!string.IsNullOrWhiteSpace(key) && !dict.ContainsKey(key))
247+
{
248+
dict[key] = value;
249+
}
250+
}
251+
catch (Exception)
252+
{
253+
// ignore individual key-value pair errors
254+
}
154255
}
155256

156-
return JsonNode.Parse(quotedField)?.GetValue<string>() ?? "";
257+
return dict;
157258
}
158259

159260
/// <summary>
@@ -213,7 +314,4 @@ public static GenerationParameters GetSample()
213314
Sampler = "DPM++ 2M Karras"
214315
};
215316
}
216-
217-
[GeneratedRegex("""\s*([\w ]+):\s*("(?:\\.|[^\\"])+"|[^,]*)(?:,|$)""")]
218-
private static partial Regex ParametersFieldsRegex();
219317
}

‎StabilityMatrix.Tests/Models/GenerationParametersTests.cs‎

Lines changed: 161 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -51,20 +51,169 @@ public void TestParse_NoNegative()
5151
}
5252

5353
[TestMethod]
54-
public void TestParseLineFields()
54+
// basic data
55+
[DataRow(
56+
"""Steps: 30, Sampler: DPM++ 2M Karras, CFG scale: 7, Seed: 2216407431, Size: 640x896, Model hash: eb2h052f91, Model: anime_v1""",
57+
7,
58+
"30",
59+
"DPM++ 2M Karras",
60+
"7",
61+
"2216407431",
62+
"640x896",
63+
"eb2h052f91",
64+
"anime_v1",
65+
new string[] { "Steps", "Sampler", "CFG scale", "Seed", "Size", "Model hash", "Model" }
66+
)]
67+
// duplicated keys
68+
[DataRow(
69+
"""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""",
70+
7,
71+
"30",
72+
"DPM++ 2M Karras",
73+
"7",
74+
"2216407431",
75+
"640x896",
76+
"eb2h052f91",
77+
"anime_v1",
78+
new string[] { "Steps", "Sampler", "CFG scale", "Seed", "Size", "Model hash", "Model" }
79+
)]
80+
public void TestParseLineFields(
81+
string line,
82+
int totalFields,
83+
string? expectedSteps,
84+
string? expectedSampler,
85+
string? expectedCfgScale,
86+
string? expectedSeed,
87+
string? expectedSize,
88+
string? expectedModelHash,
89+
string? expectedModel,
90+
string[] expectedKeys
91+
)
5592
{
56-
const string lastLine =
57-
@"Steps: 30, Sampler: DPM++ 2M Karras, CFG scale: 7, Seed: 2216407431, Size: 640x896, Model hash: eb2h052f91, Model: anime_v1";
93+
var fields = GenerationParameters.ParseLine(line);
5894

59-
var fields = GenerationParameters.ParseLine(lastLine);
95+
Assert.AreEqual(totalFields, fields.Count);
96+
Assert.AreEqual(expectedSteps, fields["Steps"]);
97+
Assert.AreEqual(expectedSampler, fields["Sampler"]);
98+
Assert.AreEqual(expectedCfgScale, fields["CFG scale"]);
99+
Assert.AreEqual(expectedSeed, fields["Seed"]);
100+
Assert.AreEqual(expectedSize, fields["Size"]);
101+
Assert.AreEqual(expectedModelHash, fields["Model hash"]);
102+
Assert.AreEqual(expectedModel, fields["Model"]);
103+
CollectionAssert.AreEqual(expectedKeys, fields.Keys);
104+
}
105+
106+
[TestMethod]
107+
// empty line
108+
[DataRow("", new string[] { })]
109+
[DataRow(" ", new string[] { })]
110+
// basic data
111+
[DataRow(
112+
"Steps: 30, Sampler: DPM++ 2M Karras, CFG scale: 7, Seed: 2216407431, Size: 640x896, Model hash: eb2h052f91, Model: anime_v1",
113+
new string[] { "Steps", "Sampler", "CFG scale", "Seed", "Size", "Model hash", "Model" }
114+
)]
115+
// no spaces
116+
[DataRow(
117+
"Steps:30,Sampler:DPM++2MKarras,CFGscale:7,Seed:2216407431,Size:640x896,Modelhash:eb2h052f91,Model:anime_v1",
118+
new string[] { "Steps", "Sampler", "CFGscale", "Seed", "Size", "Modelhash", "Model" }
119+
)]
120+
// extra commas
121+
[DataRow(
122+
"Steps: 30, Sampler: DPM++ 2M Karras, CFG scale: 7, Seed: 2216407431, Size: 640x896,,,,,, Model hash: eb2h052f91, Model: anime_v1,,,,,,,",
123+
new string[] { "Steps", "Sampler", "CFG scale", "Seed", "Size", "Model hash", "Model" }
124+
)]
125+
// quoted string
126+
[DataRow(
127+
"""Name: "John, Doe", Json: {"key:with:colon": "value, with, comma"}, It still: should work""",
128+
new string[] { "Name", "Json", "It still" }
129+
)]
130+
// extra ending brackets
131+
[DataRow(
132+
"""Name: "John, Doe", Json: {"key:with:colon": "value, with, comma"}}}}}}}})))>>, It still: should work""",
133+
new string[] { "Name", "Json", "It still" }
134+
)]
135+
// civitai
136+
[DataRow(
137+
"""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}""",
138+
new string[]
139+
{
140+
"Steps",
141+
"Sampler",
142+
"CFG scale",
143+
"Seed",
144+
"Size",
145+
"Clip skip",
146+
"Created Date",
147+
"Civitai resources",
148+
"Civitai metadata"
149+
}
150+
)]
151+
// github.com/nkchocoai/ComfyUI-SaveImageWithMetaData
152+
[DataRow(
153+
"""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"}""",
154+
new string[]
155+
{
156+
"Steps",
157+
"Sampler",
158+
"CFG scale",
159+
"Seed",
160+
"Clip skip",
161+
"Size",
162+
"Model",
163+
"Model hash",
164+
"Lora_0 Model name",
165+
"Lora_0 Model hash",
166+
"Lora_0 Strength model",
167+
"Lora_0 Strength clip",
168+
"Lora_1 Model name",
169+
"Lora_1 Model hash",
170+
"Lora_1 Strength model",
171+
"Lora_1 Strength clip",
172+
"Lora_2 Model name",
173+
"Lora_2 Model hash",
174+
"Lora_2 Strength model",
175+
"Lora_2 Strength clip",
176+
"Hashes"
177+
}
178+
)]
179+
// asymmetrical bracket
180+
[DataRow(
181+
"""Steps: 20, Missing closing bracket: {"name": "Someone did not close [this bracket"}, But: the parser, should: still return, the: fields before it""",
182+
new string[] { "Steps", "Missing closing bracket" }
183+
)]
184+
public void TestParseLineEdgeCases(string line, string[] expectedKeys)
185+
{
186+
var fields = GenerationParameters.ParseLine(line);
187+
188+
Assert.AreEqual(expectedKeys.Length, fields.Count);
189+
CollectionAssert.AreEqual(expectedKeys, fields.Keys);
190+
}
191+
192+
[TestMethod]
193+
public void TestParseLine()
194+
{
195+
var fields = GenerationParameters.ParseLine(
196+
"""Steps: 8, Sampler: Euler, CFG scale: 1, Seed: 12346789098, Size: 832x1216, Clip skip: 2, """
197+
+ """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},"""
198+
+ """Hashes: {"model": "1234455678", "lora:aaaaaaa": "1234455678", "lora:bbbbbb": "1234455678", "lora:cccccccc": "1234455678"}"""
199+
);
60200

61-
Assert.AreEqual(7, fields.Count);
62-
Assert.AreEqual("30", fields["Steps"]);
63-
Assert.AreEqual("DPM++ 2M Karras", fields["Sampler"]);
64-
Assert.AreEqual("7", fields["CFG scale"]);
65-
Assert.AreEqual("2216407431", fields["Seed"]);
66-
Assert.AreEqual("640x896", fields["Size"]);
67-
Assert.AreEqual("eb2h052f91", fields["Model hash"]);
68-
Assert.AreEqual("anime_v1", fields["Model"]);
201+
Assert.AreEqual(10, fields.Count);
202+
Assert.AreEqual("8", fields["Steps"]);
203+
Assert.AreEqual("Euler", fields["Sampler"]);
204+
Assert.AreEqual("1", fields["CFG scale"]);
205+
Assert.AreEqual("12346789098", fields["Seed"]);
206+
Assert.AreEqual("832x1216", fields["Size"]);
207+
Assert.AreEqual("2", fields["Clip skip"]);
208+
Assert.AreEqual("2024-12-22T01:01:01.0222111Z", fields["Created Date"]);
209+
Assert.AreEqual(
210+
"""[{"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"}]""",
211+
fields["Civitai resources"]
212+
);
213+
Assert.AreEqual("""{"remixOfId":11111100000}""", fields["Civitai metadata"]);
214+
Assert.AreEqual(
215+
"""{"model": "1234455678", "lora:aaaaaaa": "1234455678", "lora:bbbbbb": "1234455678", "lora:cccccccc": "1234455678"}""",
216+
fields["Hashes"]
217+
);
69218
}
70219
}

0 commit comments

Comments
 (0)