Skip to content

Commit 2fa64bf

Browse files
mohnjilesmohnjiles
committed
Add connected metadata model organizer
1 parent 8ba6510 commit 2fa64bf

16 files changed

Lines changed: 2342 additions & 134 deletions

‎StabilityMatrix.Avalonia/Models/CheckpointOrganizer/ModelOrganizationModels.cs‎

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
1+
using System.IO;
2+
using StabilityMatrix.Core.Models.Database;
3+
4+
namespace StabilityMatrix.Avalonia.Models.CheckpointOrganizer;
5+
6+
public enum ModelOrganizationPreviewStatus
7+
{
8+
Ready,
9+
Conflict,
10+
Skipped,
11+
Unchanged,
12+
}
13+
14+
public enum ModelOrganizationMetadataAction
15+
{
16+
None,
17+
ScanMissing,
18+
UpdateExisting,
19+
}
20+
21+
public sealed record ModelOrganizationFileMove
22+
{
23+
public required string SourcePath { get; init; }
24+
public required string TargetPath { get; init; }
25+
}
26+
27+
public sealed record ModelOrganizationPreviewItem
28+
{
29+
public required LocalModelFile Model { get; init; }
30+
public required string SourcePath { get; init; }
31+
public string? TargetPath { get; init; }
32+
public required ModelOrganizationPreviewStatus Status { get; init; }
33+
public string? Reason { get; init; }
34+
public IReadOnlyList<ModelOrganizationFileMove> FileMoves { get; init; } = [];
35+
36+
public bool CanApply => Status == ModelOrganizationPreviewStatus.Ready;
37+
38+
public bool IsUnchanged => Status == ModelOrganizationPreviewStatus.Unchanged;
39+
40+
public string SourceFileName => Path.GetFileName(SourcePath);
41+
public string? TargetFileName => TargetPath is not null ? Path.GetFileName(TargetPath) : null;
42+
43+
public int SortOrder =>
44+
Status switch
45+
{
46+
ModelOrganizationPreviewStatus.Ready => 0,
47+
ModelOrganizationPreviewStatus.Conflict => 1,
48+
ModelOrganizationPreviewStatus.Skipped => 2,
49+
ModelOrganizationPreviewStatus.Unchanged => 3,
50+
_ => 4,
51+
};
52+
53+
public string StatusText =>
54+
Status switch
55+
{
56+
ModelOrganizationPreviewStatus.Ready => "Ready",
57+
ModelOrganizationPreviewStatus.Conflict => "Conflict",
58+
ModelOrganizationPreviewStatus.Skipped => "Skipped",
59+
ModelOrganizationPreviewStatus.Unchanged => "Unchanged",
60+
_ => Status.ToString(),
61+
};
62+
}
63+
64+
public sealed record ModelOrganizationPlan
65+
{
66+
public required string Template { get; init; }
67+
public required string ScopePath { get; init; }
68+
public required bool IncludeNested { get; init; }
69+
public string? ValidationError { get; init; }
70+
public IReadOnlyList<ModelOrganizationPreviewItem> Items { get; init; } = [];
71+
72+
public int ReadyCount => Items.Count(item => item.Status == ModelOrganizationPreviewStatus.Ready);
73+
74+
public int ConflictCount => Items.Count(item => item.Status == ModelOrganizationPreviewStatus.Conflict);
75+
76+
public int SkippedCount =>
77+
Items.Count(item =>
78+
item.Status is ModelOrganizationPreviewStatus.Skipped or ModelOrganizationPreviewStatus.Unchanged
79+
);
80+
}
81+
82+
public sealed record ModelOrganizationApplyResult
83+
{
84+
public required int MovedCount { get; init; }
85+
public required int SkippedCount { get; init; }
86+
public required int ConflictCount { get; init; }
87+
public IReadOnlyList<string> Errors { get; init; } = [];
88+
}

‎StabilityMatrix.Avalonia/Models/Inference/FileNameFormat.cs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,4 +81,5 @@ public static bool TryParse(
8181

8282
public const string DefaultTemplate = "{date}_{time}-{model_name}-{seed}";
8383
public const string DefaultModelBrowserTemplate = "{file_name}";
84+
public const string DefaultOrganizationTemplate = "{file_name}";
8485
}

‎StabilityMatrix.Avalonia/Models/Inference/FileNameFormatProvider.cs‎

Lines changed: 151 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
using StabilityMatrix.Core.Extensions;
1010
using StabilityMatrix.Core.Models;
1111
using StabilityMatrix.Core.Models.Api;
12+
using StabilityMatrix.Core.Models.Database;
1213
using StabilityMatrix.Core.Models.Inference;
1314

