swarm-assistent/AssistentInventory.cs

using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Threading.Tasks;
using Newtonsoft.Json.Linq;
using SwarmUI.Accounts;
using SwarmUI.Core;
using SwarmUI.Text2Image;
using SwarmUI.Utils;

namespace Mrleo1nid.SwarmAssistent;

/// <summary>Server-side model inventory (LoRA / checkpoint / wildcard lists).</summary>
public partial class SwarmAssistentExtension
{
    const int MaxLorasInInventoryFallback = 150;
    const int MaxWildcardsInInventoryFallback = 80;
    const int MaxCheckpointsInInventoryFallback = 60;
    const int InventoryBlurbMaxFallback = 140;

    /// <summary>Server-side LoRA / checkpoint / wildcard inventory (not DOM scrape).
    /// Pass rescan=true after downloads so new files appear (calls Program.RefreshAllModelSets).</summary>
    public async Task<JObject> AssistentListInventory(Session session, bool rescan = false)
    {
        await Task.CompletedTask;
        if (rescan)
        {
            try
            {
                Program.RefreshAllModelSets();
            }
            catch (Exception ex)
            {
                Logs.Debug($"AssistentListInventory rescan: {ex.Message}");
                try
                {
                    Program.ModelRefreshEvent?.Invoke();
                }
                catch (Exception ex2)
                {
                    Logs.Debug($"AssistentListInventory ModelRefreshEvent: {ex2.Message}");
                }
            }
        }

        JArray loras = [];
        JArray checkpoints = [];
        JArray wildcards = [];

        if (Program.T2IModelSets.TryGetValue("LoRA", out T2IModelHandler loraHandler))
        {
            foreach (T2IModel model in loraHandler.Models.Values
                .OrderByDescending(m => LooksLikeKreaArch(m))
                .ThenBy(m => m.Name)
                .Take(CfgInt("max_loras_inventory", MaxLorasInInventoryFallback)))
            {
                loras.Add(BuildInventoryModelEntry(model, "lora"));
            }
        }

        if (Program.T2IModelSets.TryGetValue("Stable-Diffusion", out T2IModelHandler ckptHandler))
        {
            foreach (T2IModel model in ckptHandler.Models.Values
                .OrderByDescending(m => LooksLikeKreaArch(m))
                .ThenBy(m => m.Name)
                .Take(CfgInt("max_checkpoints_inventory", MaxCheckpointsInInventoryFallback)))
            {
                checkpoints.Add(BuildInventoryModelEntry(model, "checkpoint"));
            }
        }

        try
        {
            foreach (string name in WildcardsHelper.ListFiles.OrderBy(n => n).Take(CfgInt("max_wildcards_inventory", MaxWildcardsInInventoryFallback)))
            {
                wildcards.Add(new JObject { ["name"] = name });
            }
        }
        catch (Exception ex)
        {
            Logs.Debug($"AssistentListInventory wildcards: {ex.Message}");
        }

        bool hasCivitaiKey = !string.IsNullOrWhiteSpace(session.User.GetGenericData("civitai_api", "key"));

        return new JObject
        {
            ["success"] = true,
            ["loras"] = loras,
            ["checkpoints"] = checkpoints,
            ["wildcards"] = wildcards,
            ["has_civitai_key"] = hasCivitaiKey,
            ["rescanned"] = rescan,
            ["inventory_at"] = DateTimeOffset.UtcNow.ToUnixTimeSeconds(),
        };
    }

    static bool LooksLikeKreaArch(T2IModel model)
    {
        string arch = model?.ModelClass?.ID ?? "";
        string compat = model?.ModelClass?.CompatClass?.ID ?? "";
        string name = model?.Name ?? "";
        string blob = $"{arch} {compat} {name}".ToLowerInvariant();
        return blob.Contains("krea");
    }

    JObject BuildInventoryModelEntry(T2IModel model, string kind)
    {
        string weight = null;
        try { weight = model.RawFilePath; } catch { /* ignore */ }

        string usage = model.Metadata?.UsageHint;
        string desc = model.Metadata?.Description;
        try
        {
            if (string.IsNullOrWhiteSpace(desc) && !string.IsNullOrWhiteSpace(model.Description))
            {
                desc = model.Description;
            }
        }
        catch
        {
            // older Swarm builds
        }

        string blurb = null;
        string raw = !string.IsNullOrWhiteSpace(usage) ? usage : desc;
        if (!string.IsNullOrWhiteSpace(raw))
        {
            blurb = Clip(CollapseWs(raw), CfgInt("inventory_blurb_max", InventoryBlurbMaxFallback));
        }

        JArray tags = null;
        if (model.Metadata?.Tags is { Length: > 0 } tagArr)
        {
            tags = new JArray(tagArr.Where(t => !string.IsNullOrWhiteSpace(t)).Take(8));
        }

        string trigger = model.Metadata?.TriggerPhrase;
        JArray triggers = null;
        if (!string.IsNullOrWhiteSpace(trigger))
        {
            triggers = new JArray(trigger.Split(new[] { ',', ';' }, StringSplitOptions.RemoveEmptyEntries | StringSplitOptions.TrimEntries).Take(12));
        }

        JObject entry = new()
        {
            ["name"] = model.Name,
            ["title"] = model.Metadata?.Title ?? model.Title ?? model.Name,
            ["kind"] = kind,
            ["trigger_phrase"] = trigger,
            ["architecture"] = model.ModelClass?.ID,
            ["compat_class"] = model.ModelClass?.CompatClass?.ID,
            ["hash"] = model.Metadata?.Hash ?? "",
            ["krea_likely"] = LooksLikeKreaArch(model),
        };
        if (!string.IsNullOrWhiteSpace(weight))
        {
            string stem = Path.GetFileNameWithoutExtension(weight);
            string dir = Path.GetDirectoryName(weight);
            string side = Path.Combine(dir ?? "", $"{stem}.civitai.json");
            entry["has_sidecar"] = File.Exists(side);
            foreach (string suffix in new[] { ".preview.jpg", ".preview.png", ".preview.jpeg", ".jpg", ".png", ".webp" })
            {
                string prev = Path.Combine(dir ?? "", stem + suffix);
                if (File.Exists(prev))
                {
                    string folder = kind == "lora" ? "Lora" : "Stable-Diffusion";
                    entry["preview_url"] = $"View/Models/{folder}/{Path.GetFileName(prev)}";
                    break;
                }
            }
        }
        else
        {
            entry["has_sidecar"] = false;
        }
        if (triggers is not null && triggers.Count > 0)
        {
            entry["triggers"] = triggers;
        }
        if (!string.IsNullOrWhiteSpace(blurb))
        {
            entry["blurb"] = blurb;
        }
        if (!string.IsNullOrWhiteSpace(usage))
        {
            entry["usage_hint"] = Clip(CollapseWs(usage), 120);
        }
        if (tags is not null && tags.Count > 0)
        {
            entry["tags"] = tags;
        }
        string defW = model.Metadata?.LoraDefaultWeight;
        if (!string.IsNullOrWhiteSpace(defW) && kind == "lora")
        {
            entry["default_weight"] = defW;
        }
        return entry;
    }
}