332 lines
14 KiB
JavaScript
332 lines
14 KiB
JavaScript
#!/usr/bin/env node
|
|
/**
|
|
* Stamp trigger words into adapters that were trained before HOT-Step embedded
|
|
* them (every Side-Step-trained adapter, and every HOT-Step one from before
|
|
* 2026-07-28).
|
|
*
|
|
* docs/plans/2026-07-28-adapter-trigger-embedding.md §6
|
|
*
|
|
* Adds `hot_step_trigger`, `hot_step_trigger_position` and
|
|
* `modelspec.trigger_phrase` to a safetensors file's `__metadata__`. Tensor
|
|
* bytes are copied verbatim, so the weights are provably unchanged — only the
|
|
* JSON header grows. Unknown metadata keys are ignored by every other consumer,
|
|
* so a stamped adapter still loads in ComfyUI / PEFT / Side-Step.
|
|
*
|
|
* DRY RUN BY DEFAULT. Nothing is written without --apply.
|
|
*
|
|
* node server/scripts/stamp-adapter-triggers.mjs
|
|
* node server/scripts/stamp-adapter-triggers.mjs --apply
|
|
*
|
|
* --adapters <dir> adapters root (default: ../../adapters from this file)
|
|
* --datasets <dir> a Side-Step corpus root; each <dir>/<name>/dataset.json
|
|
* supplies that adapter's REAL custom_tag + tag_position,
|
|
* which beats guessing from the filename
|
|
* --map <file.json> {"<adapterName>": "<trigger>"} explicit overrides
|
|
* --position <p> prepend|append fallback when no dataset says otherwise
|
|
* --only <lm|dit|all> which subtree to walk (default all)
|
|
* --no-backup skip the .bak copy (not recommended)
|
|
* --keep-backup keep every .bak instead of releasing it after verify
|
|
* --force restamp adapters that already carry a trigger
|
|
*/
|
|
import fs from 'fs';
|
|
import path from 'path';
|
|
import { fileURLToPath } from 'url';
|
|
|
|
const HERE = path.dirname(fileURLToPath(import.meta.url));
|
|
const MAX_HEADER_BYTES = 64 * 1024 * 1024;
|
|
const SIZE_SUFFIX = /-(?:0\.6B|1\.7B|4B)$/;
|
|
|
|
// ── args ────────────────────────────────────────────────────────────────────
|
|
const argv = process.argv.slice(2);
|
|
const flag = (name) => argv.includes(`--${name}`);
|
|
const opt = (name, dflt) => {
|
|
const i = argv.indexOf(`--${name}`);
|
|
return i >= 0 && i + 1 < argv.length ? argv[i + 1] : dflt;
|
|
};
|
|
|
|
const APPLY = flag('apply');
|
|
const FORCE = flag('force');
|
|
const BACKUP = !flag('no-backup');
|
|
// A .bak is taken per file and released as soon as that file verifies, so peak
|
|
// extra disk is one adapter rather than a second copy of the whole corpus.
|
|
// --keep-backup retains them all instead.
|
|
const KEEP_BACKUP = flag('keep-backup');
|
|
const POSITION = opt('position', 'prepend');
|
|
const ONLY = opt('only', 'all');
|
|
const ADAPTERS_ROOT = path.resolve(opt('adapters', path.join(HERE, '..', '..', 'adapters')));
|
|
const DATASETS_DIR = opt('datasets', '');
|
|
const MAP_FILE = opt('map', '');
|
|
|
|
if (POSITION !== 'prepend' && POSITION !== 'append') {
|
|
console.error(`--position must be prepend|append (got "${POSITION}")`);
|
|
process.exit(2);
|
|
}
|
|
|
|
let explicitMap = {};
|
|
if (MAP_FILE) {
|
|
try {
|
|
explicitMap = JSON.parse(fs.readFileSync(path.resolve(MAP_FILE), 'utf8'));
|
|
} catch (e) {
|
|
console.error(`cannot read --map ${MAP_FILE}: ${e.message}`);
|
|
process.exit(2);
|
|
}
|
|
}
|
|
|
|
// ── safetensors header I/O ──────────────────────────────────────────────────
|
|
|
|
function readHeader(file) {
|
|
const fd = fs.openSync(file, 'r');
|
|
try {
|
|
const lenBuf = Buffer.allocUnsafe(8);
|
|
if (fs.readSync(fd, lenBuf, 0, 8, 0) !== 8) return null;
|
|
const headerLen = Number(lenBuf.readBigUInt64LE(0));
|
|
if (!Number.isSafeInteger(headerLen) || headerLen <= 0 || headerLen > MAX_HEADER_BYTES) return null;
|
|
const hdr = Buffer.allocUnsafe(headerLen);
|
|
if (fs.readSync(fd, hdr, 0, headerLen, 8) !== headerLen) return null;
|
|
return { json: JSON.parse(hdr.toString('utf8')), headerLen };
|
|
} catch {
|
|
return null;
|
|
} finally {
|
|
fs.closeSync(fd);
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Rewrite `file` with `metadata` merged into its `__metadata__`.
|
|
*
|
|
* Streams `newLen(8) + newHeader + originalPayload` to a temp file, then swaps.
|
|
* Never edits in place, so an interrupted run leaves the original intact.
|
|
*/
|
|
function stamp(file, metadata) {
|
|
const head = readHeader(file);
|
|
if (!head) throw new Error('unreadable safetensors header');
|
|
|
|
const json = { ...head.json, __metadata__: { ...(head.json.__metadata__ || {}), ...metadata } };
|
|
const newHeader = Buffer.from(JSON.stringify(json), 'utf8');
|
|
// safetensors requires the header to be 8-byte aligned; pad with spaces,
|
|
// which are legal JSON whitespace outside string literals.
|
|
const pad = (8 - (newHeader.length % 8)) % 8;
|
|
const padded = Buffer.concat([newHeader, Buffer.alloc(pad, 0x20)]);
|
|
const lenBuf = Buffer.allocUnsafe(8);
|
|
lenBuf.writeBigUInt64LE(BigInt(padded.length), 0);
|
|
|
|
const originalSize = fs.statSync(file).size;
|
|
const payloadBytes = originalSize - 8 - head.headerLen;
|
|
const expectedSize = 8 + padded.length + payloadBytes;
|
|
const originalNames = Object.keys(head.json).filter(k => k !== '__metadata__').sort().join('\n');
|
|
|
|
const tmp = `${file}.stamping`;
|
|
const out = fs.openSync(tmp, 'w');
|
|
const src = fs.openSync(file, 'r');
|
|
try {
|
|
fs.writeSync(out, lenBuf);
|
|
fs.writeSync(out, padded);
|
|
const buf = Buffer.allocUnsafe(4 * 1024 * 1024);
|
|
let pos = 8 + head.headerLen;
|
|
for (;;) {
|
|
const n = fs.readSync(src, buf, 0, buf.length, pos);
|
|
if (n <= 0) break;
|
|
fs.writeSync(out, buf, 0, n);
|
|
pos += n;
|
|
}
|
|
} finally {
|
|
fs.closeSync(src);
|
|
fs.closeSync(out);
|
|
}
|
|
|
|
const bak = `${file}.bak`;
|
|
if (BACKUP) fs.copyFileSync(file, bak);
|
|
fs.rmSync(file);
|
|
fs.renameSync(tmp, file);
|
|
|
|
// Verify the swapped file before moving on. The payload is a verbatim byte
|
|
// copy, so an exact size match plus an identical tensor-entry set is enough
|
|
// to catch truncation, a short write or a mangled header — the only ways this
|
|
// operation can go wrong. On failure, put the original back.
|
|
try {
|
|
const size = fs.statSync(file).size;
|
|
const check = readHeader(file);
|
|
if (size !== expectedSize) throw new Error(`size ${size} != expected ${expectedSize}`);
|
|
if (!check) throw new Error('rewritten header does not parse');
|
|
if (check.json.__metadata__?.hot_step_trigger !== metadata.hot_step_trigger) {
|
|
throw new Error('trigger missing from the rewritten header');
|
|
}
|
|
const names = Object.keys(check.json).filter(k => k !== '__metadata__').sort().join('\n');
|
|
if (names !== originalNames) throw new Error('tensor entry set changed');
|
|
} catch (e) {
|
|
if (BACKUP && fs.existsSync(bak)) {
|
|
fs.rmSync(file, { force: true });
|
|
fs.renameSync(bak, file);
|
|
throw new Error(`${e.message} — ORIGINAL RESTORED from .bak`);
|
|
}
|
|
throw new Error(`${e.message} — NO BACKUP TO RESTORE (ran with --no-backup)`);
|
|
}
|
|
|
|
if (BACKUP && !KEEP_BACKUP) fs.rmSync(bak, { force: true });
|
|
}
|
|
|
|
// ── discovery ───────────────────────────────────────────────────────────────
|
|
|
|
/** Every adapter under `root`: PEFT dirs (unversioned, and each stamped run
|
|
* subfolder — per-run layout) plus bare .safetensors files. `name` stays the
|
|
* ARTIST name for runs, which is what the trigger proposal keys on. */
|
|
function findAdapters(root, kind) {
|
|
const out = [];
|
|
if (!fs.existsSync(root)) return out;
|
|
for (const e of fs.readdirSync(root, { withFileTypes: true })) {
|
|
if (e.name.startsWith('.')) continue;
|
|
const full = path.join(root, e.name);
|
|
if (e.isDirectory()) {
|
|
const model = path.join(full, 'adapter_model.safetensors');
|
|
if (fs.existsSync(model)) out.push({ name: e.name, file: model, kind });
|
|
for (const run of fs.readdirSync(full, { withFileTypes: true })) {
|
|
if (!run.isDirectory() || run.name.startsWith('.') || run.name === 'milestones') continue;
|
|
const runModel = path.join(full, run.name, 'adapter_model.safetensors');
|
|
if (fs.existsSync(runModel)) out.push({ name: e.name, file: runModel, kind });
|
|
}
|
|
} else if (e.isFile() && e.name.endsWith('.safetensors')) {
|
|
out.push({ name: e.name.replace(/\.safetensors$/i, ''), file: full, kind });
|
|
}
|
|
}
|
|
return out;
|
|
}
|
|
|
|
/**
|
|
* `custom_tag` + `tag_position` per dataset name — the ground truth for what an
|
|
* adapter was actually trained with. Reads both the Training Studio's own
|
|
* `dataset_meta.json` and any Side-Step corpus passed via --datasets, whose
|
|
* `dataset.json` carries the same fields under `metadata`.
|
|
*/
|
|
function loadDatasetTags() {
|
|
const tags = new Map();
|
|
|
|
const addStudio = (dsRoot) => {
|
|
if (!fs.existsSync(dsRoot)) return;
|
|
for (const e of fs.readdirSync(dsRoot, { withFileTypes: true })) {
|
|
if (!e.isDirectory()) continue;
|
|
try {
|
|
const meta = JSON.parse(fs.readFileSync(path.join(dsRoot, e.name, 'dataset_meta.json'), 'utf8'));
|
|
const tag = (meta.customTag || meta.custom_tag || '').trim();
|
|
const pos = (meta.tagPosition || meta.tag_position || '').trim();
|
|
if (tag) tags.set(e.name, { tag, pos });
|
|
} catch { /* a dataset without readable meta simply contributes nothing */ }
|
|
}
|
|
};
|
|
|
|
const addSideStep = (root) => {
|
|
if (!root || !fs.existsSync(root)) return;
|
|
for (const e of fs.readdirSync(root, { withFileTypes: true })) {
|
|
if (!e.isDirectory()) continue;
|
|
try {
|
|
const j = JSON.parse(fs.readFileSync(path.join(root, e.name, 'dataset.json'), 'utf8'));
|
|
const md = j.metadata || {};
|
|
const tag = (md.custom_tag || '').trim();
|
|
const pos = (md.tag_position || '').trim();
|
|
if (tag) tags.set(e.name, { tag, pos });
|
|
} catch { /* likewise */ }
|
|
}
|
|
};
|
|
|
|
addStudio(path.join(HERE, '..', 'data', 'training', 'datasets'));
|
|
if (DATASETS_DIR) addSideStep(path.resolve(DATASETS_DIR));
|
|
return tags;
|
|
}
|
|
|
|
/** explicit map > dataset custom_tag > dir name minus the -<size> suffix. */
|
|
function proposeTrigger(name, datasetTags) {
|
|
if (explicitMap[name]) return { trigger: String(explicitMap[name]).trim(), position: POSITION, source: 'map' };
|
|
const bare = name.replace(SIZE_SUFFIX, '');
|
|
const hit = datasetTags.get(name) || datasetTags.get(bare);
|
|
if (hit) {
|
|
// A dataset whose tag_position is "replace" never put the tag in its
|
|
// captions at all (preprocess-run.h:203), so that adapter has no trigger.
|
|
if (hit.pos === 'replace') return { trigger: '', position: '', source: 'replace' };
|
|
return { trigger: hit.tag, position: hit.pos === 'append' ? 'append' : 'prepend', source: 'dataset' };
|
|
}
|
|
return { trigger: bare, position: POSITION, source: 'filename' };
|
|
}
|
|
|
|
// ── run ─────────────────────────────────────────────────────────────────────
|
|
|
|
const datasetTags = loadDatasetTags();
|
|
const targets = [];
|
|
// Per-base layout (adapterLayout.ts): planner adapters under lm/ (legacy) and
|
|
// lm-06b/lm-17b/lm-4b; DiT adapters at the root (legacy) and under dit-*.
|
|
if (ONLY === 'all' || ONLY === 'dit') {
|
|
// Root level: LEGACY bare files / PEFT dirs only. The lm*/dit-* trees are
|
|
// walked by their own passes — descending into them here double-counts.
|
|
for (const e of fs.existsSync(ADAPTERS_ROOT) ? fs.readdirSync(ADAPTERS_ROOT, { withFileTypes: true }) : []) {
|
|
if (e.name.startsWith('.') || /^(lm|lm-|dit-)/i.test(e.name)) continue;
|
|
const full = path.join(ADAPTERS_ROOT, e.name);
|
|
if (e.isDirectory()) {
|
|
const model = path.join(full, 'adapter_model.safetensors');
|
|
if (fs.existsSync(model)) targets.push({ name: e.name, file: model, kind: 'dit' });
|
|
} else if (e.isFile() && e.name.endsWith('.safetensors')) {
|
|
targets.push({ name: e.name.replace(/\.safetensors$/i, ''), file: full, kind: 'dit' });
|
|
}
|
|
}
|
|
for (const e of fs.existsSync(ADAPTERS_ROOT) ? fs.readdirSync(ADAPTERS_ROOT, { withFileTypes: true }) : []) {
|
|
if (e.isDirectory() && /^dit-/i.test(e.name)) {
|
|
targets.push(...findAdapters(path.join(ADAPTERS_ROOT, e.name), 'dit'));
|
|
}
|
|
}
|
|
}
|
|
if (ONLY === 'all' || ONLY === 'lm') {
|
|
for (const sub of ['lm', 'lm-06b', 'lm-17b', 'lm-4b']) {
|
|
targets.push(...findAdapters(path.join(ADAPTERS_ROOT, sub), 'lm'));
|
|
}
|
|
}
|
|
|
|
if (!targets.length) {
|
|
console.error(`No adapters found under ${ADAPTERS_ROOT} (--only ${ONLY}).`);
|
|
process.exit(1);
|
|
}
|
|
|
|
const rows = [];
|
|
let skipped = 0;
|
|
for (const t of targets) {
|
|
const head = readHeader(t.file);
|
|
if (!head) { rows.push({ ...t, trigger: '', source: 'UNREADABLE', action: 'skip' }); continue; }
|
|
const existing = (head.json.__metadata__ || {}).hot_step_trigger;
|
|
if (existing && !FORCE) { skipped++; continue; }
|
|
const { trigger, position, source } = proposeTrigger(t.name, datasetTags);
|
|
rows.push({ ...t, trigger, position, source, action: trigger ? (existing ? 'restamp' : 'stamp') : 'skip' });
|
|
}
|
|
|
|
const pad = (s, n) => String(s).padEnd(n);
|
|
console.log(`\nAdapters root: ${ADAPTERS_ROOT}`);
|
|
console.log(`${targets.length} adapter(s) found, ${skipped} already stamped (skipped), ${rows.length} to review.`);
|
|
console.log(`Fallback position: ${POSITION} Mode: ${APPLY ? 'APPLY' : 'DRY RUN'}${BACKUP ? ' (keeping .bak)' : ''}`);
|
|
console.log(`Dataset tags loaded: ${datasetTags.size}${DATASETS_DIR ? ` (incl. ${DATASETS_DIR})` : ''}\n`);
|
|
console.log(`${pad('KIND', 5)} ${pad('ADAPTER', 42)} ${pad('PROPOSED TRIGGER', 28)} ${pad('POS', 8)} ${pad('SOURCE', 9)} ACTION`);
|
|
console.log('-'.repeat(110));
|
|
for (const r of rows) {
|
|
console.log(`${pad(r.kind, 5)} ${pad(r.name.slice(0, 42), 42)} ${pad(r.trigger.slice(0, 28), 28)} ${pad(r.position || '-', 8)} ${pad(r.source, 9)} ${r.action}`);
|
|
}
|
|
const bySource = rows.reduce((m, r) => ({ ...m, [r.source]: (m[r.source] || 0) + 1 }), {});
|
|
console.log(`\nBy source: ${Object.entries(bySource).map(([k, v]) => `${k}=${v}`).join(' ')}`);
|
|
|
|
if (!APPLY) {
|
|
console.log(`\nDRY RUN — nothing written. Re-run with --apply once the table above is right.`);
|
|
console.log(`Fix individual rows with --map triggers.json, e.g. {"abba-4B": "abba"}\n`);
|
|
process.exit(0);
|
|
}
|
|
|
|
let ok = 0, failed = 0;
|
|
for (const r of rows) {
|
|
if (r.action === 'skip') continue;
|
|
try {
|
|
stamp(r.file, {
|
|
hot_step_trigger: r.trigger,
|
|
hot_step_trigger_position: r.position || POSITION,
|
|
'modelspec.trigger_phrase': r.trigger,
|
|
});
|
|
ok++;
|
|
} catch (e) {
|
|
console.error(`FAILED ${r.name}: ${e.message}`);
|
|
failed++;
|
|
}
|
|
}
|
|
console.log(`\nStamped ${ok} adapter(s)${failed ? `, ${failed} failed` : ''}.`);
|
|
process.exit(failed ? 1 : 0);
|