From 3bb0363d7cc3a14d68f4234590f788953c86db9c Mon Sep 17 00:00:00 2001 From: Aimeri Baddouh Date: Sat, 31 Aug 2024 11:00:07 -0400 Subject: [PATCH 1/2] Save Models in baseModel subfolder Add the functionality of reading baseModel from civitai metadata, and save models in their base model subfolders. This will look like: Models/Stable-diffusion/ SDXL 1.0/ SD1.5/ Pony/ Flux.1 D/ Flux.1 F/ etc... --- src/WebAPI/ModelsAPI.cs | 25 +++++++++++++++++++++++-- src/wwwroot/js/genpage/utiltab.js | 1 + 2 files changed, 24 insertions(+), 2 deletions(-) diff --git a/src/WebAPI/ModelsAPI.cs b/src/WebAPI/ModelsAPI.cs index ceeb00e65..2075b3a9d 100644 --- a/src/WebAPI/ModelsAPI.cs +++ b/src/WebAPI/ModelsAPI.cs @@ -505,7 +505,21 @@ public static async Task DoModelDownloadWS(Session session, WebSocket w } try { - string outPath = $"{handler.FolderPaths[0]}/{name}.safetensors"; + string baseModel = ""; + if (!string.IsNullOrWhiteSpace(metadata)) + { + JObject metadataObj = JObject.Parse(metadata); + baseModel = metadataObj["modelspec.baseModel"]?.ToString() ?? ""; + } + string outPath; + if (!string.IsNullOrWhiteSpace(baseModel)) + { + outPath = $"{handler.FolderPaths[0]}/{baseModel}/{name}.safetensors"; + } + else + { + outPath = $"{handler.FolderPaths[0]}/{name}.safetensors"; + } if (File.Exists(outPath)) { await ws.SendJson(new JObject() { ["error"] = "Model at that save path already exists." }, API.WebsocketTimeout); @@ -559,7 +573,14 @@ public static async Task DoModelDownloadWS(Session session, WebSocket w File.Move(tempPath, outPath); if (!string.IsNullOrWhiteSpace(metadata)) { - File.WriteAllText($"{handler.FolderPaths[0]}/{name}.json", metadata); + if (!string.IsNullOrWhiteSpace(baseModel)) + { + File.WriteAllText($"{handler.FolderPaths[0]}/{baseModel}/{name}.json", metadata); + } + else + { + File.WriteAllText($"{handler.FolderPaths[0]}/{name}.json", metadata); + } } await ws.SendJson(new JObject() { ["success"] = true }, API.WebsocketTimeout); } diff --git a/src/wwwroot/js/genpage/utiltab.js b/src/wwwroot/js/genpage/utiltab.js index 7ba54bf97..95496104a 100644 --- a/src/wwwroot/js/genpage/utiltab.js +++ b/src/wwwroot/js/genpage/utiltab.js @@ -259,6 +259,7 @@ class ModelDownloaderUtil { 'modelspec.author': rawData.creator.username, 'modelspec.description': `From ${url}\n${rawVersion.description || ''}\n${rawData.description}\n`, 'modelspec.date': rawVersion.createdAt, + 'modelspec.baseModel': rawVersion.baseModel, }; if (rawVersion.trainedWords) { metadata['modelspec.trigger_phrase'] = rawVersion.trainedWords.join(", "); From 26f0b5c645fb8c8ca4bf71d023ae589ebe4cfad7 Mon Sep 17 00:00:00 2001 From: Aimeri Baddouh Date: Sun, 1 Sep 2024 09:31:12 -0400 Subject: [PATCH 2/2] 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 --- src/Core/Settings.cs | 3 +++ src/WebAPI/ModelsAPI.cs | 24 ++++++++++++++---------- 2 files changed, 17 insertions(+), 10 deletions(-) diff --git a/src/Core/Settings.cs b/src/Core/Settings.cs index 05d7376fd..34c0791bb 100644 --- a/src/Core/Settings.cs +++ b/src/Core/Settings.cs @@ -273,6 +273,9 @@ public class FileFormatData : AutoConfiguration [ConfigComment("If true, folders will be discarded from starred image paths.")] public bool StarNoFolders = false; + [ConfigComment("Whether to automatically use the base model type as part of the model path when downloading using the 'Model Download' tool")] + public bool GroupDownloadedModelsByBaseType = false; + public class ThemesImpl : SettingsOptionsAttribute.AbstractImpl { public override string[] GetOptions => [.. Program.Web.RegisteredThemes.Keys]; diff --git a/src/WebAPI/ModelsAPI.cs b/src/WebAPI/ModelsAPI.cs index 2075b3a9d..04cb0fc47 100644 --- a/src/WebAPI/ModelsAPI.cs +++ b/src/WebAPI/ModelsAPI.cs @@ -511,16 +511,17 @@ public static async Task DoModelDownloadWS(Session session, WebSocket w JObject metadataObj = JObject.Parse(metadata); baseModel = metadataObj["modelspec.baseModel"]?.ToString() ?? ""; } - string outPath; - if (!string.IsNullOrWhiteSpace(baseModel)) + string modelOutPath; + if (!string.IsNullOrWhiteSpace(baseModel) && session.User.Settings.GroupDownloadedModelsByBaseType) { - outPath = $"{handler.FolderPaths[0]}/{baseModel}/{name}.safetensors"; + modelOutPath = $"{handler.FolderPaths[0]}/{baseModel}/{name}.safetensors"; } else { - outPath = $"{handler.FolderPaths[0]}/{name}.safetensors"; + modelOutPath = $"{handler.FolderPaths[0]}/{name}.safetensors"; } - if (File.Exists(outPath)) + modelOutPath = Utilities.StrictFilenameClean(modelOutPath); + if (File.Exists(modelOutPath)) { await ws.SendJson(new JObject() { ["error"] = "Model at that save path already exists." }, API.WebsocketTimeout); return null; @@ -530,7 +531,7 @@ public static async Task DoModelDownloadWS(Session session, WebSocket w { File.Delete(tempPath); } - Directory.CreateDirectory(Path.GetDirectoryName(outPath)); + Directory.CreateDirectory(Path.GetDirectoryName(modelOutPath)); using CancellationTokenSource canceller = new(); Task downloading = Utilities.DownloadFile(url, tempPath, (progress, total, perSec) => { @@ -570,17 +571,20 @@ public static async Task DoModelDownloadWS(Session session, WebSocket w } }); await downloading; - File.Move(tempPath, outPath); + File.Move(tempPath, modelOutPath); if (!string.IsNullOrWhiteSpace(metadata)) { - if (!string.IsNullOrWhiteSpace(baseModel)) + string metadataOutPath; + if (!string.IsNullOrWhiteSpace(baseModel) && session.User.Settings.GroupDownloadedModelsByBaseType) { - File.WriteAllText($"{handler.FolderPaths[0]}/{baseModel}/{name}.json", metadata); + metadataOutPath = $"{handler.FolderPaths[0]}/{baseModel}/{name}.json"; } else { - File.WriteAllText($"{handler.FolderPaths[0]}/{name}.json", metadata); + metadataOutPath = $"{handler.FolderPaths[0]}/{name}.json"; } + metadataOutPath = Utilities.StrictFilenameClean(metadataOutPath); + File.WriteAllText(metadataOutPath, metadata); } await ws.SendJson(new JObject() { ["success"] = true }, API.WebsocketTimeout); }