diff --git a/app/ui/composer/elements/captioned.js b/app/ui/composer/elements/captioned.js
index 24aa4208..7dc28293 100644
--- a/app/ui/composer/elements/captioned.js
+++ b/app/ui/composer/elements/captioned.js
@@ -1,26 +1,16 @@
// A picture or clip with its caption beside it: the preview and the tag editor in one card.
//
-// The card shows what it is given and reports what a person edits. The caption is a comma separated
-// line; the editor inside owns the chips, and onChange hears the whole line after an edit.
+// The card shows what it is given and reports what a person edits. The caption is free text in a plain
+// box, kept exactly as typed; onChange hears the whole text as it changes.
import { define, lookup } from "../internals/register.js";
import { el } from "../internals/el.js";
import { setText } from "../internals/text.js";
-import { clean } from "./tags.js";
const LAYOUTS = ["cell", "side"];
-const tagsOf = (text) => {
- const seen = new Set(), out = [];
- for (const part of String(text ?? "").split(/[,\n]/)) {
- const tag = clean(part);
- if (tag && !seen.has(tag.toLowerCase())) { seen.add(tag.toLowerCase()); out.push(tag); }
- }
- return out;
-};
-
for (const layout of LAYOUTS) {
- define("captioned", layout, ({ src, kind = "image", title = "", caption = "", badge, tone = "neutral", selected = false, suggestions = [], onChange, onSelect, onOpen, label } = {}) => {
+ define("captioned", layout, ({ src, kind = "image", title = "", caption = "", badge, tone = "neutral", selected = false, onChange, onSelect, onOpen, label } = {}) => {
const glyph = () => el("span", { cls: "cx-captioned-glyph", text: kind === "video" ? "▶" : "▨", attrs: { "aria-hidden": "true" } });
const view = el("button", { cls: ["cx-captioned-view", "cx-focusable"], attrs: { type: "button", "aria-label": title ? `Open ${title}` : "Open" } });
const show = (next, nextKind) => {
@@ -41,20 +31,20 @@ for (const layout of LAYOUTS) {
const pick = onSelect ? el("input", { cls: ["cx-captioned-pick", "cx-focusable"], attrs: { type: "checkbox", "aria-label": title ? `Select ${title}` : "Select" } }) : null;
if (pick) { pick.checked = !!selected; pick.addEventListener("change", () => { node.classList.toggle("cx-on", pick.checked); onSelect(pick.checked); }); }
- const editor = lookup("tags", "editor")({ value: tagsOf(caption), suggestions, placeholder: "Add a tag", label: title ? `Tags for ${title}` : "Tags",
- onChange: (list) => { if (onChange) onChange(list.join(", ")); } });
+ const editor = lookup("textarea", "md")({ value: String(caption ?? ""), rows: 3, autoGrow: true, placeholder: "Caption", label: title ? `Caption for ${title}` : "Caption",
+ onInput: (text) => { if (onChange) onChange(text); } });
const head = el("div", { cls: "cx-captioned-head", children: [pick, name, chip].filter(Boolean) });
const body = el("div", { cls: "cx-captioned-body", children: [head, editor.node] });
const node = el("div", { cls: ["cx-captioned", `cx-captioned-${layout}`, selected ? "cx-on" : ""].filter(Boolean), attrs: { role: "group", "aria-label": label || title || "Captioned media" }, children: [view, body] });
return {
node,
- getCaption: () => editor.get().join(", "),
- setCaption: (text) => editor.setValue(tagsOf(text)),
+ getCaption: () => editor.value,
+ setCaption: (text) => editor.setValue(String(text ?? "")),
setSelected(on) { if (pick) pick.checked = !!on; node.classList.toggle("cx-on", !!on); },
setBadge(text, nextTone) { setText(chip, text ?? ""); chip.hidden = !text; if (nextTone) chip.className = `cx-captioned-badge cx-captioned-${nextTone}`; },
setPicture: show,
- focus: () => editor.focus(),
+ focus: () => editor.node.focus(),
destroy() { editor.destroy(); node.remove(); },
};
});
diff --git a/app/ui/composer/tests/chart_tags.test.js b/app/ui/composer/tests/chart_tags.test.js
index a1e79dcb..14d52f0d 100644
--- a/app/ui/composer/tests/chart_tags.test.js
+++ b/app/ui/composer/tests/chart_tags.test.js
@@ -242,15 +242,17 @@ test("a cloud tag is shown as text, never as markup", () => {
// --- captioned ---------------------------------------------------------------
-test("a captioned card shows its tags as chips and hands the whole line back after an edit", () => {
+test("a captioned card shows its caption in a text box as written and hands the whole text back as it is typed", () => {
const heard = [];
- const c = mount(composer.captioned.cell({ title: "a.png", caption: "red, blue, red", onChange: (v) => heard.push(v) }));
- assert.deepEqual([...c.node.querySelectorAll(".cx-tag-text")].map((n) => n.textContent), ["red", "blue"]);
- c.node.querySelector(".cx-tag-x").click();
- assert.deepEqual(heard, ["blue"]);
+ const c = mount(composer.captioned.cell({ title: "a.png", caption: "A fox runs.\nred, blue, red", onChange: (v) => heard.push(v) }));
+ const box = c.node.querySelector("textarea");
+ assert.equal(box.value, "A fox runs.\nred, blue, red");
+ box.value = "A fox sleeps, red";
+ box.dispatchEvent(new window.Event("input", { bubbles: true }));
+ assert.deepEqual(heard, ["A fox sleeps, red"]);
c.setCaption("one, two");
assert.equal(c.getCaption(), "one, two");
- assert.deepEqual(heard, ["blue"], "setCaption does not call onChange");
+ assert.deepEqual(heard, ["A fox sleeps, red"], "setCaption does not call onChange");
});
test("a captioned card takes its picture and caption as text, selects, opens and marks a clip", () => {
@@ -258,7 +260,7 @@ test("a captioned card takes its picture and caption as text, selects, opens and
const c = mount(composer.captioned.side({ title: "x", caption: "t", kind: "video", src: "/thumb.png", badge: "Clip", tone: "good",
onSelect: (on) => events.push(["select", on]), onOpen: () => events.push(["open"]) }));
assert.equal(c.node.querySelector(".cx-captioned-title").textContent, "x");
- assert.equal(c.node.querySelector(".cx-tag-text").textContent, "t");
+ assert.equal(c.node.querySelector("textarea").value, "t");
assert.ok(c.node.querySelector(".cx-captioned-play"), "a clip is marked");
assert.equal(c.node.querySelector(".cx-captioned-badge").textContent, "Clip");
c.node.querySelector(".cx-captioned-view").click();
diff --git a/app/ui/features/training/captions.js b/app/ui/features/training/captions.js
index 3e40a536..a5453944 100644
--- a/app/ui/features/training/captions.js
+++ b/app/ui/features/training/captions.js
@@ -10,6 +10,18 @@ export const splitTags = (text) => {
};
export const joinTags = (tags) => splitTags((tags || []).join(", ")).join(", ").slice(0, 4000);
+const same = (a, b) => String(a).split(/\s+/).filter(Boolean).join(" ").toLowerCase() === String(b).split(/\s+/).filter(Boolean).join(" ").toLowerCase();
+/** The comma separated pieces of a caption as typed: trimmed, no empties, nothing dropped or reworded. */
+export const parts = (text) => String(text ?? "").split(",").map((p) => p.trim()).filter(Boolean);
+
+/** The caption with these tags added at its end after a comma, skipping any it has; the caption itself stays as typed. */
+export function addTags(caption, tags) {
+ const have = parts(caption);
+ const fresh = splitTags((tags || []).join(", ")).filter((t) => !have.some((h) => same(t, h)));
+ const text = String(caption ?? "").trimEnd().replace(/,+$/, "").trimEnd();
+ return (text ? [text, ...fresh] : fresh).join(", ").slice(0, 4000);
+}
+
/** Every tag with how many captions hold it, most common first. */
export function frequency(items) {
const counts = new Map();
@@ -24,20 +36,19 @@ export function frequency(items) {
export function bulk(items, op, tag, to, ids) {
tag = String(tag ?? "").split(/\s+/).filter(Boolean).join(" ");
if (!tag || !["add", "prepend", "remove", "replace"].includes(op)) return 0;
- const low = tag.toLowerCase();
to = String(to ?? "").split(/\s+/).filter(Boolean).join(" ");
if (op === "replace" && !to) op = "remove";
let changed = 0;
for (const item of items) {
if (ids && !ids.has(item.id)) continue;
- let tags = splitTags(item.caption);
- const has = tags.some((t) => t.toLowerCase() === low);
+ let tags = parts(item.caption);
+ const has = tags.some((t) => same(t, tag));
if (op === "add" && !has) tags.push(tag);
- else if (op === "prepend") tags = [tag, ...tags.filter((t) => t.toLowerCase() !== low)];
- else if (op === "remove" && has) tags = tags.filter((t) => t.toLowerCase() !== low);
- else if (op === "replace" && has) tags = tags.map((t) => (t.toLowerCase() === low ? to : t));
+ else if (op === "prepend") tags = [tag, ...tags.filter((t) => !same(t, tag))];
+ else if (op === "remove" && has) tags = tags.filter((t) => !same(t, tag));
+ else if (op === "replace" && has) tags = tags.map((t) => (same(t, tag) ? to : t));
else continue;
- const text = joinTags(tags);
+ const text = tags.join(", ").slice(0, 4000);
if (text !== item.caption) { item.caption = text; changed += 1; }
}
return changed;
diff --git a/app/ui/features/training/index.js b/app/ui/features/training/index.js
index 3a3f4b74..fdbb7c50 100644
--- a/app/ui/features/training/index.js
+++ b/app/ui/features/training/index.js
@@ -101,6 +101,7 @@ export default {
await session.flush();
const settings = settingsOf(doc);
settings.trigger = doc.trigger || "";
+ settings.postfix = doc.postfix || "";
setRun({ shown: true, active: true, state: "queued", loss: [], smooth: [], step: 0, total: settings.steps, elapsed: 0, message: "Checking the pipeline", error: "", lora: "" });
draw();
let before = "";
diff --git a/app/ui/features/training/tags.js b/app/ui/features/training/tags.js
index 0e3ff252..f3f26a7f 100644
--- a/app/ui/features/training/tags.js
+++ b/app/ui/features/training/tags.js
@@ -1,7 +1,7 @@
// Step 2, Tags: every item with its caption beside its picture, bulk edits, a tag cloud, and automatic tagging.
import { composer as c } from "../../composer/composer.js";
import { LABEL_TONES, LABEL_WORDS, stillOf } from "./state.js";
-import { bulk, frequency, joinTags, splitTags } from "./captions.js";
+import { addTags, bulk, frequency, joinTags, splitTags } from "./captions.js";
const PAGE = 48;
const FILTERS = [
@@ -28,7 +28,7 @@ export function tagsStep({ app, session, go, ui, peek }) {
const cards = list.slice(0, ui.limit).map((item) => {
const card = c.captioned.cell({ title: item.name || item.src.filename || item.src.id, src: stillOf(item), kind: item.media_kind, caption: item.caption,
- badge: LABEL_WORDS[item.label], tone: LABEL_TONES[item.label], selected: sel.has(item.id), suggestions: cloud.map((x) => x.tag).slice(0, 400),
+ badge: LABEL_WORDS[item.label], tone: LABEL_TONES[item.label], selected: sel.has(item.id),
onChange: (text) => session.edit(() => { item.caption = text; }),
onSelect: (on) => { if (on) sel.add(item.id); else sel.delete(item.id); sync(); },
onOpen: () => peek(item) });
@@ -67,12 +67,15 @@ export function tagsStep({ app, session, go, ui, peek }) {
ui.redraw();
} });
- // --- trigger word
- const trigger = c.input.md({ label: "Trigger word", placeholder: "none", value: doc.trigger || "",
- onCommit: (v) => session.edit((d) => { d.trigger = v.split(/\s+/).filter(Boolean).join(" ").slice(0, 80); }) });
+ // --- text put before and after every caption when training (the set's own; the project's prefix and postfix are not used)
+ const phrase = (key, label, hint) => c.field.default({ label, hint, control: c.input.md({ label, placeholder: "none", value: doc[key] || "",
+ onCommit: (v) => session.edit((d) => { d[key] = v.split(/\s+/).filter(Boolean).join(" ").slice(0, 80); }) }) });
return c.region.stack({ gap: "md", children: [
- c.field.default({ label: "Trigger word", hint: "Put first in every caption when training, so you can call the LoRA up later by writing it. Optional.", control: trigger }),
+ c.field.row({ fields: [
+ phrase("trigger", "Prefix", "Put before every caption when training, for example a word to call the LoRA up later. Optional."),
+ phrase("postfix", "Postfix", "Put after every caption when training. Optional."),
+ ] }),
c.group.default({ label: "Automatic tagging", rows: [autoTag({ app, session, ui })] }),
c.group.default({ label: "Tags in this set", hint: "Click one to show only the items that have it. The x removes it from all.", rows: [cloudEl] }),
c.toolbar.default({ items: [filter, count], trailing: [everyone, none] }),
@@ -141,7 +144,7 @@ export function applyTagging(session, job, mode) {
for (const item of d.items) {
const found = (job.results || {})[item.id];
if (!found) continue;
- const text = joinTags(mode === "replace" ? found : [...splitTags(item.caption), ...found]);
+ const text = mode === "replace" ? joinTags(found) : addTags(item.caption, found);
if (text !== item.caption) { item.caption = text; n += 1; }
}
});
diff --git a/app/ui/features/training/train.js b/app/ui/features/training/train.js
index 8720b0e1..9a38acd7 100644
--- a/app/ui/features/training/train.js
+++ b/app/ui/features/training/train.js
@@ -71,7 +71,7 @@ export function trainStep({ app, session, go, ui }) {
stop.node.hidden = !busy;
return c.region.stack({ gap: "md", children: [
- c.hint.default({ text: used.length ? `${used.length} item${used.length > 1 ? "s" : ""} will be learned from (${caption} with tags). ${used.length - caption ? "Items without tags are learned from with only the trigger word." : ""}`.trim()
+ c.hint.default({ text: used.length ? `${used.length} item${used.length > 1 ? "s" : ""} will be learned from (${caption} with tags). ${used.length - caption ? "Items without a caption are learned from with only the prefix and postfix." : ""}`.trim()
: `Nothing in "${session.name}" fits "${(USES.find((u) => u.value === v.use) || {}).label}". Add media or change the choice below.` }),
hasAudioVae ? c.banner.info({ text: "This pipeline has an audio VAE, so it trains picture and sound together (H3). One clip at a time, so training is slow and needs a lot of memory." }) : null,
c.group.default({ label: "Basics", rows: [name, preset,
diff --git a/app/ui/features/training/training.test.js b/app/ui/features/training/training.test.js
index 30d9bfd6..6bd5e04f 100644
--- a/app/ui/features/training/training.test.js
+++ b/app/ui/features/training/training.test.js
@@ -1,7 +1,7 @@
import test from "node:test";
import assert from "node:assert/strict";
import { setupDom, teardownDom } from "../../composer/tests/_dom.js";
-import { bulk, frequency, joinTags, splitTags } from "./captions.js";
+import { addTags, bulk, frequency, joinTags, splitTags } from "./captions.js";
let feature;
test.before(async () => { setupDom(); feature = (await import("./index.js")).default; });
@@ -59,6 +59,12 @@ test("captions: tags are trimmed and unique, counts are per caption, bulk edits
assert.equal(bulk(items, "prepend", "zq"), 2);
assert.equal(items[1].caption, "zq, lion");
assert.equal(bulk(items, "add", ""), 0);
+ const prose = [{ id: "p", caption: "A fox runs.\nIt jumps, twice, home" }];
+ assert.equal(bulk(prose, "add", "sky"), 1);
+ assert.equal(prose[0].caption, "A fox runs.\nIt jumps, twice, home, sky", "only commas are touched: the words stay as written");
+ assert.equal(addTags("a fox, running home.", ["A Fox", "grass", "sky"]), "a fox, running home., grass, sky");
+ assert.equal(addTags("", ["x", "y"]), "x, y");
+ assert.equal(addTags("a long prompt,", ["x"]), "a long prompt, x");
});
test("the window is hidden while the add-on is off", async () => {
@@ -81,7 +87,7 @@ test("the Data step lists the sets, offers the three sources and can be skipped"
off();
});
-test("Tags step: a card per item; editing a caption saves the set", async () => {
+test("Tags step: a text box per item under its picture; typing saves the set and keeps the text as typed", async () => {
const f = fake();
const { button, off } = await open(f);
button.click();
@@ -89,10 +95,13 @@ test("Tags step: a card per item; editing a caption saves the set", async () =>
click("Next: tags").click();
await wait(20);
assert.equal(document.querySelectorAll(".cx-captioned").length, 2);
- document.querySelector(".cx-captioned .cx-tag-x").click(); // takes "cat" off the first item
+ const box = document.querySelector(".cx-captioned textarea");
+ assert.equal(box.value, "cat, grey");
+ box.value = "A cat\nsits, calmly";
+ box.dispatchEvent(new window.Event("input", { bubbles: true }));
await wait(700);
const saved = f.calls.filter((c) => c[0] === "save").pop();
- assert.equal(saved[2].items[0].caption, "grey");
+ assert.equal(saved[2].items[0].caption, "A cat\nsits, calmly");
off();
});
@@ -131,6 +140,33 @@ test("Train step: More settings stays open while an option in it is changed", as
off();
});
+test("Tags step: the prefix and postfix are the set's own, saved with it and sent to the run", async () => {
+ const f = fake();
+ const { button, off } = await open(f);
+ button.click();
+ await wait(60);
+ click("Next: tags").click();
+ await wait(20);
+ const field = (label) => [...document.querySelectorAll(".cx-field")].find((x) => x.textContent.startsWith(label)).querySelector("input");
+ for (const [label, text] of [["Prefix", "zq"], ["Postfix", "end"]]) {
+ const input = field(label);
+ input.value = text;
+ input.dispatchEvent(new window.Event("change", { bubbles: true }));
+ }
+ await wait(700);
+ const saved = f.calls.filter((c) => c[0] === "save").pop();
+ assert.equal(saved[2].trigger, "zq");
+ assert.equal(saved[2].postfix, "end");
+ click("Next: train").click();
+ await wait(20);
+ click("Start training").click();
+ await wait(40);
+ const plan = f.calls.find((c) => c[0] === "plan");
+ assert.equal(plan[2].trigger, "zq");
+ assert.equal(plan[2].postfix, "end");
+ off();
+});
+
test("a refused plan shows the reason and starts nothing", async () => {
const f = fake();
f.api.trainPlan = async () => ({ prompt: {}, problems: ["the pipeline has no clip loader"] });
diff --git a/core/tests/test_shipped_nodes.py b/core/tests/test_shipped_nodes.py
index 57dbc97b..de75edf5 100644
--- a/core/tests/test_shipped_nodes.py
+++ b/core/tests/test_shipped_nodes.py
@@ -74,7 +74,7 @@ FROZEN = {
"model", "clip", "vae", "audio_vae", "set_name", "name", "use", "steps", "learning_rate", "rank", "algorithm",
"optimizer", "loss", "batch_size", "grad_accumulation", "resolution", "aspect_buckets", "frames", "fps", "start",
"trigger", "seed", "training_dtype", "lora_dtype", "gradient_checkpointing", "offloading", "bypass",
- "model_name", "weight_dtype", "compute_dtype",
+ "model_name", "weight_dtype", "compute_dtype", "postfix",
],
"FunPackEmptyLatent": ["model", "width", "height", "length", "batch_size", "audio_vae", "frame_rate"],
"FunPackDiffusionModelLoader": [
diff --git a/docs/ui-kit-reference.md b/docs/ui-kit-reference.md
index ef8f858b..3ffe239c 100644
--- a/docs/ui-kit-reference.md
+++ b/docs/ui-kit-reference.md
@@ -152,7 +152,7 @@ All overlays use the portal, the z-ladder and the dismiss stack.
| `chart` | `line` | `series` (`{ name, values }` or `{ name, points: [{x, y}] }`, plus `tone` `accent`/`good`/`warn`/`muted`, `faint` for a raw noisy line under a smoothed one, `dashed`), `xLabel`, `yLabel`, `height` (`sm`, `md`, `lg`), `logY`, `empty`, `label`. Points that are not numbers are skipped; long series are cut to a few hundred drawn points keeping highs and lows. Hover reads the nearest values. Methods `setSeries`, `setEmpty`. |
| `tags` | `editor` | `value` (list of strings), `suggestions`, `placeholder`, `label`, `maxTags`, `onChange(list)`. Enter or comma adds, Backspace on an empty box removes the last, chips drag to reorder, Tab takes the first suggestion. Methods `get`, `setValue` (does not call `onChange`), `focus`. |
| `tags` | `cloud` | `items` (`{ tag, count }`), `selected`, `onPick(tag)`, `onRemove(tag)` (adds a remove button), `label`, `empty`. Draws the first 300 and says how many more. Methods `setItems`, `setSelected`. |
-| `captioned` | `cell`, `side` | A picture or clip with its caption: preview, name, a badge and a tag editor in one card. `cell` stacks the picture over the tags (for grids), `side` puts them side by side (for one at a time). `src` (a picture URL; a clip passes its still), `kind` (`image` or `video`, a clip gets a play mark), `title`, `caption` (comma separated), `badge` and `tone` (`neutral`, `good`, `bad`, `warn`), `selected` and `onSelect(on)` (adds a tick box), `suggestions`, `onChange(caption)`, `onOpen()` (makes the picture a button). Methods `getCaption`, `setCaption` (does not call `onChange`), `setSelected`, `setBadge`, `setPicture`, `focus`. |
+| `captioned` | `cell`, `side` | A picture or clip with its caption: preview, name, a badge and a text box for the caption in one card. `cell` stacks the picture over the caption (for grids), `side` puts them side by side (for one at a time). `src` (a picture URL; a clip passes its still), `kind` (`image` or `video`, a clip gets a play mark), `title`, `caption` (free text, kept as typed), `badge` and `tone` (`neutral`, `good`, `bad`, `warn`), `selected` and `onSelect(on)` (adds a tick box), `onChange(caption)` (as it is typed), `onOpen()` (makes the picture a button). Methods `getCaption`, `setCaption` (does not call `onChange`), `setSelected`, `setBadge`, `setPicture`, `focus`. |
### Media and time
diff --git a/modules/training/lora/captions.py b/modules/training/lora/captions.py
index 96f0e576..f218faaf 100644
--- a/modules/training/lora/captions.py
+++ b/modules/training/lora/captions.py
@@ -6,7 +6,7 @@ MAX_CAPTION = 4000
def split_tags(text) -> list[str]:
- """"a, b ,, a" -> ["a", "b"]: trimmed, no empties, first occurrence kept."""
+ """"a, b ,, a" -> ["a", "b"]: split at commas, spaces collapsed, no empties, first occurrence kept. For counting, not for rewriting."""
seen, out = set(), []
for part in str(text or "").replace("\n", ",").split(","):
tag = " ".join(part.split())
@@ -20,13 +20,27 @@ def join_tags(tags) -> str:
return ", ".join(split_tags(", ".join(str(t) for t in tags)))[:MAX_CAPTION]
-def with_trigger(caption, trigger) -> str:
- """The trigger word first, once. An empty trigger changes nothing."""
- trigger = " ".join(str(trigger or "").split())
- if not trigger:
- return str(caption or "")
- rest = [t for t in split_tags(caption) if t.lower() != trigger.lower()]
- return ", ".join([trigger, *rest])
+def parts(text) -> list[str]:
+ """The comma separated pieces of a caption as typed: trimmed, no empties, nothing dropped or reworded."""
+ return [p.strip() for p in str(text or "").split(",") if p.strip()]
+
+
+def _same(a, b) -> bool:
+ return " ".join(a.split()).lower() == " ".join(b.split()).lower()
+
+
+def with_parts(caption, prefix="", postfix="") -> str:
+ """The caption between the prefix and the postfix, joined with commas. Empty parts are left out; the caption is not reworded."""
+ pieces = [" ".join(str(x or "").split()) for x in (prefix,)] + [str(caption or "").strip()] + [" ".join(str(postfix or "").split())]
+ return ", ".join(p.strip(", ") for p in pieces if p.strip(", "))
+
+
+def add_tags(caption, tags) -> str:
+ """The caption with these tags added at its end, after a comma, skipping any it already has. The caption itself is kept as typed."""
+ have = parts(caption)
+ new = [t for t in split_tags(", ".join(str(t) for t in tags)) if not any(_same(t, h) for h in have)]
+ text = str(caption or "").rstrip().rstrip(",").rstrip()
+ return (", ".join([text, *new]) if text else ", ".join(new))[:MAX_CAPTION]
def frequency(captions) -> list[tuple[str, int]]:
@@ -42,6 +56,7 @@ def bulk(items, op, tag, new=None, ids=None) -> int:
"""Change captions in place; returns how many items changed.
`op` is add (at the end), prepend, remove or replace (tag -> new). `ids` limits it to those items.
+ Only the comma separated piece that matches is touched; the rest of a caption keeps its words.
"""
changed = 0
tag = " ".join(str(tag or "").split())
@@ -54,21 +69,20 @@ def bulk(items, op, tag, new=None, ids=None) -> int:
for item in items:
if ids is not None and item.get("id") not in ids:
continue
- tags = split_tags(item.get("caption"))
- low = [t.lower() for t in tags]
- if op == "add" and tag.lower() not in low:
- tags.append(tag)
+ pieces = parts(item.get("caption"))
+ has = any(_same(p, tag) for p in pieces)
+ if op == "add" and not has:
+ pieces.append(tag)
elif op == "prepend":
- tags = [tag] + [t for t in tags if t.lower() != tag.lower()]
- elif op == "remove" and tag.lower() in low:
- tags = [t for t in tags if t.lower() != tag.lower()]
- elif op == "replace" and tag.lower() in low:
- tags = [new if t.lower() == tag.lower() else t for t in tags]
+ pieces = [tag] + [p for p in pieces if not _same(p, tag)]
+ elif op == "remove" and has:
+ pieces = [p for p in pieces if not _same(p, tag)]
+ elif op == "replace" and has:
+ pieces = [new if _same(p, tag) else p for p in pieces]
else:
continue
- text = join_tags(tags)
+ text = ", ".join(pieces)[:MAX_CAPTION]
if text != item.get("caption"):
item["caption"] = text
changed += 1
return changed
-
diff --git a/modules/training/lora/nodes.py b/modules/training/lora/nodes.py
index fa36da9b..25db76a6 100644
--- a/modules/training/lora/nodes.py
+++ b/modules/training/lora/nodes.py
@@ -37,7 +37,7 @@ class FunPackTrainLora(io.ComfyNode):
io.Int.Input("frames", default=33, min=1, max=1025, tooltip="Frames taken from each clip, for video models. H3 rounds to its own grid (17 x n + 5)."),
io.Float.Input("fps", default=0.0, min=0.0, max=120.0, step=0.5, tooltip="Frames per second to read the clips at. 0 uses each clip's own rate (H3: 24)."),
io.Combo.Input("start", options=["start", "middle"], default="start", tooltip="Where in a clip to take the frames from."),
- io.String.Input("trigger", default="", tooltip="A word put at the start of every caption, to call the LoRA up later."),
+ io.String.Input("trigger", default="", tooltip="The prefix: text put at the start of every caption, for example a word to call the LoRA up later."),
io.Int.Input("seed", default=0, min=0, max=0xFFFFFFFFFFFFFFFF, tooltip="Seed for the noise and the LoRA's starting weights."),
io.Combo.Input("training_dtype", options=["bf16", "fp32", "none"], default="bf16", tooltip="Precision used while training. bf16 is right for nearly everything."),
io.Combo.Input("lora_dtype", options=["bf16", "fp32"], default="bf16", tooltip="Precision the LoRA is stored in."),
@@ -47,6 +47,7 @@ class FunPackTrainLora(io.ComfyNode):
io.String.Input("model_name", default="", tooltip="The diffusion model file to train on, as the loader lists it. The trainer loads it itself, once, with no loader node in front."),
io.Combo.Input("weight_dtype", options=["default", "fp8_e4m3fn", "fp8_e5m2", "fp16", "bf16", "fp32"], default="default", tooltip="How the model file's weights are stored when it is loaded here."),
io.Combo.Input("compute_dtype", options=["default", "fp16", "bf16", "fp32"], default="default", tooltip="What the maths runs in when the model is loaded here."),
+ io.String.Input("postfix", default="", tooltip="Text put at the end of every caption."),
],
outputs=[io.String.Output(display_name="lora")],
)
@@ -54,7 +55,7 @@ class FunPackTrainLora(io.ComfyNode):
@classmethod
def execute(cls, clip, vae, set_name, name, use, steps, learning_rate, rank, algorithm, optimizer, loss, batch_size,
grad_accumulation, resolution, aspect_buckets, frames, fps, start, trigger, seed, training_dtype, lora_dtype,
- gradient_checkpointing, offloading, bypass, model=None, audio_vae=None, model_name="", weight_dtype="default", compute_dtype="default") -> io.NodeOutput:
+ gradient_checkpointing, offloading, bypass, model=None, audio_vae=None, model_name="", weight_dtype="default", compute_dtype="default", postfix="") -> io.NodeOutput:
if not addons.enabled("lora_training"):
raise RuntimeError("LoRA training is switched off. Turn it on in Settings > Modules.")
cfg = {"name": name, "use": use, "steps": steps, "learning_rate": learning_rate, "rank": rank, "algorithm": algorithm,
@@ -62,5 +63,5 @@ class FunPackTrainLora(io.ComfyNode):
"resolution": resolution, "aspect_buckets": aspect_buckets, "frames": frames, "fps": fps, "start": start,
"trigger": trigger, "seed": seed, "training_dtype": training_dtype, "lora_dtype": lora_dtype,
"gradient_checkpointing": gradient_checkpointing, "offloading": offloading, "bypass": bypass,
- "model_name": model_name, "weight_dtype": weight_dtype, "compute_dtype": compute_dtype}
+ "model_name": model_name, "weight_dtype": weight_dtype, "compute_dtype": compute_dtype, "postfix": postfix}
return io.NodeOutput(train.run(model, clip, vae, audio_vae, set_name, cfg))
diff --git a/modules/training/lora/store.py b/modules/training/lora/store.py
index 6fc85fae..4936a3d3 100644
--- a/modules/training/lora/store.py
+++ b/modules/training/lora/store.py
@@ -102,7 +102,8 @@ def clean(doc, name: str) -> dict:
seen.add(src_key(item["src"]))
items.append(item)
settings = doc.get("settings")
- return {"name": name, "trigger": " ".join(str(doc.get("trigger") or "").split())[:80], "items": items,
+ return {"name": name, "trigger": " ".join(str(doc.get("trigger") or "").split())[:80],
+ "postfix": " ".join(str(doc.get("postfix") or "").split())[:80], "items": items,
"settings": settings if isinstance(settings, dict) else {}, "updated": time.time()}
diff --git a/modules/training/lora/tests/test_sets_and_tags.py b/modules/training/lora/tests/test_sets_and_tags.py
index e35aaa96..33abd3a6 100644
--- a/modules/training/lora/tests/test_sets_and_tags.py
+++ b/modules/training/lora/tests/test_sets_and_tags.py
@@ -17,8 +17,11 @@ def out(name, **extra):
def test_tags_are_trimmed_and_unique():
assert captions.split_tags(" a,B ,, a\nc ") == ["a", "B", "c"]
- assert captions.with_trigger("red, zq, tall", "zq") == "zq, red, tall"
- assert captions.with_trigger("red", "") == "red"
+ assert captions.with_parts("red, tall", "zq", "end") == "zq, red, tall, end"
+ assert captions.with_parts("red", "", "") == "red" and captions.with_parts("", "zq", "") == "zq"
+ assert captions.with_parts("A fox runs.\nIt jumps, twice.", " zq ", None) == "zq, A fox runs.\nIt jumps, twice."
+ assert captions.add_tags("a fox, running home.", ["A Fox", "grass", "sky"]) == "a fox, running home., grass, sky"
+ assert captions.add_tags("", ["x", "y"]) == "x, y" and captions.add_tags("a long prompt,", ["x"]) == "a long prompt, x"
def test_bulk_edits_only_the_chosen_items():
diff --git a/modules/training/lora/train.py b/modules/training/lora/train.py
index 276d1f2b..3fcd8356 100644
--- a/modules/training/lora/train.py
+++ b/modules/training/lora/train.py
@@ -148,8 +148,7 @@ def run(model, clip, vae, audio_vae, set_name: str, cfg: dict, prompt_id: str =
raise ValueError(f"{item['name'] or item['id']} cannot be read: {exc}") from exc
single = bool(adapt) or not cfg["aspect_buckets"]
buckets = media_io.buckets_for([w / h for w, h in sizes], cfg["resolution"], step, single)
- trigger = cfg["trigger"]
- texts = [captions.with_trigger(item["caption"], trigger) for item, _ in items]
+ texts = [captions.with_parts(item["caption"], cfg["trigger"], cfg.get("postfix", "")) for item, _ in items]
if adapt and adapt.get("video"):
cfg = _fit_frames(cfg, adapt, items)