using System; using System.Collections.Generic; using System.IO; using System.Linq; using System.Net.Http; using System.Net.Http.Headers; using System.Text; using System.Text.RegularExpressions; using System.Threading.Tasks; using FreneticUtilities.FreneticExtensions; using Newtonsoft.Json.Linq; using SwarmUI.Accounts; using SwarmUI.Core; using SwarmUI.Text2Image; using SwarmUI.Utils; using SwarmUI.WebAPI; namespace Mrleo1nid.SwarmAssistent; /// Server-side model inventory, assistant cards and Civitai lookups. public partial class SwarmAssistentExtension { const int MaxLorasInInventoryFallback = 150; const int MaxWildcardsInInventoryFallback = 80; const int MaxCheckpointsInInventoryFallback = 60; const int InventoryBlurbMaxFallback = 140; static string ModelWeightPath(string setName, string modelName) { if (string.IsNullOrWhiteSpace(modelName) || !Program.T2IModelSets.TryGetValue(setName, out T2IModelHandler handler)) { return null; } if (!handler.Models.TryGetValue(modelName, out T2IModel model) && !handler.Models.TryGetValue(modelName.Replace('\\', '/'), out model)) { // Try suffix match model = handler.Models.Values.FirstOrDefault(m => string.Equals(m.Name, modelName, StringComparison.OrdinalIgnoreCase) || m.Name.EndsWith("/" + modelName, StringComparison.OrdinalIgnoreCase) || Path.GetFileNameWithoutExtension(m.Name) == Path.GetFileNameWithoutExtension(modelName)); } if (model is null) { return null; } try { // SwarmUI T2IModel exposes RawFilePath in recent builds. return model.RawFilePath; } catch { return null; } } static string CardPathForWeight(string weightPath) { if (string.IsNullOrWhiteSpace(weightPath)) { return null; } string dir = Path.GetDirectoryName(weightPath); string stem = Path.GetFileNameWithoutExtension(weightPath); if (string.IsNullOrWhiteSpace(dir) || string.IsNullOrWhiteSpace(stem)) { return null; } return Path.Combine(dir, $"{stem}.assistent.json"); } static string SetNameForKind(string kind) { return (kind ?? "").Trim().ToLowerInvariant() switch { "lora" => "LoRA", "checkpoint" or "ckpt" or "stable-diffusion" => "Stable-Diffusion", _ => null, }; } JObject ReadCardObject(string kind, string name) { string set = SetNameForKind(kind); string weight = ModelWeightPath(set, name); string card = CardPathForWeight(weight); if (card is null || !File.Exists(card)) { return null; } try { return JObject.Parse(File.ReadAllText(card, Encoding.UTF8)); } catch { return null; } } public async Task AssistentGetCard(Session session, string kind, string name) { await Task.CompletedTask; if (string.IsNullOrWhiteSpace(kind) || string.IsNullOrWhiteSpace(name)) { return new JObject { ["error"] = "kind and name required" }; } JObject card = ReadCardObject(kind, name); string set = SetNameForKind(kind); string weight = ModelWeightPath(set, name); return new JObject { ["success"] = true, ["kind"] = kind, ["name"] = name, ["has_card"] = card is not null, ["weight_path"] = weight, ["card"] = card, }; } public async Task AssistentSaveCard(Session session, string kind, string name, JObject card, bool enqueue_wanted = false) { await Task.CompletedTask; if (card is null) { return new JObject { ["error"] = "card required" }; } kind = (kind ?? card["kind"]?.ToString() ?? "").Trim(); name = (name ?? card["name"]?.ToString() ?? "").Trim(); if (string.IsNullOrWhiteSpace(kind) || string.IsNullOrWhiteSpace(name)) { return new JObject { ["error"] = "kind and name required" }; } card["kind"] = kind; card["name"] = name; string set = SetNameForKind(kind); string weight = ModelWeightPath(set, name); if (!string.IsNullOrWhiteSpace(weight) && File.Exists(weight)) { string path = CardPathForWeight(weight); File.WriteAllText(path, card.ToString(Newtonsoft.Json.Formatting.Indented), Encoding.UTF8); _ = IngestCardToMemory(card, name); return new JObject { ["success"] = true, ["path"] = path, ["installed"] = true }; } // Not installed — draft into wanted-cards + optionally enqueue download for next up. Directory.CreateDirectory(WantedCardsDir()); string rawVid = card["version_id"]?.ToString() ?? "draft"; string vid = Regex.IsMatch(rawVid, @"^\d+$") ? rawVid : "draft"; string draft = Path.Combine(WantedCardsDir(), $"{vid}.assistent.json"); File.WriteAllText(draft, card.ToString(Newtonsoft.Json.Formatting.Indented), Encoding.UTF8); if (enqueue_wanted || !string.IsNullOrWhiteSpace(card["civitai_url"]?.ToString())) { await AssistentEnqueueWanted(session, kind, card["civitai_url"]?.ToString(), card["version_id"]?.Value() ?? 0, card["title"]?.ToString() ?? name, card); } _ = IngestCardToMemory(card, name); return new JObject { ["success"] = true, ["path"] = draft, ["installed"] = false, ["wanted"] = true }; } async Task IngestCardToMemory(JObject card, string name) { if (Memory is null || card is null) { return; } try { string kind = (card["kind"]?.ToString() ?? "lora").Trim().ToLowerInvariant(); string key = (card["name"]?.ToString() ?? name ?? "").Trim(); List bits = []; foreach (string field in new[] { "when", "avoid", "prompt_hint", "notes" }) { string v = card[field]?.ToString(); if (!string.IsNullOrWhiteSpace(v)) { bits.Add($"{field}: {v.Trim()}"); } } if (card["triggers"] is JArray tr) { string joined = string.Join(", ", tr.Select(t => t?.ToString()).Where(s => !string.IsNullOrWhiteSpace(s))); if (!string.IsNullOrWhiteSpace(joined)) { bits.Add("triggers: " + joined); } } if (bits.Count == 0 || string.IsNullOrWhiteSpace(key)) { return; } string text = $"{kind} {key}. " + string.Join(" ", bits); string baseUrl = NormalizeBaseUrl(Config.LoadSettings()["base_url"]?.ToString()); string embedModel = Config.LoadSettings()["embed_model"]?.ToString() ?? Config.LoadAssistant(Config.DefaultPersonaId())["embed_model"]?.ToString(); await Memory.UpsertTextAsync(baseUrl, "card", key, text, "user", card, embedModel, AssistentMemory.SharedPersona); } catch (Exception ex) { Logs.Debug($"IngestCardToMemory: {ex.Message}"); } } public async Task AssistentGetCardMeta(Session session, string kind, string name, int version_id = 0, bool fetch = false) { string set = SetNameForKind(kind); string weight = ModelWeightPath(set, name); JObject civitai = null; JArray exampleUrls = []; JArray previewUrls = []; bool hasSidecar = false; string fetchError = null; bool fetched = false; if (!string.IsNullOrWhiteSpace(weight)) { string stem = Path.GetFileNameWithoutExtension(weight); string dir = Path.GetDirectoryName(weight); string side = Path.Combine(dir ?? "", $"{stem}.civitai.json"); if (File.Exists(side)) { hasSidecar = true; try { civitai = JObject.Parse(File.ReadAllText(side, Encoding.UTF8)); } catch { // ignore } } 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)) { // Swarm View path — relative URL works in the same origin browser session. previewUrls.Add($"View/Models/{(kind == "lora" ? "Lora" : "Stable-Diffusion")}/{Path.GetFileName(prev)}"); break; } } } if (civitai is not null) { if (version_id <= 0) { version_id = civitai["id"]?.Value() ?? 0; } CollectExampleUrls(civitai, exampleUrls); } string hash = null; string trigger = null; try { if (Program.T2IModelSets.TryGetValue(set, out T2IModelHandler h) && (h.Models.TryGetValue(name, out T2IModel m) || h.Models.TryGetValue(name.Replace('\\', '/'), out m))) { trigger = m.Metadata?.TriggerPhrase; hash = m.Metadata?.Hash; } } catch { // ignore } if (fetch && civitai is null) { string apiKey = session.User.GetGenericData("civitai_api", "key") ?? ""; if (string.IsNullOrWhiteSpace(apiKey)) { fetchError = "Civitai: нет ключа в User Settings"; } else { try { JObject remote = null; if (version_id > 0) { remote = await FetchCivitaiModelVersion(apiKey, version_id); } if (remote is null && !string.IsNullOrWhiteSpace(hash)) { string sha = hash.Trim().ToLowerInvariant(); if (sha.StartsWith("sha256:")) { sha = sha["sha256:".Length..]; } if (sha.Length == 64) { remote = await FetchCivitaiByHash(apiKey, sha); } else { fetchError ??= "Civitai: хеш модели не SHA256"; } } if (remote is not null) { civitai = remote; fetched = true; version_id = remote["id"]?.Value() ?? version_id; CollectExampleUrls(remote, exampleUrls); if (!string.IsNullOrWhiteSpace(weight)) { try { string stem = Path.GetFileNameWithoutExtension(weight); string dir = Path.GetDirectoryName(weight); string side = Path.Combine(dir ?? "", $"{stem}.civitai.json"); File.WriteAllText(side, remote.ToString(Newtonsoft.Json.Formatting.Indented), Encoding.UTF8); hasSidecar = true; } catch (Exception ex) { Logs.Debug($"AssistentGetCardMeta write sidecar: {ex.Message}"); } } } else if (fetchError is null) { fetchError = string.IsNullOrWhiteSpace(hash) ? "Civitai: нет hash и version_id" : "Хеш не найден на Civitai"; } } catch (Exception ex) { fetchError = $"Civitai: {ex.Message}"; } } } JObject card = ReadCardObject(kind, name); return new JObject { ["success"] = true, ["kind"] = kind, ["name"] = name, ["version_id"] = version_id, ["trigger_phrase"] = trigger, ["has_card"] = card is not null, ["has_sidecar"] = hasSidecar, ["fetched"] = fetched, ["fetch_error"] = fetchError, ["card"] = card, ["civitai"] = civitai, ["example_urls"] = exampleUrls, ["preview_urls"] = previewUrls, ["weight_path"] = weight, ["hash"] = hash, }; } static void CollectExampleUrls(JObject civitai, JArray exampleUrls) { if (civitai?["images"] is not JArray imgs) { return; } foreach (JToken img in imgs.Take(6)) { string u = img?["url"]?.ToString(); if (!string.IsNullOrWhiteSpace(u)) { exampleUrls.Add(u); } } } async Task FetchCivitaiByHash(string apiKey, string sha) { string[] hosts = ["civitai.red", "civitai.com"]; Exception last = null; foreach (string host in hosts) { try { string url = $"https://{host}/api/v1/model-versions/by-hash/{sha}"; using HttpRequestMessage req = new(HttpMethod.Get, url); if (!string.IsNullOrWhiteSpace(apiKey)) { req.Headers.Authorization = new AuthenticationHeaderValue("Bearer", apiKey.Trim()); } using HttpResponseMessage resp = await HttpClient.SendAsync(req); string body = await resp.Content.ReadAsStringAsync(); if (resp.StatusCode == System.Net.HttpStatusCode.NotFound) { continue; } if (!resp.IsSuccessStatusCode) { last = new Exception($"HTTP {(int)resp.StatusCode}: {Clip(body, 160)}"); if ((int)resp.StatusCode is 401 or 403) { throw last; } continue; } return JObject.Parse(body); } catch (Exception ex) when (ex is not HttpRequestException && ex.Message.Contains("401")) { throw; } catch (Exception ex) { last = ex; } } if (last is not null) { throw last; } return null; } async Task FetchCivitaiModelVersion(string apiKey, int versionId) { string[] hosts = ["civitai.red", "civitai.com"]; Exception last = null; foreach (string host in hosts) { try { string url = $"https://{host}/api/v1/model-versions/{versionId}"; using HttpRequestMessage req = new(HttpMethod.Get, url); if (!string.IsNullOrWhiteSpace(apiKey)) { req.Headers.Authorization = new AuthenticationHeaderValue("Bearer", apiKey.Trim()); } using HttpResponseMessage resp = await HttpClient.SendAsync(req); string body = await resp.Content.ReadAsStringAsync(); if (!resp.IsSuccessStatusCode) { last = new Exception($"HTTP {(int)resp.StatusCode}: {Clip(body, 160)}"); if ((int)resp.StatusCode is 401 or 403) { throw last; } continue; } return JObject.Parse(body); } catch (Exception ex) { last = ex; if (ex.Message.Contains("401") || ex.Message.Contains("403")) { throw; } } } if (last is not null) { throw last; } return null; } /// Server-side LoRA / checkpoint / wildcard inventory (not DOM scrape). /// Pass rescan=true after downloads so new files appear (calls Program.RefreshAllModelSets). public async Task 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 cardPath = CardPathForWeight(weight); bool hasCard = !string.IsNullOrWhiteSpace(cardPath) && File.Exists(cardPath); 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; if (hasCard) { try { JObject card = JObject.Parse(File.ReadAllText(cardPath, Encoding.UTF8)); string fromCard = (card["notes"] ?? card["when"] ?? card["prompt_hint"])?.ToString(); if (!string.IsNullOrWhiteSpace(fromCard)) { blurb = Clip(fromCard.Trim(), CfgInt("inventory_blurb_max", InventoryBlurbMaxFallback)); } } catch { // ignore bad card json } } if (string.IsNullOrWhiteSpace(blurb)) { 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 ?? "", ["has_card"] = hasCard, ["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; } /// Search Civitai for LoRAs (prefers Krea 2 base). Uses Swarm-stored civitai_api key. public async Task AssistentSearchCivitai(Session session, string query, int limit = 8) { string q = (query ?? "").Trim(); if (string.IsNullOrWhiteSpace(q)) { return new JObject { ["error"] = "query is required" }; } limit = Math.Clamp(limit, 1, 20); string apiKey = session.User.GetGenericData("civitai_api", "key") ?? ""; HashSet installedNames = CollectInstalledLoraNames(); HashSet installedHashes = CollectInstalledLoraHashes(); string[] hosts = ["civitai.red", "civitai.com"]; Exception lastEx = null; foreach (string host in hosts) { try { string url = $"https://{host}/api/v1/models?limit={limit}&types=LORA&query={Uri.EscapeDataString(q)}"; using HttpRequestMessage req = new(HttpMethod.Get, url); if (!string.IsNullOrWhiteSpace(apiKey)) { req.Headers.Authorization = new AuthenticationHeaderValue("Bearer", apiKey.Trim()); } using HttpResponseMessage resp = await HttpClient.SendAsync(req); string body = await resp.Content.ReadAsStringAsync(); if (!resp.IsSuccessStatusCode) { lastEx = new Exception($"HTTP {(int)resp.StatusCode}: {Clip(body, 200)}"); continue; } JObject parsed = JObject.Parse(body); JArray items = parsed["items"] as JArray ?? []; JArray results = []; foreach (JToken item in items) { if (item is not JObject mo) { continue; } JObject card = BuildCivitaiCard(mo, installedNames, installedHashes); if (card is not null) { results.Add(card); } } // Prefer Krea-compatible first JArray sorted = new(results.OrderByDescending(t => LooksLikeKrea(t["base_model"]?.ToString())).ThenBy(t => t["name"]?.ToString())); return new JObject { ["success"] = true, ["query"] = q, ["host"] = host, ["results"] = sorted, ["has_civitai_key"] = !string.IsNullOrWhiteSpace(apiKey), }; } catch (Exception ex) { lastEx = ex; } } return new JObject { ["error"] = $"Civitai search failed: {lastEx?.Message ?? "unknown"}" }; } static bool LooksLikeKrea(string text) => !string.IsNullOrEmpty(text) && Regex.IsMatch(text, @"krea", RegexOptions.IgnoreCase); static HashSet CollectInstalledLoraNames() { HashSet names = new(StringComparer.OrdinalIgnoreCase); if (!Program.T2IModelSets.TryGetValue("LoRA", out T2IModelHandler handler)) { return names; } foreach (T2IModel m in handler.Models.Values) { names.Add(m.Name); string leaf = m.Name.Replace('\\', '/').AfterLast('/'); if (!string.IsNullOrEmpty(leaf)) { names.Add(leaf); names.Add(Path.GetFileNameWithoutExtension(leaf)); } } return names; } static HashSet CollectInstalledLoraHashes() { HashSet hashes = new(StringComparer.OrdinalIgnoreCase); if (!Program.T2IModelSets.TryGetValue("LoRA", out T2IModelHandler handler)) { return hashes; } foreach (T2IModel m in handler.Models.Values) { string h = m.Metadata?.Hash; if (!string.IsNullOrWhiteSpace(h)) { hashes.Add(h.Trim().ToLowerInvariant()); } } return hashes; } static JObject BuildCivitaiCard(JObject model, HashSet installedNames, HashSet installedHashes) { string name = model["name"]?.ToString() ?? ""; JArray versions = model["modelVersions"] as JArray; JObject ver = versions?.FirstOrDefault() as JObject; if (ver is null) { return null; } string baseModel = ver["baseModel"]?.ToString() ?? ""; JArray trained = ver["trainedWords"] as JArray ?? []; List triggers = trained.Select(t => t?.ToString()).Where(s => !string.IsNullOrWhiteSpace(s)).Take(8).ToList(); JObject file = null; foreach (JToken f in ver["files"] as JArray ?? []) { if (f is JObject fo && (fo["primary"]?.Value() == true || (fo["name"]?.ToString() ?? "").EndsWith(".safetensors", StringComparison.OrdinalIgnoreCase))) { file = fo; break; } } file ??= (ver["files"] as JArray)?.FirstOrDefault() as JObject; string downloadUrl = file?["downloadUrl"]?.ToString() ?? ver["downloadUrl"]?.ToString() ?? ""; string fileName = file?["name"]?.ToString() ?? ""; string sha = file?["hashes"]?["SHA256"]?.ToString() ?? file?["hashes"]?["AutoV2"]?.ToString() ?? ""; string saveName = string.IsNullOrWhiteSpace(fileName) ? Regex.Replace(name, @"[^\w\-.]+", "_").Trim('_') : Path.GetFileNameWithoutExtension(fileName); bool already = false; if (!string.IsNullOrWhiteSpace(sha) && installedHashes.Contains(sha.Trim().ToLowerInvariant())) { already = true; } else if (installedNames.Contains(saveName) || installedNames.Contains(name) || installedNames.Contains(fileName)) { already = true; } return new JObject { ["id"] = model["id"], ["version_id"] = ver["id"], ["name"] = name, ["base_model"] = baseModel, ["krea_likely"] = LooksLikeKrea(baseModel), ["triggers"] = new JArray(triggers), ["download_url"] = downloadUrl, ["file_name"] = saveName, ["sha256"] = sha, ["already_installed"] = already, ["n_sfw"] = model["nsfw"]?.Value() ?? false, }; } }