Skip to content

Commit 3cd3552

Browse files
committed
Make baseModel folder a user setting
By making this a default to off user setting, we avoid the problem of forcing users to use a base model folder when downloading models using the model downloader tool. Check path with Utilities.StrictFilenameClean() for safety
1 parent 3bb0363 commit 3cd3552

File tree

2 files changed

+17
-10
lines changed

2 files changed

+17
-10
lines changed

src/Core/Settings.cs

+3
Original file line numberDiff line numberDiff line change
@@ -273,6 +273,9 @@ public class FileFormatData : AutoConfiguration
273273
[ConfigComment("If true, folders will be discarded from starred image paths.")]
274274
public bool StarNoFolders = false;
275275

276+
[ConfigComment("Whether to automatically use the base model type as part of the model path when downloading using the 'Model Download' tool")]
277+
public bool GroupDownloadedModelsByBaseType = false;
278+
276279
public class ThemesImpl : SettingsOptionsAttribute.AbstractImpl
277280
{
278281
public override string[] GetOptions => [.. Program.Web.RegisteredThemes.Keys];

src/WebAPI/ModelsAPI.cs

+14-10
Original file line numberDiff line numberDiff line change
@@ -511,16 +511,17 @@ public static async Task<JObject> DoModelDownloadWS(Session session, WebSocket w
511511
JObject metadataObj = JObject.Parse(metadata);
512512
baseModel = metadataObj["modelspec.baseModel"]?.ToString() ?? "";
513513
}
514-
string outPath;
515-
if (!string.IsNullOrWhiteSpace(baseModel))
514+
string modelOutPath;
515+
if (!string.IsNullOrWhiteSpace(baseModel) && session.User.Settings.GroupDownloadedModelsByBaseType)
516516
{
517-
outPath = $"{handler.FolderPaths[0]}/{baseModel}/{name}.safetensors";
517+
modelOutPath = $"{handler.FolderPaths[0]}/{baseModel}/{name}.safetensors";
518518
}
519519
else
520520
{
521-
outPath = $"{handler.FolderPaths[0]}/{name}.safetensors";
521+
modelOutPath = $"{handler.FolderPaths[0]}/{name}.safetensors";
522522
}
523-
if (File.Exists(outPath))
523+
modelOutPath = Utilities.StrictFilenameClean(outPath);
524+
if (File.Exists(modelOutPath))
524525
{
525526
await ws.SendJson(new JObject() { ["error"] = "Model at that save path already exists." }, API.WebsocketTimeout);
526527
return null;
@@ -530,7 +531,7 @@ public static async Task<JObject> DoModelDownloadWS(Session session, WebSocket w
530531
{
531532
File.Delete(tempPath);
532533
}
533-
Directory.CreateDirectory(Path.GetDirectoryName(outPath));
534+
Directory.CreateDirectory(Path.GetDirectoryName(modelOutPath));
534535
using CancellationTokenSource canceller = new();
535536
Task downloading = Utilities.DownloadFile(url, tempPath, (progress, total, perSec) =>
536537
{
@@ -570,17 +571,20 @@ public static async Task<JObject> DoModelDownloadWS(Session session, WebSocket w
570571
}
571572
});
572573
await downloading;
573-
File.Move(tempPath, outPath);
574+
File.Move(tempPath, modelOutPath);
574575
if (!string.IsNullOrWhiteSpace(metadata))
575576
{
576-
if (!string.IsNullOrWhiteSpace(baseModel))
577+
string metadataOutPath;
578+
if (!string.IsNullOrWhiteSpace(baseModel) && session.User.Settings.GroupDownloadedModelsByBaseType)
577579
{
578-
File.WriteAllText($"{handler.FolderPaths[0]}/{baseModel}/{name}.json", metadata);
580+
metadataOutPath = $"{handler.FolderPaths[0]}/{baseModel}/{name}.json";
579581
}
580582
else
581583
{
582-
File.WriteAllText($"{handler.FolderPaths[0]}/{name}.json", metadata);
584+
metadataOutPath = $"{handler.FolderPaths[0]}/{name}.json";
583585
}
586+
metadataOutPath = Utilities.StrictFilenameClean(metadataOutPath);
587+
File.WriteAllText(metadataOutPath, metadata);
584588
}
585589
await ws.SendJson(new JObject() { ["success"] = true }, API.WebsocketTimeout);
586590
}

0 commit comments

Comments
 (0)