1415
namespace StabilityMatrix.Avalonia.Models.Inference;
@@ -24,6 +25,25 @@ public partial class FileNameFormatProvider
2425
public CivitModel? CivitModel { get; init; }
2526
public CivitModelVersion? CivitModelVersion { get; init; }
2627
public CivitFile? CivitFile { get; init; }
28+
public LocalModelFile? LocalModelFile { get; init; }
29+
30+
public static ISet<string> LocalOrganizationVariables { get; } =
31+
new HashSet<string>(
32+
[
33+
"date",
34+
"time",
35+
"author",
36+
"base_model",
37+
"file_name",
38+
"file_id",
39+
"model_id",
40+
"model_name",
41+
"model_version_id",
42+
"model_version_name",
43+
"model_type",
44+
],
45+
StringComparer.Ordinal
46+
);
2747

2848
private Dictionary<string, Func<string?>>? _substitutions;
2949

@@ -35,7 +55,11 @@ public partial class FileNameFormatProvider
3555
{ "negative_prompt", () => GenerationParameters?.NegativePrompt },
3656
{
3757
"model_name",
38-
() => Path.GetFileNameWithoutExtension(GenerationParameters?.ModelName) ?? CivitModel?.Name
58+
() =>
59+
Path.GetFileNameWithoutExtension(GenerationParameters?.ModelName)
60+
?? CivitModel?.Name
61+
?? LocalModelFile?.ConnectedModelInfo?.ModelName
62+
?? LocalModelFile?.FileNameWithoutExtension
3963
},
4064
{ "model_hash", () => GenerationParameters?.ModelHash },
4165
{ "sampler", () => GenerationParameters?.Sampler },
@@ -47,14 +71,45 @@ public partial class FileNameFormatProvider
4771
{ "project_name", () => ProjectName },
4872
{ "date", () => DateTime.Now.ToString("yyyy-MM-dd") },
4973
{ "time", () => DateTime.Now.ToString("HH-mm-ss") },
50-
{ "author", () => CivitModel?.Creator?.Username },
51-
{ "base_model", () => CivitModelVersion?.BaseModel },
52-
{ "file_name", () => Path.GetFileNameWithoutExtension(CivitFile?.Name) },
53-
{ "file_id", () => CivitFile?.Id.ToString() },
54-
{ "model_id", () => CivitModel?.Id.ToString() },
55-
{ "model_version_id", () => CivitModelVersion?.Id.ToString() },
56-
{ "model_version_name", () => CivitModelVersion?.Name },
57-
{ "model_type", () => CivitModel?.Type.ToString() },
74+
{
75+
"author",
76+
() => CivitModel?.Creator?.Username ?? LocalModelFile?.ConnectedModelInfo?.AuthorUsername
77+
},
78+
{
79+
"base_model",
80+
() => CivitModelVersion?.BaseModel ?? LocalModelFile?.ConnectedModelInfo?.BaseModel
81+
},
82+
{
83+
"file_name",
84+
() =>
85+
Path.GetFileNameWithoutExtension(CivitFile?.Name)
86+
?? Path.GetFileNameWithoutExtension(LocalModelFile?.ConnectedModelInfo?.RemoteFileName)
87+
?? LocalModelFile?.FileNameWithoutExtension
88+
},
89+
{
90+
"file_id",
91+
() => CivitFile?.Id.ToString() ?? LocalModelFile?.ConnectedModelInfo?.RemoteFileId?.ToString()
92+
},
93+
{
94+
"model_id",
95+
() => CivitModel?.Id.ToString() ?? LocalModelFile?.ConnectedModelInfo?.ModelId?.ToString()
96+
},
97+
{
98+
"model_version_id",
99+
() =>
100+
CivitModelVersion?.Id.ToString()
101+
?? LocalModelFile?.ConnectedModelInfo?.VersionId?.ToString()
102+
},
103+
{
104+
"model_version_name",
105+
() => CivitModelVersion?.Name ?? LocalModelFile?.ConnectedModelInfo?.VersionName
106+
},
107+
{
108+
"model_type",
109+
() =>
110+
CivitModel?.Type.ToString()
111+
?? (LocalModelFile?.ConnectedModelInfo is { } cmInfo ? cmInfo.ModelType.ToString() : null)
112+
},
58113
};
59114

