indirection for model preview icons

Author: mcmonkey4eva

Diff for: docs/API.md

@@ -72,6 +72,7 @@ The following `GET` routes are also available:
 - `/fonts/*.woff2`, `/imgs/*.jpg`, `/imgs/*.png`, `/favicon.ico` - Various asset files for the interface.
 - `/View/*.*` - Returns saved outputs. By default usually in a format like `/View/(user)/raw/(year-month-day)/(file).(ext)`, but is user-customizable.
 - `/Output/*.*` - legacy output call format. Do not use.
+- `/ViewSpecial/*.*` - Returns special data views. Notable example, `/ViewSpecial/Stable-Diffusion/OfficialStableDiffusion/sd_xl_base_1.0.safetensors` gets the preview icon for that model.
 - `/Audio/*.*` - Returns saved user audio files (in `Data/Audio`).
 - `/ExtensionFile/(extension)/*.*` - gets a web asset file from an extension.
 - `/ComfyBackendDirect/*.*` - direct pass-through to a comfy instance, if the [ComfyUI Backend Extension](/src/BuiltinExtensions/ComfyUIBackend/README.md) is in use.

Diff for: src/Core/WebServer.cs

@@ -7,6 +7,7 @@
 using Microsoft.Extensions.Hosting;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Primitives;
+using Newtonsoft.Json.Linq;
 using SwarmUI.Accounts;
 using SwarmUI.Text2Image;
 using SwarmUI.Utils;
@@ -239,6 +240,7 @@ public void Prep()
         WebApp.Map("/API/{*Call}", API.HandleAsyncRequest);
         WebApp.MapGet("/Output/{*Path}", ViewOutput);
         WebApp.MapGet("/View/{*Path}", ViewOutput);
+        WebApp.MapGet("/ViewSpecial/{*Path}", ViewSpecial);
         WebApp.MapGet("/ExtensionFile/{*f}", ViewExtensionScript);
         WebApp.MapGet("/Audio/{*f}", ViewAudio);
         timer.Check("[Web] core maps");
@@ -552,4 +554,69 @@ public async Task ViewOutput(HttpContext context)
         await context.Response.Body.WriteAsync(data, Program.GlobalProgramCancel);
         await context.Response.CompleteAsync();
     }
+
+    /// <summary>Web route for viewing special images (eg model icons).</summary>
+    public async Task ViewSpecial(HttpContext context)
+    {
+        string path = context.Request.Path.ToString();
+        path = Uri.UnescapeDataString(path).Replace('\\', '/').Trim('/');
+        while (path.Contains("//"))
+        {
+            path = path.Replace("//", "/");
+        }
+        (string subtype, string name) = path.After("ViewSpecial/").BeforeAndAfter('/');
+        User user = GetUserFor(context);
+        if (user is null)
+        {
+            await context.YieldJsonOutput(null, 400, Utilities.ErrorObj("invalid or unauthorized", "invalid_user"));
+            return;
+        }
+        async Task yieldResult(string imageData)
+        {
+            Image img = Image.FromDataString(imageData);
+            context.Response.ContentType = img.MimeType();
+            context.Response.StatusCode = 200;
+            context.Response.ContentLength = img.ImageData.Length;
+            context.Response.Headers.CacheControl = $"private, max-age=2";
+            await context.Response.Body.WriteAsync(img.ImageData, Program.GlobalProgramCancel);
+            await context.Response.CompleteAsync();
+        }
+        if (!user.IsAllowedModel(name))
+        {
+            Logs.Verbose($"Not showing user '{user.UserID}' sub-type '{subtype}' model image '{name}': user restriction");
+        }
+        else
+        {
+            if (subtype == "Wildcards")
+            {
+                WildcardsHelper.Wildcard card = WildcardsHelper.GetWildcard(name);
+                if (card is not null && card.Image.StartsWithFast("data:"))
+                {
+                    await yieldResult(card.Image);
+                    return;
+                }
+            }
+            if (Program.T2IModelSets.TryGetValue(subtype, out T2IModelHandler handler))
+            {
+                if (handler.Models.TryGetValue(name + ".safetensors", out T2IModel model) || handler.Models.TryGetValue(name, out model))
+                {
+                    if (model.Metadata?.PreviewImage?.StartsWithFast("data:") ?? false)
+                    {
+                        await yieldResult(model.Metadata.PreviewImage);
+                        return;
+                    }
+                }
+                else if (ModelsAPI.InternalExtraModels(subtype).TryGetValue(name + ".safetensors", out JObject remoteModel) || ModelsAPI.InternalExtraModels(subtype).TryGetValue(name, out remoteModel))
+                {
+                    if (remoteModel.TryGetValue("preview_image", out JToken previewImg) && previewImg.ToString().StartsWithFast("data:"))
+                    {
+                        await yieldResult(previewImg.ToString());
+                        return;
+                    }
+                }
+            }
+            Logs.Verbose($"Not showing user '{user.UserID}' sub-type '{subtype}' model image '{name}': not found");
+        }
+        await context.YieldJsonOutput(null, 404, Utilities.ErrorObj("404, file not found.", "file_not_found"));
+    }
 }

Diff for: src/Text2Image/T2IModel.cs

@@ -3,6 +3,7 @@
 using Newtonsoft.Json.Linq;
 using SwarmUI.Core;
 using SwarmUI.Utils;
+using SwarmUI.WebAPI;
 using System.IO;
 using System.Security.Cryptography;
 
@@ -265,15 +266,20 @@ void HandleResave(string path)
     }
 
     /// <summary>Gets a networkable copy of this model's data.</summary>
