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
mcmonkeyprojects · Sep 1, 2024 · 26f0b5c · 26f0b5c
1 parent 3bb0363
commit 26f0b5c
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 10 deletions.
diff --git 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
@@ -511,16 +511,17 @@ public static async Task<JObject> 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<JObject> 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<JObject> 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);
         }