60115
/// <summary>
@@ -65,15 +120,13 @@ public partial class FileNameFormatProvider
65120
[Pure]
66121
public ValidationResult Validate(string format)
67122
{
68-
var regex = BracketRegex();
69-
var matches = regex.Matches(format);
70-
var variables = matches.Select(m => m.Groups[1].Value);
123+
var variables = GetVariableTexts(format);
71124

72125
foreach (var variableText in variables)
73126
{
74127
try
75128
{
76-
var (variable, _) = ExtractVariableAndSlice(variableText);
129+
var variable = GetVariableName(variableText);
77130

78131
if (!Substitutions.ContainsKey(variable))
79132
{
@@ -89,6 +142,47 @@ public ValidationResult Validate(string format)
89142
return ValidationResult.Success!;
90143
}
91144

145+
public IEnumerable<string> GetVariableTexts(string template)
146+
{
147+
return BracketRegex().Matches(template).Select(m => m.Groups[1].Value);
148+
}
149+
150+
public string GetVariableName(string variableText)
151+
{
152+
var (variable, _) = ExtractVariableAndSlice(variableText);
153+
return variable;
154+
}
155+
156+
public bool TryResolveVariable(string variableText, out string? value, out string? error)
157+
{
158+
try
159+
{
160+
var (variable, slice) = ExtractVariableAndSlice(variableText);
161+
if (!Substitutions.TryGetValue(variable, out var substitution))
162+
{
163+
value = null;
164+
error = $"Unknown variable '{variable}'";
165+
return false;
166+
}
167+
168+
value = ApplySlice(substitution(), slice);
169+
if (value is null)
170+
{
171+
error = $"Variable '{variable}' is not available";
172+
return false;
173+
}
174+
175+
error = null;
176+
return true;
177+
}
178+
catch (Exception e)
179+
{
180+
value = null;
181+
error = $"Invalid variable '{variableText}': {e.Message}";
182+
return false;
183+
}
184+
}
185+
92186
public IEnumerable<FileNameFormatPart> GetParts(string template)
93187
{
94188
var regex = BracketRegex();
@@ -114,36 +208,7 @@ public IEnumerable<FileNameFormatPart> GetParts(string template)
114208
var (variable, slice) = ExtractVariableAndSlice(result.Groups[1].Value);
115209
var substitution = Substitutions[variable];
116210

117-
// Slice string if necessary
118-
if (slice is not null)
119-
{
120-
parts.Add(
121-
(FileNameFormatPart)(
122-
() =>
123-
{
124-
var value = substitution();
125-
if (value is null)
126-
return null;
127-
128-
if (slice.End is null)
129-
{
130-
value = value[(slice.Start ?? 0)..];
131-
}
132-
else
133-
{
134-
var length = Math.Min(value.Length, slice.End.Value) - (slice.Start ?? 0);
135-
value = value.Substring(slice.Start ?? 0, length);
136-
}
137-
138-
return value;
139-
}
140-
)
141-
);
142-
}
143-
else
144-
{
145-
parts.Add(substitution);
146-
}
211+
parts.Add((FileNameFormatPart)(() => ApplySlice(substitution(), slice)));
147212

148213
currentIndex += result.Length;
149214
}
@@ -198,6 +263,35 @@ public static FileNameFormatProvider GetSampleForModelBrowser()
198263
};
199264
}
200265

266+
public static FileNameFormatProvider GetSampleForOrganization()
267+
{
268+
return new FileNameFormatProvider
269+
{
270+
LocalModelFile = new LocalModelFile
271+
{
272+
RelativePath = "StableDiffusion/sample_file.safetensors",
273+
SharedFolderType = SharedFolderType.StableDiffusion,
274+
ConnectedModelInfo = new ConnectedModelInfo
275+
{
276+
ModelId = 1234,
277+
ModelName = "Sample Model",
278+
ModelDescription = string.Empty,
279+
Nsfw = false,
280+
Tags = [],
281+
ModelType = CivitModelType.Checkpoint,
282+
VersionId = 5678,
283+
VersionName = "v1.0",
284+
AuthorUsername = "SampleUser",
285+
BaseModel = "Illustrious",
286+
RemoteFileName = "sample_file.safetensors",
287+
RemoteFileId = 910,
288+
Hashes = new CivitFileHashes(),
289+
Source = ConnectedModelSource.Civitai,
290+
},
291+
},
292+
};
293+
}
294+
201295
/// <summary>
202296
/// Extract variable and index from a combined string
203297
/// </summary>
@@ -224,6 +318,20 @@ private static (string Variable, Slice? Slice) ExtractVariableAndSlice(string co
224318
return (variable, slice);
225319
}
226320

321+
private static string? ApplySlice(string? value, Slice? slice)
322+
{
323+
if (value is null || slice is null)
324+
return value;
325+
326+
if (slice.End is null)
327+
{
328+
return value[(slice.Start ?? 0)..];
329+
}
330+
331+
var length = Math.Min(value.Length, slice.End.Value) - (slice.Start ?? 0);
332+
return value.Substring(slice.Start ?? 0, length);
333+
}
334+
227335
/// <summary>
228336
/// Regex for matching contents within a curly brace.
229337
/// </summary>

0 commit comments

Comments
 (0)