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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01MpnF7GwNGMbngU11wjym7P

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,
   ] });
 }

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); }

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();
 });

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 }),

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]

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):

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"