From cfb55eff452b2a4ea4273a9e58c663dee36ab94e Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 17 Oct 2026 15:55:24 +0000 Subject: [PATCH] Training: choose which prompt parts become captions (prefix, main, postfix, variables) Variables expand when ticked, stay as written otherwise. Co-Authored-By: Claude Sonnet 5.5 Claude-Session: https://claude.ai/code/session_01MpnF7GwNGMbngU11wjym7P --- app/ui/features/training/data.js | 16 +++++--- app/ui/features/training/index.js | 4 +- app/ui/features/training/training.test.js | 15 +++++--- app/ui/shell/api.js | 2 +- modules/training/lora/__init__.py | 3 +- modules/training/lora/rated.py | 38 ++++++++++++------- .../training/lora/tests/test_sets_and_tags.py | 17 ++++++--- 7 files changed, 61 insertions(+), 34 deletions(-) diff --git a/app/ui/features/training/data.js b/app/ui/features/training/data.js index 39536c1e..00cc2166 100644 --- a/app/ui/features/training/data.js +++ b/app/ui/features/training/data.js @@ -13,8 +13,10 @@ export const RATED = [ const ALL_RATED = ["liked", "disliked", "bad_composition", "bad_image"]; export const binEntry = (m) => ({ src: { kind: "bin", id: m.id }, media_kind: m.kind === "video" ? "video" : "image", label: "own", name: m.name }); -export const ratedEntry = (r, withCaption = true) => ({ src: { kind: "output", type: r.media.type, subfolder: r.media.subfolder, filename: r.media.filename }, - media_kind: /\.(mp4|webm|mov|mkv|avi)$/i.test(r.media.filename) ? "video" : "image", label: r.label, caption: withCaption ? r.caption || "" : "", name: r.media.filename }); +export const INCLUDE = [{ value: "prefix", label: "Prefix" }, { value: "main", label: "Main prompt" }, { value: "postfix", label: "Postfix" }, { value: "variables", label: "Variables" }]; + +export const ratedEntry = (r) => ({ src: { kind: "output", type: r.media.type, subfolder: r.media.subfolder, filename: r.media.filename }, + media_kind: /\.(mp4|webm|mov|mkv|avi)$/i.test(r.media.filename) ? "video" : "image", label: r.label, caption: r.caption || "", name: r.media.filename }); const countsLine = (counts) => Object.keys(LABEL_WORDS).filter((k) => counts[k]).map((k) => `${counts[k]} ${LABEL_WORDS[k].toLowerCase()}`).join(", "); @@ -74,14 +76,18 @@ export function dataStep({ app, session, go, ui }) { const ratedCount = () => want().reduce((a, k) => a + ((ui.rated.counts || {})[k] || 0), 0); const ratedNote = c.text.sm({ text: "" }); const addRated = c.button.sm({ label: "Add these results", tone: "primary", disabled: true, onClick: async () => { - await added((ui.rated.results || []).filter((r) => want().includes(r.label)).map((r) => ratedEntry(r, ui.useCaptions)), "your results"); + await added((ui.rated.results || []).filter((r) => want().includes(r.label)).map(ratedEntry), "your results"); } }); const showRated = () => { const n = ratedCount(); ratedNote.setText(n ? `${n} rated result${n > 1 ? "s" : ""} found in your projects.` : "None found yet. Rate results in your projects and they appear here."); addRated.node.disabled = !n || !open; }; - const useCaptions = c.toggle.default({ label: "Use their prompts as captions", hint: "The prompt each result was made from, with shortcuts expanded. Variables, prefixes and postfixes are left out.", checked: ui.useCaptions, onChange: (v) => { ui.useCaptions = v; } }); + const include = c.field.default({ label: "Captions include", hint: "What goes into the caption from each result's prompt. Shortcuts are expanded. Variables are filled in only when ticked, otherwise $name stays as written. Nothing ticked adds them with no caption.", + control: c.checklist.default({ label: "Captions include", items: INCLUDE, values: ui.include, onChange: async (v) => { + ui.include = v; + try { ui.rated = await app.api.trainRated(v.length ? v : ["none"]); } catch { /* the old captions stay */ } + } }) }); const ratedSel = c.segmented.sm({ label: "Which results", options: RATED, value: ui.ratedKind, onChange: (v) => { ui.ratedKind = v; showRated(); } }); showRated(); @@ -94,7 +100,7 @@ export function dataStep({ app, session, go, ui }) { status, c.group.default({ label: "Drop files", rows: [dropzone] }), c.group.default({ label: "Choose in your Bin", rows: [gallery, c.toolbar.default({ items: [pickedNote], trailing: [addBin] })] }), - c.group.default({ label: "Choose your own results", hint: "Results you rated in your projects. Liked ones teach what you want; disliked ones can teach what to avoid.", rows: [ratedSel, useCaptions, c.toolbar.default({ items: [ratedNote], trailing: [addRated] })] }), + c.group.default({ label: "Choose your own results", hint: "Results you rated in your projects. Liked ones teach what you want; disliked ones can teach what to avoid.", rows: [ratedSel, include, c.toolbar.default({ items: [ratedNote], trailing: [addRated] })] }), next, ] }); } diff --git a/app/ui/features/training/index.js b/app/ui/features/training/index.js index f721dc8d..90a69add 100644 --- a/app/ui/features/training/index.js +++ b/app/ui/features/training/index.js @@ -25,7 +25,7 @@ export default { const api = app.api; const session = makeSession(api); const ui = { - tab: "data", picked: new Set(), bin: [], rated: { results: [], counts: {} }, ratedKind: "liked", useCaptions: true, selected: new Set(), filter: "all", tag: "", limit: 48, + tab: "data", picked: new Set(), bin: [], rated: { results: [], counts: {} }, ratedKind: "liked", include: ["main"], selected: new Set(), filter: "all", tag: "", limit: 48, tagger: { id: "wd_vit_v3", list: [], job: null, general: 0.35, limit: 0, characters: false, mode: "add", scope: "untagged" }, advanced: false, runs: [], chosenRun: "", record: null, run: { shown: false, active: false, state: "idle", loss: [], smooth: [], step: 0, total: 0, elapsed: 0, message: "", error: "", lora: "", before: "" }, @@ -65,7 +65,7 @@ export default { // --- the Bin and rated results, read when the window opens async function loadSources() { try { ui.bin = (await api.media()).media || []; } catch { ui.bin = []; } - try { ui.rated = await api.trainRated(); } catch { ui.rated = { results: [], counts: {} }; } + try { ui.rated = await api.trainRated(ui.include); } catch { ui.rated = { results: [], counts: {} }; } draw(); } async function loadRuns() { try { ui.runs = await api.trainRuns(); } catch { ui.runs = []; } if (ui.runs.length) loadRecord(ui.chosenRun || ui.runs[0].run); } diff --git a/app/ui/features/training/training.test.js b/app/ui/features/training/training.test.js index 3eba560e..6c5de582 100644 --- a/app/ui/features/training/training.test.js +++ b/app/ui/features/training/training.test.js @@ -154,23 +154,26 @@ test("a run in progress draws its loss, then reports the saved file; Stop asks t off(); }); -test("rated results are added with their prompts as captions, or without when that is switched off", async () => { +test("rated results are added with the captions the server built, and changing what to include asks again", async () => { const f = fake(); - f.api.trainRated = async () => ({ results: [{ label: "liked", media: { type: "output", subfolder: "", filename: "x.mp4" }, caption: "a red fox" }], counts: { liked: 1 } }); + const asked = []; + f.api.trainRated = async (include) => { asked.push(include); return { results: [{ label: "liked", media: { type: "output", subfolder: "", filename: "x.mp4" }, caption: include.includes("postfix") ? "a red fox 35mm" : "a red fox" }], counts: { liked: 1 } }; }; const added = []; f.api.trainAdd = async (name, entries) => { added.push(entries); return { set: f.doc, added: entries.length }; }; const { button, off } = await open(f); button.click(); await wait(60); + assert.deepEqual(asked[0], ["main"], "the main prompt alone is the default"); click("Add these results").click(); await wait(40); assert.equal(added[0][0].caption, "a red fox"); assert.equal(added[0][0].media_kind, "video"); - const sw = [...document.querySelectorAll("input[role=switch]")].find((i) => /prompts as captions/.test(i.closest(".cx-toggle-row").textContent)); - sw.click(); - await wait(10); + const box = [...document.querySelectorAll(".cx-checklist .cx-check-row")].find((r) => r.textContent.includes("Postfix")).querySelector("input"); + box.click(); + await wait(40); + assert.deepEqual(asked[asked.length - 1].sort(), ["main", "postfix"]); click("Add these results").click(); await wait(40); - assert.equal(added[1][0].caption, ""); + assert.equal(added[1][0].caption, "a red fox 35mm"); off(); }); diff --git a/app/ui/shell/api.js b/app/ui/shell/api.js index 9ecc4b71..f8d44cbe 100644 --- a/app/ui/shell/api.js +++ b/app/ui/shell/api.js @@ -87,7 +87,7 @@ export const api = { trainDeleteSet: (name) => call("DELETE", `${TRAIN}/sets/${encodeURIComponent(name)}`), trainAdd: (name, entries) => call("POST", `${TRAIN}/sets/${encodeURIComponent(name)}/add`, { entries }), trainBulk: (name, body) => call("POST", `${TRAIN}/sets/${encodeURIComponent(name)}/bulk`, body), - trainRated: (label) => call("GET", `${TRAIN}/rated${label ? `?label=${encodeURIComponent(label)}` : ""}`), + trainRated: (include) => call("GET", `${TRAIN}/rated${include && include.length ? `?include=${encodeURIComponent(include.join(","))}` : ""}`), trainTaggers: () => call("GET", `${TRAIN}/taggers`), trainDownloadTagger: (tagger) => call("POST", `${TRAIN}/taggers/download`, { tagger }), trainRemoveTagger: (tagger) => call("POST", `${TRAIN}/taggers/remove`, { tagger }), diff --git a/modules/training/lora/__init__.py b/modules/training/lora/__init__.py index ab27af54..7de3311e 100644 --- a/modules/training/lora/__init__.py +++ b/modules/training/lora/__init__.py @@ -107,7 +107,8 @@ def routes(table, base, web): project = projects.get(entry["id"]) if project is not None: docs.append({**project.to_dict(), "id": entry["id"]}) - found = rated.scan(docs, lambda t: shortcuts.expand(t, seed=1)) + parts = [p for p in (req.rel_url.query.get("include") or "main").split(",") if p in rated.INCLUDE] + found = rated.scan(docs, lambda t: shortcuts.expand(t, seed=1), shortcuts.resolve_variables, parts) want = req.rel_url.query.get("label") if want: found = [r for r in found if r["label"] == want] diff --git a/modules/training/lora/rated.py b/modules/training/lora/rated.py index e72b731d..f95940d0 100644 --- a/modules/training/lora/rated.py +++ b/modules/training/lora/rated.py @@ -8,8 +8,6 @@ one disliked, and "Disliked: bad image" / "Disliked: bad composition" say which from __future__ import annotations -import re - from .store import kind_of FORGET = "-Just forget it-" @@ -39,25 +37,37 @@ def _media(raw): return {"filename": name, "subfolder": str(raw.get("subfolder") or ""), "type": raw.get("type") if raw.get("type") in ("output", "temp", "input") else "output"} -_VARIABLE = re.compile(r"\$[A-Za-z_][A-Za-z0-9_]*") +INCLUDE = ("prefix", "main", "postfix", "variables") -def prompt_caption(text, expand=None) -> str: - """A rated result's prompt as a caption: the scene's own words with shortcuts expanded, and nothing else. +def prompt_caption(text, doc=None, include=("main",), expand=None, resolve=None) -> str: + """A rated result's prompt as a caption, from the parts asked for. - `expand(text)` is the app's shortcut expansion. `$variables` are dropped, not filled in, and the project's - anchor and postfix were never part of the stored text, so they are not here either. + `include` names the parts: prefix and postfix are the project's anchor and postfix, main is the scene's own words, + variables fills in `$name` (left as written when not asked for). `expand(text)` is the app's shortcut + expansion; `resolve(text, variables)` fills the variables in. """ - text = " ".join(str(text or "").split()) - if expand and text: - text = " ".join(str(expand(text)).split()) - return " ".join(_VARIABLE.sub(" ", text).split()) + doc = doc if isinstance(doc, dict) else {} + want = set(include or ()) + parts = [] + if "prefix" in want: + parts.append(doc.get("anchor")) + if "main" in want: + parts.append(text) + if "postfix" in want and doc.get("postfix_enabled") is not False: + parts.append(doc.get("postfix")) + joined = " ".join(" ".join(str(p or "").split()) for p in parts).strip() + if expand and joined: + joined = " ".join(str(expand(joined)).split()) + if "variables" in want and resolve: + joined = resolve(joined, doc.get("variables")) + return " ".join(joined.split()) -def scan(projects, expand=None) -> list[dict]: +def scan(projects, expand=None, resolve=None, include=("main",)) -> list[dict]: """[{project, project_name, label, rating, media, caption}] for every rated result in these project documents. - The caption is the prompt that made the result, run through `prompt_caption` with `expand`. + The caption is the prompt that made the result, built by `prompt_caption` from the parts in `include`. """ out, seen = [], set() @@ -70,7 +80,7 @@ def scan(projects, expand=None) -> list[dict]: return seen.add(key) out.append({"project": doc.get("id", ""), "project_name": doc.get("name", ""), "label": label, "rating": str(rating), - "media": media, "caption": prompt_caption(caption, expand)}) + "media": media, "caption": prompt_caption(caption, doc, include, expand, resolve)}) for doc in projects: if not isinstance(doc, dict): diff --git a/modules/training/lora/tests/test_sets_and_tags.py b/modules/training/lora/tests/test_sets_and_tags.py index 4e45a018..a3c67a76 100644 --- a/modules/training/lora/tests/test_sets_and_tags.py +++ b/modules/training/lora/tests/test_sets_and_tags.py @@ -89,9 +89,16 @@ def test_buckets_share_one_shape_or_follow_the_picture(): assert len(set(media_io.buckets_for([1.0, 2.0, 0.5], 512, 64, False))) == 3 -def test_a_rated_prompt_becomes_a_caption_with_shortcuts_expanded_and_variables_dropped(): - assert rated.prompt_caption(" a fox\n runs $hero home ", lambda t: t.replace("fox", "red fox")) == "a red fox runs home" - assert rated.prompt_caption("", lambda t: "x") == "" and rated.prompt_caption(None) == "" +def test_a_rated_prompt_becomes_a_caption_from_the_parts_asked_for(): + doc = {"anchor": "cinematic", "postfix": "35mm", "postfix_enabled": True, "variables": [{"name": "hero", "value": "Ann"}]} + expand = lambda t: t.replace("fox", "red fox") + resolve = lambda t, v: t.replace("$hero", "Ann") + ask = lambda *include: rated.prompt_caption(" a fox\n runs $hero home ", doc, include, expand, resolve) + assert ask("main") == "a red fox runs $hero home", "variables stay as written unless asked for" + assert ask("main", "variables") == "a red fox runs Ann home" + assert ask("prefix", "main", "postfix", "variables") == "cinematic a red fox runs Ann home 35mm" + assert ask() == "" + assert rated.prompt_caption("x", {**doc, "postfix_enabled": False}, ("main", "postfix")) == "x" media = {"filename": "a.mp4", "subfolder": "", "type": "output"} - doc = {"id": "p", "scenes": [{"id": "s", "rating": "10", "rated": {"text": "a $hero fox"}}], "scene_renders": {"s": {"media": media}}} - assert rated.scan([doc], lambda t: t.replace("fox", "red fox"))[0]["caption"] == "a red fox" + project = {"id": "p", "anchor": "A", "scenes": [{"id": "s", "rating": "10", "rated": {"text": "a $hero fox"}}], "scene_renders": {"s": {"media": media}}} + assert rated.scan([project], expand, resolve, ("prefix", "main"))[0]["caption"] == "A a $hero red fox"