Expand Assistent to v0.3: Generate loop, Civitai Confirm, img2img/inpaint.

Server inventory and streaming chat, auto-apply/generate with Interrupt, Civitai search cards (Confirm-only download), plus Init/Mask wiring and inpaint_edit pack.

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Leonid Pershin
2026-08-21 13:29:26 +03:00
co-authored by Cursor
parent 06d5df9e93
commit 3380206c6a
11 changed files with 1701 additions and 156 deletions
+54
View File
@@ -201,6 +201,11 @@
max-width: 12rem;
}
.sa-header-right .sa-model-select {
min-width: 14rem;
max-width: 24rem;
}
.sa-check {
flex-direction: row !important;
align-items: center;
@@ -279,6 +284,55 @@
margin-top: 0.45rem;
}
.sa-danger {
opacity: 0.85;
}
.sa-civitai-list {
display: flex;
flex-direction: column;
gap: 0.45rem;
margin-top: 0.55rem;
}
.sa-civitai-card {
padding: 0.5rem 0.6rem;
border-radius: 0.45rem;
border: 1px solid color-mix(in srgb, currentColor 22%, transparent);
background: color-mix(in srgb, currentColor 5%, transparent);
font-size: 0.88rem;
}
.sa-civitai-card.sa-installed {
opacity: 0.7;
}
.sa-civitai-title {
font-weight: 600;
margin-bottom: 0.2rem;
}
.sa-civitai-meta {
opacity: 0.8;
font-size: 0.82rem;
margin-bottom: 0.35rem;
}
.sa-civitai-actions {
display: flex;
flex-wrap: wrap;
gap: 0.35rem;
}
.sa-streaming {
opacity: 0.9;
font-style: italic;
}
.sa-interrupt-active {
border-color: color-mix(in srgb, #c44 55%, transparent) !important;
}
.sa-composer {
border-top: 1px solid color-mix(in srgb, currentColor 18%, transparent);
padding: 0.6rem 0.7rem 0.7rem;
+944 -109
View File
File diff suppressed because it is too large Load Diff
+38 -4
View File
@@ -15,15 +15,18 @@ You are **Swarm Assistent**, a collaborative art director for **Krea 2** image g
A JSON block named "Live SwarmUI context" is attached. Treat it as ground truth:
- Use only LoRAs listed in `available_loras` (by exact `name`).
- Use only LoRAs listed in `available_loras` (by exact `name`), or candidates from a Civitai search round.
- Prefer listed `trigger_phrase` / `triggers`**never invent** trigger words.
- When enabling a LoRA, include its triggers in `prompt` if missing.
- Respect current width/height/steps/cfg unless the user asks to change them or the pack is `fix_params`.
- Respect current width/height/steps/cfg/seed/sigma_shift/sampler unless the user asks to change them or the pack is `fix_params`.
- `wildcards` lists installed wildcard names (`__name__` syntax in prompts).
- `prompt_image_count` > 0 means Prompt Images are attached — warn if they may dominate.
- **Init / inpaint:** `has_init_image`, `has_mask_image`, `init_creativity` (aka denoise, 01), `mask_blur`, `mask_grow`. `has_vision_image` is the Assistent pane reference (can become Init/Mask).
## Output contract (mandatory)
1. Write a short helpful reply in the user's language (RU or EN).
2. Then emit **one** fenced JSON patch (and only fields you want to change):
2. Then emit **one** fenced JSON patch (only fields you want to change):
```json
{
@@ -34,14 +37,45 @@ A JSON block named "Live SwarmUI context" is attached. Treat it as ground truth:
"height": 1280,
"steps": 8,
"cfg": 1,
"seed": -1,
"sigma_shift": 1.15,
"sampler": null,
"use_init_image": false,
"clear_init_image": false,
"init_creativity": 0.45,
"use_mask_image": false,
"clear_mask_image": false,
"mask_blur": null,
"mask_grow": null,
"actions": ["generate"],
"search_query": null,
"notes": "one-line why"
}
```
Rules for the patch:
### Patch rules
- Omit keys you are not changing.
- `loras` replaces the intended LoRA set for Apply (list all that should be on).
- width/height between 128 and 4096; prefer multiples near 1024 for Turbo.
- Do not invent model or LoRA filenames.
- If you cannot help (wrong architecture / no Krea 2), say so and omit the JSON patch.
### Init image / inpaint
- **img2img:** set `use_init_image: true` (uses Assistent vision / current Generate image) and `init_creativity` (0 = almost copy, 1 = almost new). Typical edits: **0.250.45**; restyle: **0.50.7**. Alias `denoise` is accepted.
- **Inpaint:** needs Init + Mask. Set `use_init_image: true` and `use_mask_image: true` only when the vision pane holds a proper mask (white = edit, black = keep). If the user has not painted a mask, tell them to use Swarm **Edit Image** / paint a mask, or press **As Mask** with a prepared mask — do not invent pixel masks.
- `clear_init_image` / `clear_mask_image` to leave img2img mode.
- Prompt Images ≠ Init Image. Prefer Init for structural edits; Prompt Images for style refs (and warn they can dominate).
### Actions (auto-safe)
- `"generate"` — after Apply, start generation (UI auto-generate is on by default).
- `"use_init"` / `"use_mask"` — same as the boolean flags (optional).
- `"search_civitai"` — Civitai search; user must **Confirm** downloads.
- `"interrupt"` — stop generation.
- Pure Q&A with no prompt/param change: omit the JSON patch entirely (do not burn GPU).
### Auto-apply note
The UI may auto-apply your patch and auto-generate when `actions` contains `generate` or when you change prompt/loras/size/init. Keep patches intentional.
+2
View File
@@ -8,7 +8,9 @@ Goal: co-create a scene / moodboard direction for **Krea 2**.
- Propose one strong prompt (not five weak ones).
- Optionally suggest which available LoRAs fit — only from the live list, with triggers.
- Mention Prompt Images only if a reference would help, and warn that refs can overpower text.
- Missing style LoRA → `search_civitai` + `search_query` (Krea-compatible).
## Deliverable
- Scene brief + JSON patch (`prompt`, optional `loras`, optional aspect).
- Add `actions: ["generate"]` when ready to try the scene.
+4 -1
View File
@@ -9,8 +9,11 @@ Goal: look at the attached image (vision) and improve the next generation for **
- If a LoRA trigger was missing or too strong, adjust weight or prompt placement.
- If the frame needs a different aspect (too tight / too wide), change width/height.
- Prompt Images overpower text on Krea 2 — if the user relied on a ref, suggest weaker reliance or clearer text.
- For **local fixes** (hands, face, object): prefer **inpaint** (`use_init_image` + mask) over rewriting the whole prompt; if no mask yet, say so and suggest painting one / pack `inpaint_edit`.
- For **global restyle**: img2img with moderate `init_creativity` (≈0.40.6) can be better than from-scratch.
## Deliverable
- Short critique in the user's language.
- JSON patch with improved `prompt` and any `loras` / size tweaks.
- JSON patch with improved `prompt` and any `loras` / size / init tweaks.
- Include `actions: ["generate"]` when proposing a revised generation (default for this mode).
+5 -1
View File
@@ -7,10 +7,14 @@ Goal: adjust **generation parameters** for Krea 2 Turbo (or Raw if context says
- Turbo: prefer steps 412 (default 8), CFG ~1, sigma shift ~1.15.
- Raw/base: higher steps (20+) and higher CFG may apply — only if context indicates Raw.
- Aspect: change width/height for framing (portrait/landscape/square); keep near 1024 unless asked for higher res.
- Seed: set `seed` when the user wants reproducibility; `-1` for random.
- Sampler: only change if the form exposes it and the user asks.
- **Init creativity** (`init_creativity` / denoise 01): only when `has_init_image` or enabling img2img — see pack `inpaint_edit`.
- Do not change the prompt unless needed to match the new framing.
- Keep LoRAs unless the user asks to drop them.
## Deliverable
- Explain the param change.
- JSON patch focusing on `width`, `height`, `steps`, `cfg` (and `prompt` only if necessary).
- JSON patch focusing on `width`, `height`, `steps`, `cfg`, `seed`, `sigma_shift`, optional `init_creativity` (and `prompt` only if necessary).
- Include `actions: ["generate"]` if the user wants to re-roll with the new params.
+37
View File
@@ -0,0 +1,37 @@
# Mode: inpaint_edit
Goal: guide **img2img** (Init Image) and **inpainting** (Init + Mask) on **Krea 2**.
## When to use which
| Need | Setup |
| --- | --- |
| Soft edit / restyle whole frame | Init only + `init_creativity` |
| Change one region (face, hand, logo) | Init + Mask (white = edit) |
| Fresh image from text | Clear init/mask; normal t2i |
## Creativity (denoise)
- **0.20.35** — small fixes, keep composition
- **0.40.55** — noticeable edit, still related
- **0.60.8** — strong restyle; structure may drift
- Always set `use_init_image: true` when enabling img2img from the Assistent vision / current image.
## Mask rules
- White = regenerate, black = preserve. Gray = partial.
- Only set `use_mask_image: true` if context shows a vision image that is meant as a mask, or the user said they prepared one.
- If `has_mask_image` is false and the user wants regional edit: ask them to paint a mask in Swarm Image Editor (or Assistent **As Mask**), then continue.
- Optional: `mask_blur` / `mask_grow` for softer edges.
## Prompting
- Describe **what should appear in the edited region**, not the whole scene dump.
- Keep LoRA triggers if the subject depends on them.
- Match width/height to the init image when possible.
## Deliverable
- Short plan (img2img vs inpaint) in the user's language.
- JSON patch with `use_init_image` / `init_creativity` (and mask fields when applicable) + improved `prompt`.
- `actions: ["generate"]` when ready to run.
+2
View File
@@ -9,9 +9,11 @@ Goal: craft or improve a **Krea 2** prompt that will generate well on Turbo.
- Prefer clarity over keyword stuffing. Krea 2 understands sentences.
- If the user wants a style covered by an available LoRA, enable that LoRA and weave its triggers in.
- Keep Turbo defaults unless the user asks otherwise (steps 8, cfg 1).
- If a needed LoRA is missing from `available_loras`, use `actions: ["search_civitai"]` with a clear `search_query` (and prefer Krea base).
## Deliverable
- Explain briefly what you changed.
- Emit a JSON patch with at least `prompt`, and `loras` when relevant.
- Include `actions: ["generate"]` when the user wants to see a new image.
- Include `width`/`height` only if aspect should change for the scene (e.g. portrait → taller).
+48 -7
View File
@@ -1,13 +1,13 @@
# Swarm Assistent
SwarmUI extension for **collaborative Krea 2** prompting via **Ollama**: chat + vision, LoRA/trigger awareness, and applyable prompt/size patches.
SwarmUI extension for **collaborative Krea 2** prompting via **Ollama**: chat + vision, LoRA/trigger awareness, applyable patches, **img2img / inpaint**, auto Generate, and Civitai search with Confirm.
## Layout
- **Left:** vision reference (drop / paste / *Send to Assistent* from Generate)
- **Splitter:** drag to resize panes
- **Right (wider):** chat
- **Top-right:** prompt pack + settings (Ollama URL, model, auto-attach)
- **Top-right:** prompt pack + settings (Ollama URL, model, auto-apply / auto-generate)
## UX
@@ -15,13 +15,19 @@ SwarmUI extension for **collaborative Krea 2** prompting via **Ollama**: chat +
- Drag images from Generate/History or the OS onto the left pane
- Paste (`Ctrl+V`) while the Assistent tab is open
- **Use current** / **Clear** for the vision reference
- **As Init** / **As Mask** / **Clear Init** — wire the vision image into Swarm `Init Image` / `Mask Image` (img2img & inpaint)
- Enter sends; Shift+Enter newline
- **Auto-apply** + **Auto-generate** (default on): patch from the LLM is applied and Generate runs when the patch changes prompt/params or includes `actions: ["generate"]`
- Pure Q&A without a patch does **not** start Generate
- **Interrupt** stops Swarm generation / clears busy state
- **Civitai** search cards require **Confirm download** (uses Swarm `DoModelDownloadWS` + stored `civitai_api` key). Auto-download is off by default.
## Requirements
- SwarmUI with a **Krea 2** checkpoint selected
- Ollama on `http://127.0.0.1:11434` (gpu-rent default when `LLM_RUNTIME=ollama`)
- A chat+vision-capable Ollama model recommended for critique mode
- Ollama on `http://127.0.0.1:11434` **on the GPU VM** (gpu-rent `LLM_RUNTIME=ollama`). The browser talks to SwarmUI; SwarmUI proxies `/api/tags` and `/api/chat`. URL in settings must stay `127.0.0.1:11434`, not the laptop tunnel port 17811.
- At least one pulled model (`ollama pull` / `ollama-models.yaml`). Empty `/api/tags` → empty Model dropdown.
- Optional: Civitai API key in SwarmUI User Settings for search/download.
## Install
@@ -41,13 +47,48 @@ Restart / rebuild SwarmUI after clone.
| Pack | Role |
| --- | --- |
| `base_krea2` | Always injected: Krea 2 rules + JSON patch contract |
| `base_krea2` | Always injected: Krea 2 rules + JSON patch / actions contract |
| `write_prompt` | Craft / improve prompts |
| `critique_image` | Vision critique → fixes |
| `compose_scene` | Scene / moodboard |
| `fix_params` | Width/height/steps/CFG |
| `fix_params` | Width/height/steps/CFG/seed/σ-shift |
| `inpaint_edit` | Init Image img2img + Mask inpaint |
Live context (checkpoint, LoRAs + triggers, current params) is injected every request.
Live context (checkpoint, server inventory LoRAs + triggers, wildcards, current params) is injected every request.
### Patch actions
```json
{
"prompt": "...",
"loras": [{"name": "exact", "weight": 0.8, "triggers": ["..."]}],
"width": 1024, "height": 1280, "steps": 8, "cfg": 1,
"seed": -1, "sigma_shift": 1.15,
"use_init_image": true,
"init_creativity": 0.45,
"use_mask_image": false,
"actions": ["generate"],
"search_query": null
}
```
- `generate` — auto-generate after apply (when enabled)
- `use_init` / `use_mask` — set Assistent vision (or current Generate) as Init / Mask
- `search_civitai` + `search_query` — server searches Civitai, second LLM hop, Confirm cards in UI
- `interrupt` — stop current generation
Mask convention: **white = edit**, black = keep. Creativity ≈ denoise (01).
## API routes
| Route | Role |
| --- | --- |
| `AssistentListModels` | Ollama `/api/tags` |
| `AssistentListInventory` | LoRA / checkpoint / wildcard inventory from Swarm |
| `AssistentSearchCivitai` | Civitai LoRA search |
| `AssistentGetPacks` | Prompt pack texts |
| `AssistentChat` | HTTP chat (+ Civitai hop) |
| `AssistentChatWS` | Streaming chat WebSocket |
## License
+553 -30
View File
@@ -3,18 +3,22 @@ using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Net.Http;
using System.Net.Http.Headers;
using System.Net.WebSockets;
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;
/// <summary>Krea 2 collaborative assistant: Ollama chat + vision + prompt/LoRA/params patches.</summary>
/// <summary>Krea 2 collaborative assistant: Ollama chat + vision + prompt/LoRA/params patches + Generate/Civitai.</summary>
public class SwarmAssistentExtension : Extension
{
public static PermInfo PermUse = Permissions.Register(new(
@@ -33,17 +37,24 @@ public class SwarmAssistentExtension : Extension
"critique_image",
"compose_scene",
"fix_params",
"inpaint_edit",
];
const int MaxCivitaiHops = 2;
const int MaxLorasInInventory = 120;
const int MaxWildcardsInInventory = 80;
static readonly Regex JsonFenceRe = new(@"```(?:json)?\s*([\s\S]*?)```", RegexOptions.IgnoreCase | RegexOptions.Compiled);
public override void OnPreInit()
{
ScriptFiles.Add("Assets/assistent.js");
StyleSheetFiles.Add("Assets/assistent.css");
ExtensionAuthor = "mrleo1nid";
Description = "Collaborative Krea 2 assistant via Ollama: chat, vision, prompts, LoRA triggers, size patches.";
Description = "Collaborative Krea 2 assistant via Ollama: chat, vision, img2img/inpaint, Generate loop, Civitai Confirm.";
License = "MIT";
Version = "0.2.0";
Tags = ["tabs", "ui", "llm", "ollama", "krea"];
Version = "0.3.1";
Tags = ["tabs", "ui", "llm", "ollama", "krea", "inpaint"];
}
public override void OnInit()
@@ -51,8 +62,11 @@ public class SwarmAssistentExtension : Extension
HttpClient ??= new HttpClient { Timeout = TimeSpan.FromMinutes(10) };
API.RegisterAPICall(AssistentListModels, false, PermUse);
API.RegisterAPICall(AssistentGetPacks, false, PermUse);
API.RegisterAPICall(AssistentListInventory, false, PermUse);
API.RegisterAPICall(AssistentSearchCivitai, false, PermUse);
API.RegisterAPICall(AssistentChat, true, PermUse);
Logs.Init("Swarm Assistent extension loaded (Ollama proxy + Krea 2 packs)");
API.RegisterAPICall(AssistentChatWS, true, PermUse);
Logs.Init("Swarm Assistent extension loaded (Ollama proxy + Krea 2 packs + inventory/Civitai)");
}
static string Clip(string text, int max)
@@ -128,24 +142,231 @@ public class SwarmAssistentExtension : Extension
return new JObject { ["success"] = true, ["packs"] = packs, ["order"] = new JArray(PackNames) };
}
/// <summary>
/// Proxy to Ollama /api/chat (non-stream).
/// <paramref name="raw"/> must include messages (JArray) and optional context_json.
/// </summary>
public async Task<JObject> AssistentChat(Session session, string baseUrl, string model, string pack, bool includeBase, JObject raw)
/// <summary>Server-side LoRA / checkpoint / wildcard inventory (not DOM scrape).</summary>
public async Task<JObject> AssistentListInventory(Session session)
{
string root = NormalizeBaseUrl(baseUrl ?? raw?["base_url"]?.ToString());
string modelName = (model ?? raw?["model"]?.ToString() ?? "").Trim();
if (string.IsNullOrWhiteSpace(modelName))
await Task.CompletedTask;
JArray loras = [];
JArray checkpoints = [];
JArray wildcards = [];
if (Program.T2IModelSets.TryGetValue("LoRA", out T2IModelHandler loraHandler))
{
return new JObject { ["error"] = "model is required" };
}
JArray userMessages = raw?["messages"] as JArray;
if (userMessages is null || userMessages.Count == 0)
{
return new JObject { ["error"] = "messages required" };
foreach (T2IModel model in loraHandler.Models.Values.OrderBy(m => m.Name).Take(MaxLorasInInventory))
{
loras.Add(new JObject
{
["name"] = model.Name,
["title"] = model.Metadata?.Title ?? model.Title ?? model.Name,
["trigger_phrase"] = model.Metadata?.TriggerPhrase,
["architecture"] = model.ModelClass?.ID,
["compat_class"] = model.ModelClass?.CompatClass?.ID,
["hash"] = model.Metadata?.Hash ?? "",
});
}
}
if (Program.T2IModelSets.TryGetValue("Stable-Diffusion", out T2IModelHandler ckptHandler))
{
foreach (T2IModel model in ckptHandler.Models.Values.OrderBy(m => m.Name).Take(60))
{
checkpoints.Add(new JObject
{
["name"] = model.Name,
["title"] = model.Metadata?.Title ?? model.Title ?? model.Name,
["architecture"] = model.ModelClass?.ID,
["compat_class"] = model.ModelClass?.CompatClass?.ID,
});
}
}
try
{
foreach (string name in WildcardsHelper.ListFiles.OrderBy(n => n).Take(MaxWildcardsInInventory))
{
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,
};
}
/// <summary>Search Civitai for LoRAs (prefers Krea 2 base). Uses Swarm-stored civitai_api key.</summary>
public async Task<JObject> 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<string> installedNames = CollectInstalledLoraNames();
HashSet<string> 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<string> CollectInstalledLoraNames()
{
HashSet<string> 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<string> CollectInstalledLoraHashes()
{
HashSet<string> 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<string> installedNames, HashSet<string> 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<string> 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<bool>() == 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<bool>() ?? false,
};
}
List<JObject> BuildOllamaMessages(string packName, bool includeBase, string contextJson, JArray userMessages, string extraSystem = null)
{
List<JObject> ollamaMessages = [];
StringBuilder system = new();
if (includeBase)
@@ -156,7 +377,6 @@ public class SwarmAssistentExtension : Extension
system.AppendLine(basePack);
}
}
string packName = (pack ?? raw?["pack"]?.ToString() ?? "write_prompt").Trim();
if (!string.IsNullOrWhiteSpace(packName) && packName != "base_krea2")
{
string situational = ReadPackFile(packName);
@@ -167,7 +387,6 @@ public class SwarmAssistentExtension : Extension
system.AppendLine(situational);
}
}
string contextJson = raw?["context_json"]?.ToString();
if (!string.IsNullOrWhiteSpace(contextJson))
{
system.AppendLine();
@@ -176,6 +395,11 @@ public class SwarmAssistentExtension : Extension
system.AppendLine(contextJson);
system.AppendLine("```");
}
if (!string.IsNullOrWhiteSpace(extraSystem))
{
system.AppendLine();
system.AppendLine(extraSystem);
}
if (system.Length > 0)
{
ollamaMessages.Add(new JObject
@@ -184,7 +408,7 @@ public class SwarmAssistentExtension : Extension
["content"] = system.ToString(),
});
}
foreach (JToken msg in userMessages)
foreach (JToken msg in userMessages ?? [])
{
if (msg is not JObject mo)
{
@@ -201,24 +425,264 @@ public class SwarmAssistentExtension : Extension
}
ollamaMessages.Add(copy);
}
return ollamaMessages;
}
static JObject TryParsePatch(string reply)
{
if (string.IsNullOrWhiteSpace(reply))
{
return null;
}
foreach (Match match in JsonFenceRe.Matches(reply))
{
string raw = match.Groups[1].Value.Trim();
try
{
JObject obj = JObject.Parse(raw);
if (obj is not null && (obj["prompt"] != null || obj["loras"] != null || obj["width"] != null
|| obj["height"] != null || obj["steps"] != null || obj["cfg"] != null
|| obj["seed"] != null || obj["sigma_shift"] != null || obj["sampler"] != null
|| obj["actions"] != null || obj["search_query"] != null || obj["civitai_query"] != null
|| obj["use_init_image"] != null || obj["clear_init_image"] != null
|| obj["init_creativity"] != null || obj["denoise"] != null
|| obj["use_mask_image"] != null || obj["clear_mask_image"] != null
|| obj["mask_blur"] != null || obj["mask_grow"] != null))
{
return obj;
}
}
catch
{
// not json
}
}
return null;
}
static string ExtractSearchQuery(JObject patch)
{
if (patch is null)
{
return null;
}
string q = (patch["search_query"] ?? patch["civitai_query"])?.ToString()?.Trim();
if (!string.IsNullOrWhiteSpace(q))
{
return q;
}
if (patch["actions"] is JArray acts)
{
foreach (JToken a in acts)
{
if (string.Equals(a?.ToString(), "search_civitai", StringComparison.OrdinalIgnoreCase))
{
return q; // may still be null — caller checks
}
}
}
return null;
}
static bool WantsCivitaiSearch(JObject patch)
{
if (patch is null)
{
return false;
}
if (!string.IsNullOrWhiteSpace(ExtractSearchQuery(patch)))
{
return true;
}
if (patch["actions"] is JArray acts)
{
foreach (JToken a in acts)
{
if (string.Equals(a?.ToString(), "search_civitai", StringComparison.OrdinalIgnoreCase))
{
return true;
}
}
}
return false;
}
async Task<(string reply, JObject raw, JArray civitaiResults)> RunChatWithHops(
Session session,
string root,
string modelName,
string packName,
bool includeBase,
string contextJson,
JArray userMessages,
Func<string, Task> onDelta = null,
Func<int, Task> onHopStart = null)
{
List<JObject> messages = BuildOllamaMessages(packName, includeBase, contextJson, userMessages);
JArray civitaiResults = [];
string reply = "";
JObject lastRaw = null;
for (int hop = 0; hop < MaxCivitaiHops; hop++)
{
if (onHopStart is not null)
{
await onHopStart(hop);
}
(reply, lastRaw) = await CallOllamaChat(root, modelName, messages, stream: onDelta is not null, onDelta);
JObject patch = TryParsePatch(reply);
if (hop + 1 >= MaxCivitaiHops || !WantsCivitaiSearch(patch))
{
break;
}
string query = ExtractSearchQuery(patch);
if (string.IsNullOrWhiteSpace(query))
{
query = userMessages.LastOrDefault(m => m["role"]?.ToString() == "user")?["content"]?.ToString() ?? "";
}
if (string.IsNullOrWhiteSpace(query))
{
break;
}
JObject search = await AssistentSearchCivitai(session, query, 8);
if (search["error"] is not null)
{
messages.Add(new JObject { ["role"] = "assistant", ["content"] = reply });
messages.Add(new JObject
{
["role"] = "user",
["content"] = $"Civitai search failed: {search["error"]}. Continue without download — use only available_loras from context.",
});
continue;
}
civitaiResults = search["results"] as JArray ?? [];
messages.Add(new JObject { ["role"] = "assistant", ["content"] = reply });
messages.Add(new JObject
{
["role"] = "user",
["content"] =
"Civitai search results (JSON). Prefer `krea_likely: true`. Do NOT download yourself — the UI shows Confirm cards. " +
"Pick useful LoRAs from results or available_loras, emit a normal patch (prompt/loras). " +
"Omit search_civitai from actions unless you need a different query.\n```json\n" +
civitaiResults.ToString(Newtonsoft.Json.Formatting.None) + "\n```",
});
}
return (reply, lastRaw, civitaiResults);
}
async Task<(string reply, JObject raw)> CallOllamaChat(
string root,
string modelName,
List<JObject> ollamaMessages,
bool stream,
Func<string, Task> onDelta)
{
JObject payload = new()
{
["model"] = modelName,
["stream"] = false,
["stream"] = stream,
["messages"] = new JArray(ollamaMessages),
};
try
using StringContent content = new(payload.ToString(Newtonsoft.Json.Formatting.None), Encoding.UTF8, "application/json");
using HttpRequestMessage req = new(HttpMethod.Post, $"{root}/api/chat") { Content = content };
using HttpResponseMessage resp = await HttpClient.SendAsync(req, stream
? HttpCompletionOption.ResponseHeadersRead
: HttpCompletionOption.ResponseContentRead);
if (!resp.IsSuccessStatusCode)
{
string errBody = await resp.Content.ReadAsStringAsync();
throw new Exception($"Ollama /api/chat HTTP {(int)resp.StatusCode}: {Clip(errBody, 800)}");
}
if (!stream)
{
using StringContent content = new(payload.ToString(Newtonsoft.Json.Formatting.None), Encoding.UTF8, "application/json");
using HttpResponseMessage resp = await HttpClient.PostAsync($"{root}/api/chat", content);
string body = await resp.Content.ReadAsStringAsync();
if (!resp.IsSuccessStatusCode)
{
return new JObject { ["error"] = $"Ollama /api/chat HTTP {(int)resp.StatusCode}: {Clip(body, 800)}" };
}
JObject parsed = JObject.Parse(body);
string reply = parsed["message"]?["content"]?.ToString() ?? parsed["response"]?.ToString() ?? "";
return (reply, parsed);
}
StringBuilder full = new();
await using Stream streamBody = await resp.Content.ReadAsStreamAsync();
using StreamReader reader = new(streamBody, Encoding.UTF8);
JObject last = null;
while (true)
{
string line = await reader.ReadLineAsync();
if (line is null)
{
break;
}
if (string.IsNullOrWhiteSpace(line))
{
continue;
}
JObject chunk = JObject.Parse(line);
last = chunk;
string delta = chunk["message"]?["content"]?.ToString() ?? "";
if (!string.IsNullOrEmpty(delta))
{
full.Append(delta);
if (onDelta is not null)
{
await onDelta(delta);
}
}
if (chunk["done"]?.Value<bool>() == true)
{
break;
}
}
return (full.ToString(), last ?? new JObject());
}
/// <summary>
/// SwarmUI passes the whole request as the JObject param (not only a nested key).
/// Support both flat fields and legacy nested <c>raw</c>.
/// </summary>
static void ExtractChatPayload(JObject raw, ref string baseUrl, ref string model, ref string pack, ref bool includeBase, out JArray userMessages, out string contextJson)
{
JObject whole = raw ?? [];
JObject nested = whole["raw"] as JObject;
if (string.IsNullOrWhiteSpace(baseUrl))
{
baseUrl = whole["base_url"]?.ToString()
?? whole["baseUrl"]?.ToString()
?? nested?["base_url"]?.ToString()
?? nested?["baseUrl"]?.ToString();
}
if (string.IsNullOrWhiteSpace(model))
{
model = whole["model"]?.ToString() ?? nested?["model"]?.ToString();
}
if (string.IsNullOrWhiteSpace(pack))
{
pack = whole["pack"]?.ToString() ?? nested?["pack"]?.ToString();
}
if (whole["includeBase"] is not null)
{
includeBase = whole.Value<bool?>("includeBase") ?? includeBase;
}
userMessages = (whole["messages"] as JArray) ?? (nested?["messages"] as JArray);
contextJson = whole["context_json"]?.ToString() ?? nested?["context_json"]?.ToString();
}
/// <summary>Proxy to Ollama /api/chat (non-stream), with optional Civitai search hop.</summary>
public async Task<JObject> AssistentChat(Session session, string baseUrl, string model, string pack, bool includeBase, JObject raw)
{
ExtractChatPayload(raw, ref baseUrl, ref model, ref pack, ref includeBase, out JArray userMessages, out string contextJson);
string root = NormalizeBaseUrl(baseUrl);
string modelName = (model ?? "").Trim();
if (string.IsNullOrWhiteSpace(modelName))
{
return new JObject { ["error"] = "model is required" };
}
if (userMessages is null || userMessages.Count == 0)
{
return new JObject { ["error"] = "messages required" };
}
string packName = (pack ?? "write_prompt").Trim();
try
{
(string reply, JObject parsed, JArray civitai) = await RunChatWithHops(
session, root, modelName, packName, includeBase, contextJson, userMessages);
return new JObject
{
["success"] = true,
@@ -226,6 +690,7 @@ public class SwarmAssistentExtension : Extension
["model"] = modelName,
["pack"] = packName,
["raw"] = parsed,
["civitai_results"] = civitai,
};
}
catch (Exception ex)
@@ -233,4 +698,62 @@ public class SwarmAssistentExtension : Extension
return new JObject { ["error"] = $"Ollama chat failed: {ex.Message}" };
}
}
/// <summary>WebSocket streaming chat (Ollama stream:true) + Civitai hops.</summary>
public async Task<JObject> AssistentChatWS(Session session, WebSocket ws, string baseUrl, string model, string pack, bool includeBase, JObject raw)
{
ExtractChatPayload(raw, ref baseUrl, ref model, ref pack, ref includeBase, out JArray userMessages, out string contextJson);
string root = NormalizeBaseUrl(baseUrl);
string modelName = (model ?? "").Trim();
if (string.IsNullOrWhiteSpace(modelName))
{
await ws.SendJson(new JObject { ["error"] = "model is required" }, API.WebsocketTimeout);
return null;
}
if (userMessages is null || userMessages.Count == 0)
{
await ws.SendJson(new JObject { ["error"] = "messages required" }, API.WebsocketTimeout);
return null;
}
string packName = (pack ?? "write_prompt").Trim();
try
{
async Task OnDelta(string delta)
{
if (ws.State == WebSocketState.Open)
{
await ws.SendJson(new JObject { ["delta"] = delta }, API.WebsocketTimeout);
}
}
async Task OnHopStart(int hop)
{
if (ws.State == WebSocketState.Open && hop > 0)
{
await ws.SendJson(new JObject
{
["clear_stream"] = true,
["hop"] = hop + 1,
["notice"] = "Civitai search done — refining…",
}, API.WebsocketTimeout);
}
}
(string reply, JObject parsed, JArray civitai) = await RunChatWithHops(
session, root, modelName, packName, includeBase, contextJson, userMessages, OnDelta, OnHopStart);
await ws.SendJson(new JObject
{
["success"] = true,
["done"] = true,
["reply"] = reply,
["model"] = modelName,
["pack"] = packName,
["raw"] = parsed,
["civitai_results"] = civitai,
}, API.WebsocketTimeout);
}
catch (Exception ex)
{
await ws.SendJson(new JObject { ["error"] = $"Ollama chat failed: {ex.Message}" }, API.WebsocketTimeout);
}
return null;
}
}
+14 -4
View File
@@ -15,6 +15,9 @@
</div>
<div class="sa-image-actions">
<button type="button" class="basic-button" id="sa_btn_use_current" title="Grab the current Generate preview">Use current</button>
<button type="button" class="basic-button" id="sa_btn_as_init" title="Set vision image as Swarm Init Image (img2img)">As Init</button>
<button type="button" class="basic-button" id="sa_btn_as_mask" title="Set vision image as Mask (inpaint — white=edit)">As Mask</button>
<button type="button" class="basic-button" id="sa_btn_clear_init" title="Clear Init Image + Mask on Generate tab">Clear Init</button>
<button type="button" class="basic-button" id="sa_btn_clear_image" title="Clear vision image">Clear</button>
<label class="sa-check"><input type="checkbox" id="sa_attach_vision" checked /> Attach next</label>
</div>
@@ -29,23 +32,30 @@
<option value="critique_image">Critique image</option>
<option value="compose_scene">Compose scene</option>
<option value="fix_params">Fix params</option>
<option value="inpaint_edit">Inpaint / img2img</option>
</select>
<select id="sa_model" class="sa-select sa-model-select" title="Ollama model">
<option value="">Loading models…</option>
</select>
<button type="button" class="basic-button sa-icon-btn" id="sa_btn_settings" title="Settings" aria-label="Settings"></button>
</div>
</header>
<div class="sa-settings" id="sa_settings" hidden>
<label>Ollama URL <input type="text" id="sa_base_url" value="http://127.0.0.1:11434" /></label>
<label>Model
<select id="sa_model" class="sa-select"></select>
</label>
<button type="button" class="basic-button" id="sa_btn_refresh_models">Refresh models</button>
<button type="button" class="basic-button" id="sa_btn_refresh_inventory">Refresh inventory</button>
<label class="sa-check"><input type="checkbox" id="sa_auto_vision" /> Auto-attach current image</label>
<label class="sa-check"><input type="checkbox" id="sa_auto_apply" checked /> Auto-apply patch</label>
<label class="sa-check"><input type="checkbox" id="sa_auto_generate" checked /> Auto-generate after patch</label>
<label class="sa-check"><input type="checkbox" id="sa_auto_critique" /> Auto-critique after generate</label>
<label class="sa-check sa-danger" title="Dangerous — downloads without Confirm"><input type="checkbox" id="sa_auto_download" /> Auto-download Civitai (off)</label>
</div>
<div class="sa-messages" id="sa_messages"></div>
<div class="sa-composer">
<textarea id="sa_input" rows="3" placeholder="Ask for a prompt, critique the image, change aspect… (Enter = send, Shift+Enter = newline)"></textarea>
<textarea id="sa_input" rows="3" placeholder="Ask for a prompt, img2img, inpaint, critique… (Enter = send, Shift+Enter = newline)"></textarea>
<div class="sa-composer-actions">
<button type="button" class="basic-button sa-primary" id="sa_btn_send">Send</button>
<button type="button" class="basic-button" id="sa_btn_interrupt" title="Interrupt generation / chat" hidden>Interrupt</button>
<button type="button" class="basic-button" id="sa_btn_clear">Clear chat</button>
<span class="sa-status" id="sa_status"></span>
</div>