105 lines
4.1 KiB
TypeScript
105 lines
4.1 KiB
TypeScript
/**
|
|
* Read the `__metadata__` header of a safetensors file — the slot HOT-Step's
|
|
* trainers use to record an adapter's trigger word.
|
|
*
|
|
* The format is: 8 bytes little-endian u64 header length, then that many bytes
|
|
* of JSON, then the tensor payload. Unknown metadata keys are ignored by every
|
|
* other consumer (PEFT, ComfyUI, LyCORIS, Side-Step, diffusers), which is why
|
|
* the trigger lives here and not in `adapter_config.json`.
|
|
*
|
|
* docs/plans/2026-07-28-adapter-trigger-embedding.md §2.1 / §4
|
|
*
|
|
* Nothing in this module throws. An unreadable, truncated or non-safetensors
|
|
* file degrades to "no metadata" — an adapter without a trigger is a normal
|
|
* state, not an error.
|
|
*/
|
|
import fs from 'fs';
|
|
import path from 'path';
|
|
|
|
/** Refuse to buffer a bogus header length. Real headers are well under a MB. */
|
|
const MAX_HEADER_BYTES = 64 * 1024 * 1024;
|
|
|
|
export interface AdapterTrigger {
|
|
/** '' when the adapter carries no embedded trigger. */
|
|
trigger: string;
|
|
/** 'prepend' | 'append', or '' when there is no trigger. */
|
|
position: 'prepend' | 'append' | '';
|
|
}
|
|
|
|
const EMPTY: AdapterTrigger = { trigger: '', position: '' };
|
|
|
|
/** `path -> {size, mtimeMs, value}` so a directory scan reads each file once. */
|
|
const cache = new Map<string, { size: number; mtimeMs: number; value: AdapterTrigger }>();
|
|
|
|
/** Read only the `__metadata__` object. Returns `{}` on any failure. */
|
|
export function readSafetensorsMetadata(file: string): Record<string, string> {
|
|
let fd: number | undefined;
|
|
try {
|
|
fd = fs.openSync(file, 'r');
|
|
const lenBuf = Buffer.allocUnsafe(8);
|
|
if (fs.readSync(fd, lenBuf, 0, 8, 0) !== 8) return {};
|
|
const headerLen = Number(lenBuf.readBigUInt64LE(0));
|
|
if (!Number.isSafeInteger(headerLen) || headerLen <= 0 || headerLen > MAX_HEADER_BYTES) return {};
|
|
|
|
const hdr = Buffer.allocUnsafe(headerLen);
|
|
if (fs.readSync(fd, hdr, 0, headerLen, 8) !== headerLen) return {};
|
|
|
|
const parsed: unknown = JSON.parse(hdr.toString('utf8'));
|
|
if (!parsed || typeof parsed !== 'object') return {};
|
|
const md = (parsed as Record<string, unknown>).__metadata__;
|
|
if (!md || typeof md !== 'object') return {};
|
|
|
|
const out: Record<string, string> = {};
|
|
for (const [k, v] of Object.entries(md as Record<string, unknown>)) {
|
|
if (typeof v === 'string') out[k] = v;
|
|
}
|
|
return out;
|
|
} catch {
|
|
return {};
|
|
} finally {
|
|
if (fd !== undefined) {
|
|
try { fs.closeSync(fd); } catch { /* ignore */ }
|
|
}
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Resolve an adapter's embedded trigger. Accepts either a bare `.safetensors`
|
|
* file or an adapter directory (in which case `adapter_model.safetensors` inside
|
|
* it is read, falling back to `lokr_weights.safetensors` — the LyCORIS leaf a
|
|
* DiT LoKR export writes instead, carrying the same trigger keys). Memoised on
|
|
* path + size + mtime, so re-scanning a folder of 200 adapters costs one stat
|
|
* each after the first pass.
|
|
*/
|
|
export function readAdapterTrigger(pathOrDir: string): AdapterTrigger {
|
|
try {
|
|
let file = pathOrDir;
|
|
const st = fs.statSync(pathOrDir);
|
|
if (st.isDirectory()) {
|
|
file = path.join(pathOrDir, 'adapter_model.safetensors');
|
|
if (!fs.existsSync(file)) file = path.join(pathOrDir, 'lokr_weights.safetensors');
|
|
}
|
|
|
|
const fst = file === pathOrDir ? st : fs.statSync(file);
|
|
if (!fst.isFile()) return EMPTY;
|
|
|
|
const hit = cache.get(file);
|
|
if (hit && hit.size === fst.size && hit.mtimeMs === fst.mtimeMs) return hit.value;
|
|
|
|
const md = readSafetensorsMetadata(file);
|
|
// `hot_step_trigger` is authoritative; `modelspec.trigger_phrase` is the
|
|
// ecosystem convention we also write, and lets us read adapters stamped by
|
|
// kohya-style tooling that never heard of HOT-Step.
|
|
const trigger = (md.hot_step_trigger || md['modelspec.trigger_phrase'] || '').trim();
|
|
const raw = (md.hot_step_trigger_position || '').trim();
|
|
const position: AdapterTrigger['position'] =
|
|
!trigger ? '' : raw === 'append' ? 'append' : 'prepend';
|
|
|
|
const value: AdapterTrigger = trigger ? { trigger, position } : EMPTY;
|
|
cache.set(file, { size: fst.size, mtimeMs: fst.mtimeMs, value });
|
|
return value;
|
|
} catch {
|
|
return EMPTY;
|
|
}
|
|
}
|