-    public JObject ToNetObject(string prefix = "")
+    public JObject ToNetObject(string prefix = "", bool dataImgs = true)
     {
+        string previewImg = PreviewImage;
+        if (!dataImgs && previewImg is not null && previewImg.StartsWithFast("data:"))
+        {
+            previewImg = $"/ViewSpecial/{Handler.ModelType}/{name}?editid={ModelsAPI.ModelEditID}";
+        }
         return new JObject()
         {
             [$"{prefix}name"] = Name,
             [$"{prefix}title"] = Metadata?.Title,
             [$"{prefix}author"] = Metadata?.Author,
             [$"{prefix}description"] = Description,
-            [$"{prefix}preview_image"] = PreviewImage,
+            [$"{prefix}preview_image"] = previewImg,
             [$"{prefix}loaded"] = AnyBackendsHaveLoaded,
             [$"{prefix}architecture"] = ModelClass?.ID,
             [$"{prefix}class"] = ModelClass?.Name,

Diff for: src/Utils/WildcardsHelper.cs

@@ -2,6 +2,7 @@
 using FreneticUtilities.FreneticToolkit;
 using Newtonsoft.Json.Linq;
 using SwarmUI.Core;
+using SwarmUI.WebAPI;
 using System.IO;
 
 namespace SwarmUI.Utils;
@@ -23,14 +24,19 @@ public class Wildcard
 
         public long TimeCreated;
 
-        public JObject GetNetObject()
+        public JObject GetNetObject(bool dataImgs = true)
         {
+            string previewImg = Image ?? "imgs/model_placeholder.jpg";
+            if (!dataImgs && previewImg is not null && previewImg.StartsWithFast("data:"))
+            {
+                previewImg = $"/ViewSpecial/Wildcards/{Name}?editid={ModelsAPI.ModelEditID}";
+            }
             return new()
             {
                 ["name"] = Name,
                 ["options"] = JArray.FromObject(Options),
                 ["raw"] = Raw,
-                ["image"] = Image ?? "imgs/model_placeholder.jpg"
+                ["image"] = previewImg
             };
         }
     }

Diff for: src/WebAPI/ModelsAPI.cs

@@ -15,6 +15,8 @@ namespace SwarmUI.WebAPI;
 [API.APIClass("API routes related to handling models (including loras, wildcards, etc).")]
 public static class ModelsAPI
 {
+    public static long ModelEditID = 0;
+
     public static void Register()
     {
         API.RegisterAPICall(ListModels, false, Permissions.FundamentalModelAccess);
@@ -198,7 +200,7 @@ bool tryMatch(string name)
                 if (tryMatch(file))
                 {
                     WildcardsHelper.Wildcard card = WildcardsHelper.GetWildcard(file);
-                    files.Add(new(card.Name, card.Name.AfterLast('/'), card.TimeCreated, card.TimeModified, card.GetNetObject()));
+                    files.Add(new(card.Name, card.Name.AfterLast('/'), card.TimeCreated, card.TimeModified, card.GetNetObject(false)));
                     if (files.Count > sanityCap)
                     {
                         break;
@@ -212,7 +214,7 @@ bool tryMatch(string name)
             {
                 if (tryMatch(possible.Name))
                 {
-                    files.Add(new(possible.Name, possible.Title, possible.Metadata?.TimeCreated ?? long.MaxValue, possible.Metadata?.TimeModified ?? long.MaxValue, possible.ToNetObject()));
+                    files.Add(new(possible.Name, possible.Title, possible.Metadata?.TimeCreated ?? long.MaxValue, possible.Metadata?.TimeModified ?? long.MaxValue, possible.ToNetObject(dataImgs: false)));
                     if (files.Count > sanityCap)
                     {
                         break;
@@ -226,7 +228,13 @@ bool tryMatch(string name)
             {
                 if (tryMatch(name))
                 {
-                    files.Add(new(name, name.AfterLast('/'), long.MaxValue, long.MaxValue, possible));
+                    JObject toAdd = possible;
+                    if (toAdd.TryGetValue("preview_image", out JToken previewImg) && previewImg.ToString().StartsWith("data:"))
+                    {
+                        toAdd = toAdd.DeepClone() as JObject;
+                        toAdd["preview_image"] = $"/ViewSpecial/{subtype}/{name}";
+                    }
+                    files.Add(new(name, name.AfterLast('/'), long.MaxValue, long.MaxValue, toAdd));
                     if (files.Count > sanityCap)
                     {
                         break;
@@ -403,6 +411,7 @@ public static async Task<JObject> EditWildcard(Session session,
             File.WriteAllBytes($"{WildcardsHelper.Folder}/{card}.jpg", img.ImageData);
             WildcardsHelper.WildcardFiles[card] = new WildcardsHelper.Wildcard() { Name = card };
         }
+        Interlocked.Increment(ref ModelEditID);
         return new JObject() { ["success"] = true };
     }
 
@@ -473,6 +482,7 @@ public static async Task<JObject> EditModelMetadata(Session session,
         }
         handler.ResetMetadataFrom(actualModel);
         _ = Utilities.RunCheckedTask(() => actualModel.ResaveModel());
+        Interlocked.Increment(ref ModelEditID);
         return new JObject() { ["success"] = true };
     }
 
@@ -804,6 +814,7 @@ void doMoveNow(string oldPath)
                 doMoveNow(altPath);
             }
         }
+        Interlocked.Increment(ref ModelEditID);
         return new JObject() { ["success"] = true };
     }
 }

Diff for: src/wwwroot/js/genpage/models.js

@@ -102,7 +102,7 @@ function save_edit_wildcard() {
         }
         genericRequest('EditWildcard', data, resData => {
             wildcardsBrowser.browser.refresh();
-            if (card.name != data.card) {
+            if (card.name && card.name != data.card) {
                 genericRequest('DeleteWildcard', { card: card.name }, data => {});
             }
         });