utkucoban's picture
Upload 4 files
fc78499 verified
Raw
History Blame Contribute Delete
20 kB
// NanoMaestro Web Worker for isolated model inference and parsing
try {
importScripts("https://cdn.jsdelivr.net/npm/onnxruntime-web@1.20.1/dist/ort.min.js");
} catch (err) {
console.error("Worker failed to import ONNX Runtime script:", err);
}
// ONNX runtime configuration
if (typeof ort !== "undefined") {
ort.env.wasm.wasmPaths = "https://cdn.jsdelivr.net/npm/onnxruntime-web@1.20.1/dist/";
ort.env.wasm.numThreads = Math.min(4, navigator.hardwareConcurrency || 1);
}
// State variables
let session = null;
let stoi = null;
let itos = null;
let vocabSize = 0;
let h = null;
let c = null;
let currentId = 0;
let activeModelName = "";
// Model parameters
let hiddenSize = 2048;
let numLayers = 2;
// Playback settings
let temperature = 0.85;
let topK = 40;
let currentBpm = 120;
let midiBpmMode = "lock";
// Seeding settings
let playingStyle = "default";
let midiTokens = [];
let midiStartBar = 0;
let isWarmingUp = false;
// Temporal queues for gathering parsed events
let tempQueue = [];
let parser = null;
function tokenId(token) {
return stoi?.[token] ?? null;
}
function makeTensorId(id) {
return new ort.Tensor("int64", BigInt64Array.from([BigInt(id)]), [1, 1]);
}
function zeroState() {
return new ort.Tensor("float32", new Float32Array(numLayers * 1 * hiddenSize), [numLayers, 1, hiddenSize]);
}
function updateModelDimensions(session) {
try {
const inputNames = session.inputNames;
if (inputNames.includes("h") && session.handler && session.handler._model) {
const inputs = session.handler._model.graph.inputs;
const hInput = inputs.find(i => i.name === "h");
if (hInput && hInput.type && hInput.type.tensorType && hInput.type.tensorType.shape) {
const shape = hInput.type.tensorType.shape.dim;
const layers = Number(shape[0].dimValue);
const hidden = Number(shape[2].dimValue);
if (layers > 0 && hidden > 0) {
numLayers = layers;
hiddenSize = hidden;
console.log(`Worker set model dimensions: layers=${numLayers}, hiddenSize=${hiddenSize}`);
}
}
}
} catch (e) {
console.error("Worker failed to detect model dimensions:", e);
}
}
// Model Step Execution with Memory Management
async function stepModel(record = true) {
if (!session) return "?";
const inputTensor = makeTensorId(currentId);
const output = await session.run({ input: inputTensor, h, c });
// Dispose input tensor to prevent memory leak
inputTensor.dispose();
// Dispose previous recurrent states to prevent memory leaks during continuous generation
if (h) {
try { h.dispose(); } catch (e) {}
}
if (c) {
try { c.dispose(); } catch (e) {}
}
h = output.h_out;
c = output.c_out;
const logitsData = output.logits.data;
currentId = sampleFromLogits(logitsData);
output.logits.dispose();
const token = itos[currentId] ?? "?";
// Reset state on EOS
if (token === "EOS") {
if (h) { try { h.dispose(); } catch (e) {} }
if (c) { try { c.dispose(); } catch (e) {} }
h = zeroState();
c = zeroState();
}
if (record && parser) {
parser.feed(token);
}
return token;
}
function sampleFromLogits(logits) {
const temp = Math.max(0.05, temperature);
const k = Math.max(1, topK);
const scored = [];
for (let i = 0; i < logits.length; i += 1) {
const token = itos[i];
if (token === undefined) continue;
let score = logits[i] / temp;
if (token === "<EOP>" || token === "EOS") score -= 1.0;
scored.push([i, score]);
}
scored.sort((a, b) => b[1] - a[1]);
const picked = scored.slice(0, Math.min(k, scored.length));
const maxScore = picked[0]?.[1] ?? 0;
let sum = 0;
for (const item of picked) {
item[2] = Math.exp(item[1] - maxScore);
sum += item[2];
}
let r = Math.random() * sum;
for (const item of picked) {
r -= item[2];
if (r <= 0) return item[0];
}
return picked[picked.length - 1][0];
}
// Duration converters
function durationToSecondsFromEventSteps(steps, grid) {
const quarterSeconds = 60 / currentBpm;
return Math.max(0.035, (Math.max(1, steps) * 4 * quarterSeconds) / Math.max(1, grid));
}
function midiToTonePitch(midi) {
const names = ["C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B"];
return `${names[((midi % 12) + 12) % 12]}${Math.floor(midi / 12) - 1}`;
}
function pushEvent(event) {
if (isWarmingUp) return; // Discard prompt history events
tempQueue.push(event);
}
// Token Slicers
function getTokensUpToBar(tokens, targetBar) {
let barCount = 0;
const sliced = [];
for (const t of tokens) {
if (t === "BAR") {
barCount += 1;
if (barCount > targetBar) {
break;
}
}
sliced.push(t);
}
return sliced;
}
// Event Parser
function makeEventParser() {
return {
grid: 64,
bar: -1,
pos: 0,
pendingPosition: null,
pendingNotes: [],
pendingNote: null,
lastQ: 0,
reset() {
this.grid = 64;
this.bar = -1;
this.pos = 0;
this.pendingPosition = null;
this.pendingNotes = [];
this.pendingNote = null;
this.lastQ = 0;
},
feed(token) {
if (token.startsWith("BPM_")) {
const bpm = Number(token.slice(4));
if (Number.isFinite(bpm) && bpm >= 40 && bpm <= 220) {
if (playingStyle !== "midi" || isWarmingUp || midiBpmMode === "model") {
currentBpm = bpm;
self.postMessage({ action: "tempo", bpm: bpm });
}
}
return;
}
if (token.startsWith("GRID_")) {
const grid = Number(token.slice(5));
if (Number.isFinite(grid) && grid > 0) this.grid = grid;
return;
}
if (token === "BAR") {
this.flushTo(this.absoluteQFor(this.bar + 1, 0));
this.bar += 1;
this.pos = 0;
return;
}
if (token.startsWith("POS_")) {
const pos = Number(token.slice(4));
if (!Number.isFinite(pos)) return;
this.flushTo(this.absoluteQFor(this.bar, pos));
this.pos = pos;
return;
}
if (token.startsWith("NOTE_")) {
const midi = Number(token.slice(5));
if (Number.isFinite(midi) && midi >= 0 && midi <= 127) {
this.pendingNote = { midi, durationSteps: 1, velocity: 0.72 };
}
return;
}
if (token.startsWith("DUR_") && this.pendingNote) {
const steps = Number(token.slice(4));
if (Number.isFinite(steps)) this.pendingNote.durationSteps = Math.max(1, steps);
return;
}
if (token.startsWith("VEL_") && this.pendingNote) {
const bucket = Number(token.slice(4));
if (Number.isFinite(bucket)) this.pendingNote.velocity = Math.max(0.2, Math.min(0.95, bucket / 8));
const q = this.absoluteQFor(this.bar, this.pos);
this.pendingPosition ??= q;
this.pendingNotes.push(this.pendingNote);
this.pendingNote = null;
}
},
absoluteQFor(bar, pos) {
return Math.max(0, bar) * 4 + (Math.max(0, pos) * 4) / Math.max(1, this.grid);
},
flushTo(nextQ) {
if (this.pendingNote) {
const q = this.absoluteQFor(this.bar, this.pos);
this.pendingPosition ??= q;
this.pendingNotes.push(this.pendingNote);
this.pendingNote = null;
}
if (this.pendingNotes.length && this.pendingPosition !== null) {
const gap = Math.max(0, this.pendingPosition - this.lastQ);
if (gap > 0) pushEvent({ type: "rest", duration: this.quartersToSeconds(gap) });
pushEvent({
type: "note",
notes: this.pendingNotes.map((note) => midiToTonePitch(note.midi)),
perNoteDurations: this.pendingNotes.map((note) => durationToSecondsFromEventSteps(note.durationSteps, this.grid)),
velocities: this.pendingNotes.map((note) => note.velocity),
duration: 0,
advance: 0,
});
this.lastQ = this.pendingPosition;
}
const finalGap = Math.max(0, nextQ - this.lastQ);
if (finalGap > 0) pushEvent({ type: "rest", duration: this.quartersToSeconds(finalGap) });
this.lastQ = Math.max(this.lastQ, nextQ);
this.pendingPosition = null;
this.pendingNotes = [];
},
quartersToSeconds(quarters) {
return (quarters * 60) / currentBpm;
},
};
}
function makeParser() {
return makeEventParser();
}
async function warmPrompt() {
h = zeroState();
c = zeroState();
let allTokensStr;
if (playingStyle === "midi" && midiTokens.length > 0) {
allTokensStr = getTokensUpToBar(midiTokens, midiStartBar);
} else if (playingStyle === "bach") {
allTokensStr = [
"BOS", "BPM_100", "GRID_64",
"BAR",
"POS_0", "NOTE_81", "DUR_4", "VEL_6",
"POS_4", "NOTE_79", "DUR_4", "VEL_6",
"POS_8", "NOTE_77", "DUR_4", "VEL_6",
"POS_12", "NOTE_76", "DUR_4", "VEL_6",
"POS_16", "NOTE_74", "DUR_4", "VEL_6",
"POS_20", "NOTE_73", "DUR_4", "VEL_6",
"POS_24", "NOTE_74", "DUR_16", "VEL_6",
"BAR"
];
} else if (playingStyle === "mozart") {
allTokensStr = [
"BOS", "BPM_132", "GRID_64",
"BAR",
// Bright G-major chamber-style opening inspired by Mozart's serenades.
"POS_0", "NOTE_55", "DUR_8", "VEL_6",
"NOTE_67", "DUR_8", "VEL_6",
"POS_8", "NOTE_74", "DUR_4", "VEL_7",
"POS_12", "NOTE_71", "DUR_4", "VEL_6",
"POS_16", "NOTE_79", "DUR_8", "VEL_7",
"POS_24", "NOTE_78", "DUR_4", "VEL_6",
"POS_28", "NOTE_76", "DUR_4", "VEL_6",
"POS_32", "NOTE_74", "DUR_8", "VEL_7",
"POS_40", "NOTE_72", "DUR_4", "VEL_6",
"POS_44", "NOTE_71", "DUR_4", "VEL_6",
"POS_48", "NOTE_69", "DUR_8", "VEL_6",
"POS_56", "NOTE_67", "DUR_8", "VEL_7",
"BAR",
"POS_0", "NOTE_60", "DUR_16", "VEL_5",
"NOTE_64", "DUR_16", "VEL_5",
"NOTE_69", "DUR_16", "VEL_5",
"POS_16", "NOTE_72", "DUR_4", "VEL_6",
"POS_20", "NOTE_76", "DUR_4", "VEL_6",
"POS_24", "NOTE_81", "DUR_8", "VEL_7",
"POS_32", "NOTE_79", "DUR_4", "VEL_6",
"POS_36", "NOTE_78", "DUR_4", "VEL_6",
"POS_40", "NOTE_76", "DUR_4", "VEL_6",
"POS_44", "NOTE_74", "DUR_4", "VEL_6",
"POS_48", "NOTE_71", "DUR_8", "VEL_6",
"POS_56", "NOTE_67", "DUR_8", "VEL_7",
"BAR"
];
} else if (playingStyle === "chopin") {
allTokensStr = [
"BOS", "BPM_68", "GRID_64",
"BAR",
// Original C-minor nocturne texture: rolling bass beneath a lyrical upper line.
"POS_0", "NOTE_36", "DUR_16", "VEL_4", "NOTE_72", "DUR_12", "VEL_6",
"POS_8", "NOTE_43", "DUR_8", "VEL_4",
"POS_12", "NOTE_75", "DUR_8", "VEL_6",
"POS_16", "NOTE_48", "DUR_8", "VEL_4",
"POS_24", "NOTE_51", "DUR_8", "VEL_4", "NOTE_77", "DUR_12", "VEL_6",
"POS_32", "NOTE_39", "DUR_16", "VEL_4", "NOTE_79", "DUR_16", "VEL_7",
"POS_40", "NOTE_46", "DUR_8", "VEL_4",
"POS_48", "NOTE_51", "DUR_8", "VEL_4", "NOTE_77", "DUR_8", "VEL_6",
"POS_56", "NOTE_55", "DUR_8", "VEL_4", "NOTE_75", "DUR_8", "VEL_5",
"BAR",
"POS_0", "NOTE_41", "DUR_16", "VEL_4", "NOTE_72", "DUR_16", "VEL_6",
"POS_8", "NOTE_48", "DUR_8", "VEL_4",
"POS_16", "NOTE_53", "DUR_8", "VEL_4", "NOTE_74", "DUR_8", "VEL_6",
"POS_24", "NOTE_56", "DUR_8", "VEL_4", "NOTE_75", "DUR_8", "VEL_6",
"POS_32", "NOTE_43", "DUR_16", "VEL_4", "NOTE_79", "DUR_12", "VEL_7",
"POS_40", "NOTE_50", "DUR_8", "VEL_4",
"POS_44", "NOTE_77", "DUR_8", "VEL_6",
"POS_48", "NOTE_55", "DUR_8", "VEL_4",
"POS_56", "NOTE_60", "DUR_8", "VEL_4", "NOTE_75", "DUR_8", "VEL_5",
"BAR"
];
} else if (playingStyle === "debussy") {
allTokensStr = [
"BOS", "BPM_74", "GRID_64",
"BAR",
// Original impressionist sketch: open fifths and slowly shifting parallel colors.
"POS_0", "NOTE_38", "DUR_32", "VEL_3", "NOTE_57", "DUR_16", "VEL_4", "NOTE_62", "DUR_16", "VEL_4", "NOTE_69", "DUR_16", "VEL_5",
"POS_16", "NOTE_59", "DUR_16", "VEL_4", "NOTE_64", "DUR_16", "VEL_4", "NOTE_71", "DUR_16", "VEL_5",
"POS_32", "NOTE_45", "DUR_32", "VEL_3", "NOTE_61", "DUR_16", "VEL_4", "NOTE_66", "DUR_16", "VEL_4", "NOTE_73", "DUR_16", "VEL_5",
"POS_48", "NOTE_64", "DUR_16", "VEL_4", "NOTE_69", "DUR_16", "VEL_4", "NOTE_76", "DUR_16", "VEL_5",
"BAR",
"POS_0", "NOTE_40", "DUR_32", "VEL_3", "NOTE_59", "DUR_16", "VEL_4", "NOTE_64", "DUR_16", "VEL_4", "NOTE_71", "DUR_16", "VEL_5",
"POS_16", "NOTE_61", "DUR_16", "VEL_4", "NOTE_66", "DUR_16", "VEL_4", "NOTE_73", "DUR_16", "VEL_5",
"POS_32", "NOTE_43", "DUR_32", "VEL_3", "NOTE_62", "DUR_16", "VEL_4", "NOTE_67", "DUR_16", "VEL_4", "NOTE_74", "DUR_16", "VEL_5",
"POS_48", "NOTE_66", "DUR_16", "VEL_4", "NOTE_71", "DUR_16", "VEL_4", "NOTE_78", "DUR_16", "VEL_5",
"BAR"
];
} else if (playingStyle === "beethoven") {
allTokensStr = [
"BOS", "BPM_126", "GRID_64",
"BAR",
// Original dramatic sonata gesture: compact motif, octave bass, accented response.
"POS_0", "NOTE_36", "DUR_8", "VEL_7", "NOTE_48", "DUR_8", "VEL_7", "NOTE_67", "DUR_4", "VEL_7",
"POS_4", "NOTE_70", "DUR_4", "VEL_7",
"POS_8", "NOTE_68", "DUR_8", "VEL_8",
"POS_16", "NOTE_43", "DUR_8", "VEL_6", "NOTE_55", "DUR_8", "VEL_6", "NOTE_65", "DUR_4", "VEL_7",
"POS_20", "NOTE_68", "DUR_4", "VEL_7",
"POS_24", "NOTE_67", "DUR_8", "VEL_8",
"POS_32", "NOTE_39", "DUR_8", "VEL_7", "NOTE_51", "DUR_8", "VEL_7", "NOTE_63", "DUR_4", "VEL_7",
"POS_36", "NOTE_67", "DUR_4", "VEL_7",
"POS_40", "NOTE_72", "DUR_8", "VEL_8",
"POS_48", "NOTE_44", "DUR_8", "VEL_6", "NOTE_56", "DUR_8", "VEL_6", "NOTE_70", "DUR_4", "VEL_7",
"POS_52", "NOTE_68", "DUR_4", "VEL_7",
"POS_56", "NOTE_67", "DUR_8", "VEL_8",
"BAR",
"POS_0", "NOTE_41", "DUR_16", "VEL_7", "NOTE_53", "DUR_16", "VEL_7", "NOTE_65", "DUR_8", "VEL_7",
"POS_8", "NOTE_68", "DUR_4", "VEL_7",
"POS_12", "NOTE_72", "DUR_4", "VEL_8",
"POS_16", "NOTE_77", "DUR_8", "VEL_8",
"POS_24", "NOTE_75", "DUR_8", "VEL_7",
"POS_32", "NOTE_43", "DUR_16", "VEL_7", "NOTE_55", "DUR_16", "VEL_7", "NOTE_74", "DUR_8", "VEL_7",
"POS_40", "NOTE_72", "DUR_8", "VEL_7",
"POS_48", "NOTE_67", "DUR_16", "VEL_8", "NOTE_72", "DUR_16", "VEL_8", "NOTE_75", "DUR_16", "VEL_8",
"BAR"
];
} else if (playingStyle === "satie") {
allTokensStr = [
"BOS", "BPM_62", "GRID_64",
"BAR",
// Original sparse meditation: low bass, suspended chords, unhurried melody.
"POS_0", "NOTE_38", "DUR_32", "VEL_3", "NOTE_69", "DUR_20", "VEL_5",
"POS_16", "NOTE_57", "DUR_24", "VEL_3", "NOTE_62", "DUR_24", "VEL_3", "NOTE_65", "DUR_24", "VEL_3",
"POS_32", "NOTE_45", "DUR_32", "VEL_3", "NOTE_72", "DUR_16", "VEL_5",
"POS_48", "NOTE_55", "DUR_24", "VEL_3", "NOTE_60", "DUR_24", "VEL_3", "NOTE_64", "DUR_24", "VEL_3",
"BAR",
"POS_0", "NOTE_40", "DUR_32", "VEL_3", "NOTE_71", "DUR_20", "VEL_5",
"POS_16", "NOTE_59", "DUR_24", "VEL_3", "NOTE_64", "DUR_24", "VEL_3", "NOTE_67", "DUR_24", "VEL_3",
"POS_32", "NOTE_43", "DUR_32", "VEL_3", "NOTE_74", "DUR_16", "VEL_5",
"POS_48", "NOTE_57", "DUR_24", "VEL_3", "NOTE_62", "DUR_24", "VEL_3", "NOTE_66", "DUR_24", "VEL_3",
"BAR"
];
} else {
allTokensStr = ["BOS", "BPM_120", "GRID_64", "BAR", "POS_0"];
}
isWarmingUp = true;
// 1. Feed all prompt tokens to parser to advance its clock
if (parser) {
for (const token of allTokensStr) {
parser.feed(token);
}
}
// 2. Slice model warm-up tokens to last 256 (matching model's context window)
const modelTokensStr = allTokensStr.slice(-256);
const ids = modelTokensStr.map(tokenId).filter((id) => id !== null);
for (const id of ids) {
currentId = id;
await stepModel(false);
}
isWarmingUp = false;
}
async function pumpTokens(targetCount = 96) {
tempQueue = [];
let steps = 0;
// Keep stepping the model until we collect enough parsed events or hit loop bounds
while (tempQueue.length < targetCount && steps < 300) {
await stepModel(true);
steps += 1;
}
return {
events: tempQueue,
lastToken: itos[currentId] ?? "?"
};
}
// Message Router
self.onmessage = async function (e) {
const data = e.data;
switch (data.action) {
case "init":
try {
console.log(`Worker loading model: ${data.activeModelName}, requested EP: ${data.ep || "wasm"}`);
activeModelName = data.activeModelName;
stoi = data.vocab.stoi;
itos = Object.fromEntries(Object.entries(data.vocab.itos).map(([key, value]) => [Number(key), value]));
vocabSize = data.vocab.vocab_size;
if (data.modelConfig) {
numLayers = data.modelConfig.numLayers || 2;
hiddenSize = data.modelConfig.hiddenSize || 2048;
}
if (session) {
try { session.dispose(); } catch (e) {}
session = null;
}
let epUsed = data.ep || "wasm";
if (data.ep === "webgpu") {
try {
console.log("Attempting WebGPU execution provider...");
session = await ort.InferenceSession.create(data.modelBuffer, {
executionProviders: ["webgpu"],
});
console.log("WebGPU session created successfully!");
epUsed = "webgpu";
} catch (webgpuErr) {
console.warn("WebGPU EP failed or unsupported by browser, falling back to CPU WASM:", webgpuErr);
session = await ort.InferenceSession.create(data.modelBuffer, {
executionProviders: ["wasm"],
graphOptimizationLevel: "all",
});
epUsed = "wasm (CPU fallback)";
}
} else {
session = await ort.InferenceSession.create(data.modelBuffer, {
executionProviders: ["wasm"],
graphOptimizationLevel: "all",
});
epUsed = "wasm";
}
updateModelDimensions(session);
if (h) { try { h.dispose(); } catch (e) {} h = null; }
if (c) { try { c.dispose(); } catch (e) {} c = null; }
self.postMessage({
action: "initialized",
activeModelName: activeModelName,
vocabSize: vocabSize,
epUsed: epUsed
});
} catch (err) {
console.error("Worker initialization failed:", err);
self.postMessage({ action: "error", message: `Init failed: ${err.message}` });
}
break;
case "start":
try {
temperature = data.temperature;
topK = data.topK;
currentBpm = data.bpm;
playingStyle = data.playingStyle || "default";
midiTokens = data.midiTokens;
midiStartBar = data.midiStartBar;
midiBpmMode = data.midiBpmMode || "lock";
parser = makeParser();
parser.reset();
await warmPrompt();
const initialBatch = await pumpTokens(96);
self.postMessage({
action: "started",
events: initialBatch.events,
lastToken: initialBatch.lastToken
});
} catch (err) {
console.error("Worker start failed:", err);
self.postMessage({ action: "error", message: `Start failed: ${err.message}` });
}
break;
case "pump":
try {
temperature = data.temperature;
topK = data.topK;
const batch = await pumpTokens(96);
self.postMessage({
action: "events",
events: batch.events,
lastToken: batch.lastToken
});
} catch (err) {
console.error("Worker pump failed:", err);
self.postMessage({ action: "error", message: `Pump failed: ${err.message}` });
}
break;
case "stop":
// Halts ongoing actions, resets states if needed
break;
}
};