@tryhamster/gerbil 1.3.1 → 1.4.0
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/dist/architectures-Ol7yUmKf.mjs.map +1 -1
- package/dist/browser/index.d.ts +7 -1
- package/dist/browser/index.d.ts.map +1 -1
- package/dist/cli.mjs +7 -7
- package/dist/cli.mjs.map +1 -1
- package/dist/{defaults-CudhglS1.mjs → defaults-KnrRSIeJ.mjs} +20 -7
- package/dist/defaults-KnrRSIeJ.mjs.map +1 -0
- package/dist/frameworks/express.d.mts +2 -2
- package/dist/frameworks/express.mjs +1 -1
- package/dist/frameworks/fastify.d.mts +2 -2
- package/dist/frameworks/fastify.mjs +1 -1
- package/dist/frameworks/hono.d.mts +2 -2
- package/dist/frameworks/hono.mjs +1 -1
- package/dist/frameworks/next.d.mts +4 -4
- package/dist/frameworks/next.mjs +1 -1
- package/dist/frameworks/react.d.mts +2 -2
- package/dist/frameworks/trpc.d.mts +2 -2
- package/dist/frameworks/trpc.mjs +1 -1
- package/dist/gerbil-CUzVDAU7.mjs +4 -0
- package/dist/{gerbil-Du-1M2Wq.mjs → gerbil-DGrCwGtJ.mjs} +26 -3
- package/dist/{gerbil-Du-1M2Wq.mjs.map → gerbil-DGrCwGtJ.mjs.map} +1 -1
- package/dist/{gerbil-4Xid85_Z.d.mts → gerbil-_68_JSGV.d.mts} +15 -4
- package/dist/gerbil-_68_JSGV.d.mts.map +1 -0
- package/dist/gpu/hooks.d.mts +21 -16
- package/dist/gpu/hooks.d.mts.map +1 -1
- package/dist/gpu/hooks.mjs +20 -15
- package/dist/gpu/hooks.mjs.map +1 -1
- package/dist/gpu/index.d.mts +3 -3
- package/dist/gpu/index.mjs +4 -4
- package/dist/{gpu-BECDib8p.mjs → gpu-DbB321DN.mjs} +1530 -235
- package/dist/gpu-DbB321DN.mjs.map +1 -0
- package/dist/{index-C4kJrqs1.d.mts → index-BdbeBDTz.d.mts} +210 -9
- package/dist/index-BdbeBDTz.d.mts.map +1 -0
- package/dist/index.d.mts +4 -4
- package/dist/index.mjs +5 -5
- package/dist/integrations/ai-sdk.d.mts +2 -2
- package/dist/integrations/ai-sdk.mjs +1 -1
- package/dist/integrations/langchain.d.mts +2 -2
- package/dist/integrations/langchain.mjs +1 -1
- package/dist/integrations/llamaindex.d.mts +2 -2
- package/dist/integrations/llamaindex.mjs +1 -1
- package/dist/integrations/mcp.d.mts +4 -4
- package/dist/integrations/mcp.mjs +4 -4
- package/dist/{mcp-4imtBQLX.mjs → mcp-DcE3wn-L.mjs} +3 -3
- package/dist/{mcp-4imtBQLX.mjs.map → mcp-DcE3wn-L.mjs.map} +1 -1
- package/dist/{moonshine-stt-BO5gZ_D6.mjs → moonshine-stt-BGmctcdO.mjs} +555 -7
- package/dist/moonshine-stt-BGmctcdO.mjs.map +1 -0
- package/dist/moonshine-stt-QoNJeCzR.mjs +4 -0
- package/dist/{one-liner-X5vMbWzs.mjs → one-liner-DTGg6zdm.mjs} +2 -2
- package/dist/{one-liner-X5vMbWzs.mjs.map → one-liner-DTGg6zdm.mjs.map} +1 -1
- package/dist/{outetts-DmR0NfSG.d.mts → outetts-Cr5oGk4z.d.mts} +9 -6
- package/dist/{outetts-DmR0NfSG.d.mts.map → outetts-Cr5oGk4z.d.mts.map} +1 -1
- package/dist/{repl-CGadrKcY.mjs → repl-Ci0zXCT_.mjs} +3 -3
- package/dist/skills/index.d.mts +6 -6
- package/dist/skills/index.mjs +3 -3
- package/dist/{skills-BFGAXYkm.mjs → skills-Dsf0op66.mjs} +2 -2
- package/dist/{skills-BFGAXYkm.mjs.map → skills-Dsf0op66.mjs.map} +1 -1
- package/dist/{types-Cpf-aOxX.d.mts → types-CjG36Vzj.d.mts} +9 -3
- package/dist/{types-Cpf-aOxX.d.mts.map → types-CjG36Vzj.d.mts.map} +1 -1
- package/package.json +1 -1
- package/dist/defaults-CudhglS1.mjs.map +0 -1
- package/dist/gerbil-4Xid85_Z.d.mts.map +0 -1
- package/dist/gerbil-BikD0xFU.mjs +0 -4
- package/dist/gpu-BECDib8p.mjs.map +0 -1
- package/dist/index-C4kJrqs1.d.mts.map +0 -1
- package/dist/moonshine-stt-BO5gZ_D6.mjs.map +0 -1
- package/dist/moonshine-stt-fVw-UB4s.mjs +0 -4
|
@@ -1,6 +1,6 @@
|
|
|
1
|
-
import { i as resolveDefaultRepo, n as OUTETTS_ASSETS, r as OUTETTS_PRESET_VOICES, t as DEFAULT_MODELS } from "./defaults-
|
|
1
|
+
import { i as resolveDefaultRepo, n as OUTETTS_ASSETS, r as OUTETTS_PRESET_VOICES, t as DEFAULT_MODELS } from "./defaults-KnrRSIeJ.mjs";
|
|
2
2
|
import { E as GEMMA4_VIS_KEYS, S as DEFAULT_GROUP_SIZE, T as DTYPE_BYTES, _ as kaniCosTensor, c as KANI_END_OF_HUMAN, d as buildKaniLayerCosSin, f as computeKaniPositions, g as kaniAttentionLayerIndices, h as generateNanoCodecDecoderGraph, l as KANI_START_OF_HUMAN, m as generateKaniTtsGraph, u as audioTokensToCodes, v as kaniLayerAlpha, w as CANONICAL_KEYS, x as parseKaniConfig, y as kaniSinTensor } from "./architectures-Ol7yUmKf.mjs";
|
|
3
|
-
import {
|
|
3
|
+
import { S as verifyGPU, _ as createStorageBuffer, b as getOrCreatePipeline, c as quantizeBackboneInt4, f as Executor, g as createBindGroup, h as clearPipelineCache, i as loadModel, l as quantizeKaniBackbone, m as MATMUL_BIAS_F16C_SPEC, o as loadOuteTTS, p as KERNEL_REGISTRY, r as loadKaniTTS, s as loadParlerTTS, u as remapPrunedToken, v as createUniformBuffer, x as initGPU, y as destroyBuffers } from "./moonshine-stt-BGmctcdO.mjs";
|
|
4
4
|
|
|
5
5
|
//#region src/gpu/architectures/gemma4_vision.ts
|
|
6
6
|
/**
|
|
@@ -767,7 +767,7 @@ const KANI_CODEC_LOOKBACK_FRAMES = 32;
|
|
|
767
767
|
/** Hard cap so a stuck decode cannot loop forever (reference uses 3000). */
|
|
768
768
|
const DEFAULT_MAX_NEW_TOKENS$1 = 3e3;
|
|
769
769
|
/** Build a fresh map of just the constant weights a graph references (see moonshine-stt). */
|
|
770
|
-
function selectGraphWeights$
|
|
770
|
+
function selectGraphWeights$2(graph, weights) {
|
|
771
771
|
const out = /* @__PURE__ */ new Map();
|
|
772
772
|
for (const [name, desc] of Object.entries(graph.tensors)) {
|
|
773
773
|
if (desc.storage !== "constant") continue;
|
|
@@ -795,10 +795,10 @@ function nucleusSample$1(probs, ids, topP) {
|
|
|
795
795
|
}
|
|
796
796
|
let keepSum = 0;
|
|
797
797
|
for (let r = 0; r < cutoff; r++) keepSum += probs[order[r]];
|
|
798
|
-
let pick = Math.random() * keepSum;
|
|
798
|
+
let pick$1 = Math.random() * keepSum;
|
|
799
799
|
for (let r = 0; r < cutoff; r++) {
|
|
800
|
-
pick -= probs[order[r]];
|
|
801
|
-
if (pick <= 0) return ids[order[r]];
|
|
800
|
+
pick$1 -= probs[order[r]];
|
|
801
|
+
if (pick$1 <= 0) return ids[order[r]];
|
|
802
802
|
}
|
|
803
803
|
return ids[order[0]];
|
|
804
804
|
}
|
|
@@ -891,7 +891,7 @@ var KaniTTS = class KaniTTS {
|
|
|
891
891
|
maxSeqLen,
|
|
892
892
|
kvMode: "f32"
|
|
893
893
|
});
|
|
894
|
-
exec.uploadWeightsMap(selectGraphWeights$
|
|
894
|
+
exec.uploadWeightsMap(selectGraphWeights$2(graph, loaded.backboneWeights));
|
|
895
895
|
exec.initBindGroups();
|
|
896
896
|
return exec;
|
|
897
897
|
}
|
|
@@ -1105,7 +1105,7 @@ var KaniTTS = class KaniTTS {
|
|
|
1105
1105
|
kvMode: "f32"
|
|
1106
1106
|
});
|
|
1107
1107
|
try {
|
|
1108
|
-
exec.uploadWeightsMap(selectGraphWeights$
|
|
1108
|
+
exec.uploadWeightsMap(selectGraphWeights$2(graph, this.loaded.codecWeights));
|
|
1109
1109
|
exec.initBindGroups();
|
|
1110
1110
|
exec.reset();
|
|
1111
1111
|
exec.writeInput("audio_codes", winCodes);
|
|
@@ -1212,6 +1212,10 @@ function dacKey(name) {
|
|
|
1212
1212
|
}
|
|
1213
1213
|
function generateDacSpeechDecoderGraph(opts) {
|
|
1214
1214
|
const T = opts.numFrames;
|
|
1215
|
+
const NCB = opts.numCodebooks ?? DAC_NUM_CODEBOOKS;
|
|
1216
|
+
const rates = opts.decoderRates ?? DAC_DECODER_RATES;
|
|
1217
|
+
const decoderDim = opts.decoderDim ?? DAC_DECODER_DIM;
|
|
1218
|
+
const latentDim = opts.latentDim ?? DAC_LATENT_DIM;
|
|
1215
1219
|
const tensors = {};
|
|
1216
1220
|
const nodes = [];
|
|
1217
1221
|
const executionOrder = [];
|
|
@@ -1224,14 +1228,14 @@ function generateDacSpeechDecoderGraph(opts) {
|
|
|
1224
1228
|
};
|
|
1225
1229
|
addTensor({
|
|
1226
1230
|
name: "audio_codes",
|
|
1227
|
-
shape: [
|
|
1231
|
+
shape: [NCB, T],
|
|
1228
1232
|
dtype: "u32",
|
|
1229
1233
|
storage: "activation"
|
|
1230
1234
|
});
|
|
1231
1235
|
addTensor({
|
|
1232
1236
|
name: dacKey("quantizer.codebook"),
|
|
1233
1237
|
shape: [
|
|
1234
|
-
|
|
1238
|
+
NCB,
|
|
1235
1239
|
DAC_CODEBOOK_SIZE,
|
|
1236
1240
|
DAC_CODEBOOK_DIM
|
|
1237
1241
|
],
|
|
@@ -1242,8 +1246,8 @@ function generateDacSpeechDecoderGraph(opts) {
|
|
|
1242
1246
|
addTensor({
|
|
1243
1247
|
name: dacKey("quantizer.out_proj.weight"),
|
|
1244
1248
|
shape: [
|
|
1245
|
-
|
|
1246
|
-
|
|
1249
|
+
NCB,
|
|
1250
|
+
latentDim,
|
|
1247
1251
|
DAC_CODEBOOK_DIM
|
|
1248
1252
|
],
|
|
1249
1253
|
dtype: "f32",
|
|
@@ -1252,14 +1256,14 @@ function generateDacSpeechDecoderGraph(opts) {
|
|
|
1252
1256
|
});
|
|
1253
1257
|
addTensor({
|
|
1254
1258
|
name: dacKey("quantizer.out_proj.bias"),
|
|
1255
|
-
shape: [
|
|
1259
|
+
shape: [NCB, latentDim],
|
|
1256
1260
|
dtype: "f32",
|
|
1257
1261
|
storage: "constant",
|
|
1258
1262
|
safetensorsKey: dacKey("quantizer.out_proj.bias")
|
|
1259
1263
|
});
|
|
1260
1264
|
addTensor({
|
|
1261
1265
|
name: "dac_latent",
|
|
1262
|
-
shape: [
|
|
1266
|
+
shape: [latentDim, T],
|
|
1263
1267
|
dtype: "f32",
|
|
1264
1268
|
storage: "activation"
|
|
1265
1269
|
});
|
|
@@ -1274,8 +1278,8 @@ function generateDacSpeechDecoderGraph(opts) {
|
|
|
1274
1278
|
],
|
|
1275
1279
|
outputs: ["dac_latent"],
|
|
1276
1280
|
attributes: {
|
|
1277
|
-
NCB
|
|
1278
|
-
Cout:
|
|
1281
|
+
NCB,
|
|
1282
|
+
Cout: latentDim,
|
|
1279
1283
|
T,
|
|
1280
1284
|
D: DAC_CODEBOOK_DIM,
|
|
1281
1285
|
codebook_size: DAC_CODEBOOK_SIZE
|
|
@@ -1425,11 +1429,11 @@ function generateDacSpeechDecoderGraph(opts) {
|
|
|
1425
1429
|
});
|
|
1426
1430
|
}
|
|
1427
1431
|
let cur = "dac_conv1_out";
|
|
1428
|
-
let length = conv1dFull("dac_conv1", "dac_latent", dacKey("decoder.model.0.weight"), dacKey("decoder.model.0.bias"),
|
|
1429
|
-
let channels =
|
|
1430
|
-
for (let stage = 0; stage <
|
|
1432
|
+
let length = conv1dFull("dac_conv1", "dac_latent", dacKey("decoder.model.0.weight"), dacKey("decoder.model.0.bias"), latentDim, decoderDim, 7, T, 1, 3, cur);
|
|
1433
|
+
let channels = decoderDim;
|
|
1434
|
+
for (let stage = 0; stage < rates.length; stage++) {
|
|
1431
1435
|
const m = stage + 1;
|
|
1432
|
-
const rate =
|
|
1436
|
+
const rate = rates[stage];
|
|
1433
1437
|
const outCh = channels / 2;
|
|
1434
1438
|
const upK = 2 * rate;
|
|
1435
1439
|
const upPad = Math.ceil(rate / 2);
|
|
@@ -1458,15 +1462,17 @@ function generateDacSpeechDecoderGraph(opts) {
|
|
|
1458
1462
|
}
|
|
1459
1463
|
cur = h;
|
|
1460
1464
|
}
|
|
1465
|
+
const finalSnakeIdx = rates.length + 1;
|
|
1466
|
+
const finalConvIdx = rates.length + 2;
|
|
1461
1467
|
const postSnake = "dac_post_snake";
|
|
1462
|
-
snake("dac_post_snake_op", cur, channels, length, dacKey(
|
|
1468
|
+
snake("dac_post_snake_op", cur, channels, length, dacKey(`decoder.model.${finalSnakeIdx}.alpha`), postSnake);
|
|
1463
1469
|
const pcm = "dac_pcm";
|
|
1464
|
-
const pcmLen = conv1dFull("dac_conv2", postSnake, dacKey(
|
|
1470
|
+
const pcmLen = conv1dFull("dac_conv2", postSnake, dacKey(`decoder.model.${finalConvIdx}.weight`), dacKey(`decoder.model.${finalConvIdx}.bias`), channels, 1, 7, length, 1, 3, pcm);
|
|
1465
1471
|
return {
|
|
1466
1472
|
architecture: "DacSpeechDecoder",
|
|
1467
1473
|
config: {
|
|
1468
|
-
hidden_size:
|
|
1469
|
-
num_layers:
|
|
1474
|
+
hidden_size: decoderDim,
|
|
1475
|
+
num_layers: rates.length,
|
|
1470
1476
|
num_heads: 0,
|
|
1471
1477
|
num_kv_heads: 0,
|
|
1472
1478
|
head_dim: 0,
|
|
@@ -1493,10 +1499,13 @@ function generateDacSpeechDecoderGraph(opts) {
|
|
|
1493
1499
|
outputs: [pcm]
|
|
1494
1500
|
};
|
|
1495
1501
|
}
|
|
1496
|
-
/**
|
|
1497
|
-
|
|
1502
|
+
/**
|
|
1503
|
+
* Final decoded PCM length (samples) for a T-frame code grid, via the DAC stage chain.
|
|
1504
|
+
* `rates` defaults to DAC.speech.v1.0 ([8,5,4,2]); pass [8,8,4,2] for Parler's dac_44khz.
|
|
1505
|
+
*/
|
|
1506
|
+
function dacOutputLength(numFrames, rates = DAC_DECODER_RATES) {
|
|
1498
1507
|
let length = numFrames;
|
|
1499
|
-
for (const rate of
|
|
1508
|
+
for (const rate of rates) length = convTransposeLen(length, rate, 2 * rate, Math.ceil(rate / 2), 0);
|
|
1500
1509
|
return length;
|
|
1501
1510
|
}
|
|
1502
1511
|
/** Strided/dilated Conv1d output length: floor((L + 2*pad - dilation*(K-1) - 1)/stride) + 1. */
|
|
@@ -2253,13 +2262,30 @@ function generateOuteTtsBackboneGraph(rawConfig, dtype, groupSize, kvDtype) {
|
|
|
2253
2262
|
* builder emits match the reference speaker JSONs the presets ship.
|
|
2254
2263
|
*
|
|
2255
2264
|
* The DAC encode (audio → c1/c2 codes) is done bit-exact on the GPU
|
|
2256
|
-
* (generateDacSpeechEncoderGraph + RVQEncode, validated 450/450 vs torch). The
|
|
2257
|
-
*
|
|
2258
|
-
* the code stream per word
|
|
2259
|
-
*
|
|
2260
|
-
*
|
|
2261
|
-
*
|
|
2262
|
-
*
|
|
2265
|
+
* (generateDacSpeechEncoderGraph + RVQEncode, validated 450/450 vs torch). The codes
|
|
2266
|
+
* are EXACT; only the word→frame boundaries are estimated. outetts uses Whisper
|
|
2267
|
+
* word-level timestamps to slice the code stream per word. We provide two on-device
|
|
2268
|
+
* estimators (both leave the codes untouched — they only decide which frames map to
|
|
2269
|
+
* which word, and which audio span each word's prosody is computed over):
|
|
2270
|
+
*
|
|
2271
|
+
* 1. alignWordsByEnergy (default) — a forced-segmentation aligner driven by the
|
|
2272
|
+
* clip's OWN short-time energy envelope. It finds the silence gaps (low-energy
|
|
2273
|
+
* runs) between utterances and snaps word boundaries onto them, distributing the
|
|
2274
|
+
* transcript words across the detected speech runs (char-weighted within a run).
|
|
2275
|
+
* This is audio-driven: boundaries land on real pauses instead of being guessed.
|
|
2276
|
+
*
|
|
2277
|
+
* 2. alignWordsProportional (fallback) — the original audio-BLIND estimator that
|
|
2278
|
+
* spreads frames across words proportionally to char length. Used when the energy
|
|
2279
|
+
* envelope yields no usable structure (one continuous speech run, or fewer gaps
|
|
2280
|
+
* than word boundaries).
|
|
2281
|
+
*
|
|
2282
|
+
* On-device STT note: Gerbil's Moonshine STT (MoonshineSTT, src/gpu/moonshine-stt.ts)
|
|
2283
|
+
* is an encoder-decoder ASR that emits a transcript but does NOT expose Whisper-style
|
|
2284
|
+
* word timestamps — its decoder cross-attention scores are consumed inside the kernel
|
|
2285
|
+
* and never read back, so there is no per-token→frame map to slice on. Energy-based
|
|
2286
|
+
* forced segmentation is therefore the closest achievable on-device alignment; STT is
|
|
2287
|
+
* used (when wired by the caller) to obtain/confirm the WORD SEQUENCE, and the energy
|
|
2288
|
+
* aligner places the boundaries. See `OuteTTS.encodeSpeaker`.
|
|
2263
2289
|
*/
|
|
2264
2290
|
const DAC_TPS = 75;
|
|
2265
2291
|
const WHITESPACE_RE = /\s+/;
|
|
@@ -2422,6 +2448,146 @@ function alignWordsProportional(transcript, numFrames) {
|
|
|
2422
2448
|
}
|
|
2423
2449
|
return out;
|
|
2424
2450
|
}
|
|
2451
|
+
/**
|
|
2452
|
+
* Short-time RMS energy envelope sampled at the DAC frame rate (one value per code
|
|
2453
|
+
* frame). frame f covers audio samples [f·hop, f·hop + win). Used by the energy
|
|
2454
|
+
* aligner to locate silence gaps. `hop` is DAC_HOP (320 samples ≈ 13.3 ms @ 24 kHz),
|
|
2455
|
+
* matching the DAC frame stride so envelope index == code-frame index.
|
|
2456
|
+
*/
|
|
2457
|
+
function frameEnergyEnvelope(audio, hop, numFrames) {
|
|
2458
|
+
const env = new Float32Array(numFrames);
|
|
2459
|
+
const win = hop;
|
|
2460
|
+
for (let f = 0; f < numFrames; f++) {
|
|
2461
|
+
const start = f * hop;
|
|
2462
|
+
const end = Math.min(start + win, audio.length);
|
|
2463
|
+
let sum = 0;
|
|
2464
|
+
let n = 0;
|
|
2465
|
+
for (let i = start; i < end; i++) {
|
|
2466
|
+
sum += audio[i] * audio[i];
|
|
2467
|
+
n++;
|
|
2468
|
+
}
|
|
2469
|
+
env[f] = n > 0 ? Math.sqrt(sum / n) : 0;
|
|
2470
|
+
}
|
|
2471
|
+
return env;
|
|
2472
|
+
}
|
|
2473
|
+
/**
|
|
2474
|
+
* Segment the energy envelope into speech runs separated by silence gaps. A frame is
|
|
2475
|
+
* "silent" when its energy is below `floor + (peak-floor)*silenceRatio`; a run of
|
|
2476
|
+
* silent frames at least `minGapFrames` long is treated as a word/phrase boundary.
|
|
2477
|
+
* Leading/trailing silence is trimmed. Returns the speech runs in order.
|
|
2478
|
+
*/
|
|
2479
|
+
function detectSpeechRuns(env, silenceRatio, minGapFrames) {
|
|
2480
|
+
const numFrames = env.length;
|
|
2481
|
+
if (numFrames === 0) return [];
|
|
2482
|
+
let peak = 0;
|
|
2483
|
+
let floor = Number.POSITIVE_INFINITY;
|
|
2484
|
+
for (let f = 0; f < numFrames; f++) {
|
|
2485
|
+
peak = Math.max(peak, env[f]);
|
|
2486
|
+
floor = Math.min(floor, env[f]);
|
|
2487
|
+
}
|
|
2488
|
+
if (peak <= 0) return [];
|
|
2489
|
+
const threshold = floor + (peak - floor) * silenceRatio;
|
|
2490
|
+
const runs = [];
|
|
2491
|
+
let runStart = -1;
|
|
2492
|
+
let silentRun = 0;
|
|
2493
|
+
for (let f = 0; f < numFrames; f++) if (env[f] > threshold) {
|
|
2494
|
+
if (runStart < 0) runStart = f;
|
|
2495
|
+
silentRun = 0;
|
|
2496
|
+
} else {
|
|
2497
|
+
silentRun++;
|
|
2498
|
+
if (runStart >= 0 && silentRun >= minGapFrames) {
|
|
2499
|
+
runs.push({
|
|
2500
|
+
start: runStart,
|
|
2501
|
+
end: f - silentRun + 1
|
|
2502
|
+
});
|
|
2503
|
+
runStart = -1;
|
|
2504
|
+
}
|
|
2505
|
+
}
|
|
2506
|
+
if (runStart >= 0) runs.push({
|
|
2507
|
+
start: runStart,
|
|
2508
|
+
end: numFrames
|
|
2509
|
+
});
|
|
2510
|
+
return runs.filter((r) => r.end > r.start);
|
|
2511
|
+
}
|
|
2512
|
+
const ENERGY_SILENCE_RATIO = .15;
|
|
2513
|
+
const ENERGY_MIN_GAP_FRAMES = 4;
|
|
2514
|
+
/**
|
|
2515
|
+
* Distribute `numWords` whole words across speech runs proportionally to each run's
|
|
2516
|
+
* frame length, giving every run ≥1 word and summing to exactly `numWords`. Surplus
|
|
2517
|
+
* (from the per-run floor) is handed to the longest runs first; any overshoot is taken
|
|
2518
|
+
* back from the shortest runs (never below 1). Returns the per-run word count.
|
|
2519
|
+
*/
|
|
2520
|
+
function distributeWordCounts(runLens, numWords) {
|
|
2521
|
+
const totalRunLen = runLens.reduce((a, b) => a + b, 0) || 1;
|
|
2522
|
+
const targets = runLens.map((len) => Math.max(1, Math.floor(len / totalRunLen * numWords)));
|
|
2523
|
+
const order = runLens.map((_, i) => i).sort((a, b) => runLens[b] - runLens[a]);
|
|
2524
|
+
let assigned = targets.reduce((a, b) => a + b, 0);
|
|
2525
|
+
let oi = 0;
|
|
2526
|
+
while (assigned < numWords) {
|
|
2527
|
+
targets[order[oi % order.length]]++;
|
|
2528
|
+
assigned++;
|
|
2529
|
+
oi++;
|
|
2530
|
+
}
|
|
2531
|
+
while (assigned > numWords) {
|
|
2532
|
+
const idx = order[order.length - 1 - oi % order.length];
|
|
2533
|
+
if (targets[idx] > 1) {
|
|
2534
|
+
targets[idx]--;
|
|
2535
|
+
assigned--;
|
|
2536
|
+
}
|
|
2537
|
+
oi++;
|
|
2538
|
+
}
|
|
2539
|
+
return targets;
|
|
2540
|
+
}
|
|
2541
|
+
/**
|
|
2542
|
+
* Split a contiguous frame span [spanStart, spanEnd) among `groupWords` char-
|
|
2543
|
+
* proportionally, appending the resulting {@link AlignedWord}s to `out`. Boundaries
|
|
2544
|
+
* are contiguous within the span (the last word reaches spanEnd exactly).
|
|
2545
|
+
*/
|
|
2546
|
+
function splitRunWords(groupWords, spanStart, spanEnd, out) {
|
|
2547
|
+
const span = Math.max(1, spanEnd - spanStart);
|
|
2548
|
+
const chars = groupWords.map((w) => w.length || 1);
|
|
2549
|
+
const groupTotal = chars.reduce((a, b) => a + b, 0) || groupWords.length;
|
|
2550
|
+
let acc = 0;
|
|
2551
|
+
for (let g = 0; g < groupWords.length; g++) {
|
|
2552
|
+
const wStart = g === 0 ? spanStart : spanStart + Math.round(acc / groupTotal * span);
|
|
2553
|
+
acc += chars[g];
|
|
2554
|
+
const wEnd = g === groupWords.length - 1 ? spanEnd : spanStart + Math.round(acc / groupTotal * span);
|
|
2555
|
+
out.push({
|
|
2556
|
+
word: groupWords[g],
|
|
2557
|
+
startFrame: wStart,
|
|
2558
|
+
endFrame: Math.max(wEnd, wStart)
|
|
2559
|
+
});
|
|
2560
|
+
}
|
|
2561
|
+
}
|
|
2562
|
+
/**
|
|
2563
|
+
* Forced word→frame alignment driven by the clip's energy envelope. Detects the
|
|
2564
|
+
* silence gaps in `audio` (at DAC frame resolution) and distributes the transcript's
|
|
2565
|
+
* words across the detected speech runs so that word boundaries land on real pauses.
|
|
2566
|
+
* Within a run, words are split char-proportionally (same local rule as the fallback).
|
|
2567
|
+
* Run spans are made contiguous across the silence gaps (a run owns the frames up to
|
|
2568
|
+
* the next run's start) so the per-word c1/c2 slices still tile the full [0, numFrames)
|
|
2569
|
+
* code stream with no dropped frames.
|
|
2570
|
+
*
|
|
2571
|
+
* Falls back to {@link alignWordsProportional} when the envelope has no usable
|
|
2572
|
+
* structure (≤1 speech run). The DAC codes are never touched; only the boundaries move.
|
|
2573
|
+
*/
|
|
2574
|
+
function alignWordsByEnergy(transcript, numFrames, audio, hop) {
|
|
2575
|
+
const words = transcript.split(WHITESPACE_RE).filter((w) => w.length > 0);
|
|
2576
|
+
if (words.length === 0) return [];
|
|
2577
|
+
if (words.length === 1 || numFrames <= 0) return alignWordsProportional(transcript, numFrames);
|
|
2578
|
+
const runs = detectSpeechRuns(frameEnergyEnvelope(audio, hop, numFrames), ENERGY_SILENCE_RATIO, ENERGY_MIN_GAP_FRAMES);
|
|
2579
|
+
if (runs.length < 2) return alignWordsProportional(transcript, numFrames);
|
|
2580
|
+
const targets = distributeWordCounts(runs.map((r) => r.end - r.start), words.length);
|
|
2581
|
+
const out = [];
|
|
2582
|
+
let wordIdx = 0;
|
|
2583
|
+
for (let r = 0; r < runs.length; r++) {
|
|
2584
|
+
const spanStart = r === 0 ? 0 : runs[r].start;
|
|
2585
|
+
const spanEnd = r === runs.length - 1 ? numFrames : runs[r + 1].start;
|
|
2586
|
+
splitRunWords(words.slice(wordIdx, wordIdx + targets[r]), spanStart, spanEnd, out);
|
|
2587
|
+
wordIdx += targets[r];
|
|
2588
|
+
}
|
|
2589
|
+
return out;
|
|
2590
|
+
}
|
|
2425
2591
|
|
|
2426
2592
|
//#endregion
|
|
2427
2593
|
//#region src/gpu/outetts.ts
|
|
@@ -2581,7 +2747,7 @@ async function loadOuteSpeaker(voiceOrUrl, baseURL) {
|
|
|
2581
2747
|
return await resp.json();
|
|
2582
2748
|
}
|
|
2583
2749
|
/** Build a fresh map of just the constant weights a graph references (see kani-tts). */
|
|
2584
|
-
function selectGraphWeights(graph, weights) {
|
|
2750
|
+
function selectGraphWeights$1(graph, weights) {
|
|
2585
2751
|
const out = /* @__PURE__ */ new Map();
|
|
2586
2752
|
for (const [name, desc] of Object.entries(graph.tensors)) {
|
|
2587
2753
|
if (desc.storage !== "constant") continue;
|
|
@@ -2591,7 +2757,7 @@ function selectGraphWeights(graph, weights) {
|
|
|
2591
2757
|
return out;
|
|
2592
2758
|
}
|
|
2593
2759
|
/** Index of the maximum logit (greedy / argmax decode). */
|
|
2594
|
-
function argmax$
|
|
2760
|
+
function argmax$2(logits) {
|
|
2595
2761
|
let best = 0;
|
|
2596
2762
|
let bestV = logits[0];
|
|
2597
2763
|
for (let i = 1; i < logits.length; i++) if (logits[i] > bestV) {
|
|
@@ -2637,10 +2803,10 @@ function nucleusSample(probs, ids, topP) {
|
|
|
2637
2803
|
}
|
|
2638
2804
|
let keepSum = 0;
|
|
2639
2805
|
for (let r = 0; r < cutoff; r++) keepSum += probs[order[r]];
|
|
2640
|
-
let pick = Math.random() * keepSum;
|
|
2806
|
+
let pick$1 = Math.random() * keepSum;
|
|
2641
2807
|
for (let r = 0; r < cutoff; r++) {
|
|
2642
|
-
pick -= probs[order[r]];
|
|
2643
|
-
if (pick <= 0) return ids[order[r]];
|
|
2808
|
+
pick$1 -= probs[order[r]];
|
|
2809
|
+
if (pick$1 <= 0) return ids[order[r]];
|
|
2644
2810
|
}
|
|
2645
2811
|
return ids[order[0]];
|
|
2646
2812
|
}
|
|
@@ -2690,7 +2856,7 @@ var OuteTTS = class OuteTTS {
|
|
|
2690
2856
|
maxSeqLen,
|
|
2691
2857
|
kvMode: "f32"
|
|
2692
2858
|
});
|
|
2693
|
-
exec.uploadWeightsMap(selectGraphWeights(graph, loaded.backboneWeights));
|
|
2859
|
+
exec.uploadWeightsMap(selectGraphWeights$1(graph, loaded.backboneWeights));
|
|
2694
2860
|
exec.initBindGroups();
|
|
2695
2861
|
this.backboneExec = exec;
|
|
2696
2862
|
}
|
|
@@ -2765,11 +2931,14 @@ var OuteTTS = class OuteTTS {
|
|
|
2765
2931
|
* Fidelity vs outetts `create_speaker`: the DAC-encode (audio → c1/c2 codes) and the
|
|
2766
2932
|
* energy/spectral_centroid/pitch feature math are a faithful port. Two documented
|
|
2767
2933
|
* approximations: (a) word→frame ALIGNMENT — outetts uses Whisper word timestamps to
|
|
2768
|
-
* slice the code stream
|
|
2769
|
-
*
|
|
2770
|
-
* {@link
|
|
2771
|
-
*
|
|
2772
|
-
* and
|
|
2934
|
+
* slice the code stream. Gerbil's on-device Moonshine STT does not expose word
|
|
2935
|
+
* timestamps (its decoder cross-attention is never read back), so we instead run an
|
|
2936
|
+
* energy-forced aligner ({@link alignWordsByEnergy}) that snaps word boundaries onto
|
|
2937
|
+
* the clip's own silence gaps — audio-driven, tighter than the old char-proportional
|
|
2938
|
+
* guess, and exact when each word is pause-separated; it falls back to proportional
|
|
2939
|
+
* for continuous speech. (b) we skip outetts' pyloudnorm pre-normalization (it mainly
|
|
2940
|
+
* rescales absolute level / the energy bucket). The cloned voice timbre and prosody
|
|
2941
|
+
* come through; only the per-word boundaries remain an estimate.
|
|
2773
2942
|
*/
|
|
2774
2943
|
async encodeSpeaker(ref) {
|
|
2775
2944
|
if ("words" in ref && "text" in ref) return ref;
|
|
@@ -2785,7 +2954,7 @@ var OuteTTS = class OuteTTS {
|
|
|
2785
2954
|
const codes = await this.encodeAudioToCodes(pcm, numFrames);
|
|
2786
2955
|
const c1All = codes.subarray(0, numFrames);
|
|
2787
2956
|
const c2All = codes.subarray(numFrames, 2 * numFrames);
|
|
2788
|
-
const words =
|
|
2957
|
+
const words = alignWordsByEnergy(transcript, numFrames, mono, DAC_HOP).map((a) => {
|
|
2789
2958
|
const startSample = Math.min(a.startFrame * DAC_HOP, mono.length);
|
|
2790
2959
|
const endSample = Math.min(a.endFrame * DAC_HOP, mono.length);
|
|
2791
2960
|
const features = extractAudioFeatures(mono.subarray(startSample, endSample), DAC_SAMPLE_RATE);
|
|
@@ -2810,7 +2979,7 @@ var OuteTTS = class OuteTTS {
|
|
|
2810
2979
|
kvMode: "f32"
|
|
2811
2980
|
});
|
|
2812
2981
|
try {
|
|
2813
|
-
exec.uploadWeightsMap(selectGraphWeights(graph, this.loaded.codecWeights));
|
|
2982
|
+
exec.uploadWeightsMap(selectGraphWeights$1(graph, this.loaded.codecWeights));
|
|
2814
2983
|
exec.initBindGroups();
|
|
2815
2984
|
exec.reset();
|
|
2816
2985
|
exec.writeInput("pcm_in", pcm);
|
|
@@ -2852,7 +3021,7 @@ var OuteTTS = class OuteTTS {
|
|
|
2852
3021
|
const generated = [];
|
|
2853
3022
|
let logits = prefillLogits;
|
|
2854
3023
|
for (let step = 0; step < p.maxNewTokens; step++) {
|
|
2855
|
-
const next = p.greedy ? argmax$
|
|
3024
|
+
const next = p.greedy ? argmax$2(logits) : this.sampleToken(logits, {
|
|
2856
3025
|
temperature: p.temperature,
|
|
2857
3026
|
topP: p.topP,
|
|
2858
3027
|
topK: p.topK,
|
|
@@ -2903,7 +3072,7 @@ var OuteTTS = class OuteTTS {
|
|
|
2903
3072
|
kvMode: "f32"
|
|
2904
3073
|
});
|
|
2905
3074
|
try {
|
|
2906
|
-
exec.uploadWeightsMap(selectGraphWeights(graph, this.loaded.codecWeights));
|
|
3075
|
+
exec.uploadWeightsMap(selectGraphWeights$1(graph, this.loaded.codecWeights));
|
|
2907
3076
|
exec.initBindGroups();
|
|
2908
3077
|
exec.reset();
|
|
2909
3078
|
exec.writeInput("audio_codes", codes);
|
|
@@ -2925,198 +3094,1295 @@ var OuteTTS = class OuteTTS {
|
|
|
2925
3094
|
};
|
|
2926
3095
|
|
|
2927
3096
|
//#endregion
|
|
2928
|
-
//#region src/gpu/
|
|
2929
|
-
|
|
2930
|
-
|
|
3097
|
+
//#region src/gpu/architectures/parler.ts
|
|
3098
|
+
const PARLER_SAMPLE_RATE = 44100;
|
|
3099
|
+
const PARLER_NUM_CODEBOOKS = 9;
|
|
3100
|
+
const PARLER_DECODER_RATES = [
|
|
3101
|
+
8,
|
|
3102
|
+
8,
|
|
3103
|
+
4,
|
|
3104
|
+
2
|
|
3105
|
+
];
|
|
3106
|
+
const PARLER_DAC_LATENT_DIM = 1024;
|
|
3107
|
+
const PARLER_DAC_DECODER_DIM = 1536;
|
|
3108
|
+
/** Decoder audio-vocab: 1024 codes + eos(1024)+bos(1025)+pad(1024); lm_head vocab. */
|
|
3109
|
+
const PARLER_AUDIO_VOCAB = 1088;
|
|
3110
|
+
const PARLER_BOS_TOKEN_ID = 1025;
|
|
3111
|
+
const PARLER_EOS_TOKEN_ID = 1024;
|
|
3112
|
+
const PARLER_PAD_TOKEN_ID = 1024;
|
|
3113
|
+
const CAPS = {
|
|
3114
|
+
text: true,
|
|
3115
|
+
vision: false,
|
|
3116
|
+
moe: false
|
|
3117
|
+
};
|
|
3118
|
+
/** Pull Parler dims from the (nested) HF config. */
|
|
3119
|
+
function parseParlerConfig(raw) {
|
|
3120
|
+
const te = raw.text_encoder ?? {};
|
|
3121
|
+
const dec = raw.decoder ?? {};
|
|
3122
|
+
const enc_heads = te.num_heads ?? 16;
|
|
3123
|
+
const enc_d_model = te.d_model ?? 1024;
|
|
3124
|
+
const dec_heads = dec.num_attention_heads ?? 16;
|
|
3125
|
+
return {
|
|
3126
|
+
enc_d_model,
|
|
3127
|
+
enc_layers: te.num_layers ?? 24,
|
|
3128
|
+
enc_heads,
|
|
3129
|
+
enc_d_kv: te.d_kv ?? Math.floor(enc_d_model / enc_heads),
|
|
3130
|
+
enc_d_ff: te.d_ff ?? 2816,
|
|
3131
|
+
enc_vocab: te.vocab_size ?? 32128,
|
|
3132
|
+
rel_num_buckets: te.relative_attention_num_buckets ?? 32,
|
|
3133
|
+
rel_max_distance: te.relative_attention_max_distance ?? 128,
|
|
3134
|
+
ln_eps: te.layer_norm_epsilon ?? 1e-6,
|
|
3135
|
+
dec_hidden: dec.hidden_size ?? 1024,
|
|
3136
|
+
dec_layers: dec.num_hidden_layers ?? 24,
|
|
3137
|
+
dec_heads,
|
|
3138
|
+
dec_kv_heads: dec.num_key_value_heads ?? dec_heads,
|
|
3139
|
+
dec_ffn: dec.ffn_dim ?? 4096,
|
|
3140
|
+
num_codebooks: dec.num_codebooks ?? PARLER_NUM_CODEBOOKS,
|
|
3141
|
+
audio_vocab: dec.vocab_size ?? PARLER_AUDIO_VOCAB,
|
|
3142
|
+
embed_vocab: (dec.vocab_size ?? PARLER_AUDIO_VOCAB) + 1,
|
|
3143
|
+
max_position: dec.max_position_embeddings ?? 4096,
|
|
3144
|
+
bos_token_id: dec.bos_token_id ?? PARLER_BOS_TOKEN_ID,
|
|
3145
|
+
eos_token_id: dec.eos_token_id ?? PARLER_EOS_TOKEN_ID,
|
|
3146
|
+
pad_token_id: dec.pad_token_id ?? PARLER_PAD_TOKEN_ID
|
|
3147
|
+
};
|
|
3148
|
+
}
|
|
3149
|
+
function baseConfig(c, hidden, heads, layers) {
|
|
3150
|
+
return {
|
|
3151
|
+
hidden_size: hidden,
|
|
3152
|
+
num_layers: layers,
|
|
3153
|
+
num_heads: heads,
|
|
3154
|
+
num_kv_heads: heads,
|
|
3155
|
+
head_dim: Math.floor(hidden / heads),
|
|
3156
|
+
intermediate_size: c.dec_ffn,
|
|
3157
|
+
vocab_size: c.audio_vocab,
|
|
3158
|
+
context_length: c.max_position,
|
|
3159
|
+
rms_norm_eps: c.ln_eps,
|
|
3160
|
+
norm_type: "layernorm",
|
|
3161
|
+
rope_base: 0,
|
|
3162
|
+
rope_dim: 0,
|
|
3163
|
+
kv_layout: "LHSd",
|
|
3164
|
+
is_moe: false,
|
|
3165
|
+
has_vision_tower: false
|
|
3166
|
+
};
|
|
3167
|
+
}
|
|
2931
3168
|
/**
|
|
2932
|
-
*
|
|
2933
|
-
*
|
|
2934
|
-
*
|
|
3169
|
+
* T5's bidirectional relative-position bucket (verbatim port of
|
|
3170
|
+
* transformers' T5Attention._relative_position_bucket, bidirectional=True).
|
|
3171
|
+
* Maps a signed relative position (key − query) to a bucket id in [0, num_buckets).
|
|
2935
3172
|
*/
|
|
2936
|
-
function
|
|
2937
|
-
|
|
2938
|
-
|
|
2939
|
-
const
|
|
2940
|
-
|
|
2941
|
-
|
|
2942
|
-
const
|
|
2943
|
-
const
|
|
2944
|
-
|
|
2945
|
-
|
|
2946
|
-
|
|
2947
|
-
|
|
2948
|
-
const hIdx = _heapIndices;
|
|
2949
|
-
const hVal = _heapValues;
|
|
2950
|
-
let penaltySet = null;
|
|
2951
|
-
if (repetitionPenalty !== 1 && previousTokens?.length) penaltySet = new Set(previousTokens);
|
|
2952
|
-
let heapSize = 0;
|
|
2953
|
-
for (let i = 0; i < N; i++) {
|
|
2954
|
-
let s = logits[i];
|
|
2955
|
-
if (penaltySet?.has(i)) s = s > 0 ? s / repetitionPenalty : s * repetitionPenalty;
|
|
2956
|
-
s /= temperature;
|
|
2957
|
-
if (heapSize < K) {
|
|
2958
|
-
hIdx[heapSize] = i;
|
|
2959
|
-
hVal[heapSize] = s;
|
|
2960
|
-
heapSize++;
|
|
2961
|
-
if (heapSize === K) for (let j = (K >> 1) - 1; j >= 0; j--) siftDown(hIdx, hVal, j, K);
|
|
2962
|
-
} else if (s > hVal[0]) {
|
|
2963
|
-
hIdx[0] = i;
|
|
2964
|
-
hVal[0] = s;
|
|
2965
|
-
siftDown(hIdx, hVal, 0, K);
|
|
2966
|
-
}
|
|
2967
|
-
}
|
|
2968
|
-
for (let i = 1; i < heapSize; i++) {
|
|
2969
|
-
const vi = hVal[i];
|
|
2970
|
-
const ii = hIdx[i];
|
|
2971
|
-
let j = i - 1;
|
|
2972
|
-
while (j >= 0 && hVal[j] < vi) {
|
|
2973
|
-
hVal[j + 1] = hVal[j];
|
|
2974
|
-
hIdx[j + 1] = hIdx[j];
|
|
2975
|
-
j--;
|
|
2976
|
-
}
|
|
2977
|
-
hVal[j + 1] = vi;
|
|
2978
|
-
hIdx[j + 1] = ii;
|
|
2979
|
-
}
|
|
2980
|
-
const maxScore = hVal[0];
|
|
2981
|
-
let sumExp = 0;
|
|
2982
|
-
for (let i = 0; i < heapSize; i++) {
|
|
2983
|
-
const p = Math.exp(hVal[i] - maxScore);
|
|
2984
|
-
hVal[i] = p;
|
|
2985
|
-
sumExp += p;
|
|
2986
|
-
}
|
|
2987
|
-
const invSum = 1 / sumExp;
|
|
2988
|
-
for (let i = 0; i < heapSize; i++) hVal[i] *= invSum;
|
|
2989
|
-
let candidateCount = heapSize;
|
|
2990
|
-
if (topP < 1) {
|
|
2991
|
-
let cumulative$1 = 0;
|
|
2992
|
-
for (let i = 0; i < heapSize; i++) {
|
|
2993
|
-
cumulative$1 += hVal[i];
|
|
2994
|
-
if (cumulative$1 >= topP) {
|
|
2995
|
-
candidateCount = i + 1;
|
|
2996
|
-
break;
|
|
2997
|
-
}
|
|
2998
|
-
}
|
|
2999
|
-
let sum = 0;
|
|
3000
|
-
for (let i = 0; i < candidateCount; i++) sum += hVal[i];
|
|
3001
|
-
const inv = 1 / sum;
|
|
3002
|
-
for (let i = 0; i < candidateCount; i++) hVal[i] *= inv;
|
|
3003
|
-
}
|
|
3004
|
-
const r = Math.random();
|
|
3005
|
-
let cumulative = 0;
|
|
3006
|
-
for (let i = 0; i < candidateCount; i++) {
|
|
3007
|
-
cumulative += hVal[i];
|
|
3008
|
-
if (r <= cumulative) return hIdx[i];
|
|
3009
|
-
}
|
|
3010
|
-
return hIdx[candidateCount - 1];
|
|
3173
|
+
function relativePositionBucket(relativePosition, numBuckets, maxDistance) {
|
|
3174
|
+
let bucket = 0;
|
|
3175
|
+
let n = relativePosition;
|
|
3176
|
+
const halfBuckets = Math.floor(numBuckets / 2);
|
|
3177
|
+
if (n > 0) bucket += halfBuckets;
|
|
3178
|
+
else n = -n;
|
|
3179
|
+
const maxExact = Math.floor(halfBuckets / 2);
|
|
3180
|
+
const isSmall = n < maxExact;
|
|
3181
|
+
let largeVal = maxExact + Math.floor(Math.log(n / maxExact) / Math.log(maxDistance / maxExact) * (halfBuckets - maxExact));
|
|
3182
|
+
largeVal = Math.min(largeVal, halfBuckets - 1);
|
|
3183
|
+
bucket += isSmall ? n : largeVal;
|
|
3184
|
+
return bucket;
|
|
3011
3185
|
}
|
|
3012
3186
|
/**
|
|
3013
|
-
*
|
|
3187
|
+
* Build the dense per-head relative-position bias B[head, q, k] = bias_table[bucket(k-q), head]
|
|
3188
|
+
* for a length-S encoder sequence. `biasTable` is the learned
|
|
3189
|
+
* relative_attention_bias.weight, row-major [num_buckets, num_heads].
|
|
3190
|
+
* Returns a flat Float32Array of length num_heads*S*S (head-major, then q, then k).
|
|
3014
3191
|
*/
|
|
3015
|
-
function
|
|
3016
|
-
|
|
3017
|
-
|
|
3018
|
-
for (let
|
|
3019
|
-
|
|
3020
|
-
|
|
3021
|
-
|
|
3022
|
-
|
|
3023
|
-
|
|
3024
|
-
/** Min-heap sift down on parallel index/value typed arrays. */
|
|
3025
|
-
function siftDown(indices, values, i, n) {
|
|
3026
|
-
while (true) {
|
|
3027
|
-
let smallest = i;
|
|
3028
|
-
const left = 2 * i + 1;
|
|
3029
|
-
const right = 2 * i + 2;
|
|
3030
|
-
if (left < n && values[left] < values[smallest]) smallest = left;
|
|
3031
|
-
if (right < n && values[right] < values[smallest]) smallest = right;
|
|
3032
|
-
if (smallest === i) break;
|
|
3033
|
-
const ti = indices[i];
|
|
3034
|
-
indices[i] = indices[smallest];
|
|
3035
|
-
indices[smallest] = ti;
|
|
3036
|
-
const tv = values[i];
|
|
3037
|
-
values[i] = values[smallest];
|
|
3038
|
-
values[smallest] = tv;
|
|
3039
|
-
i = smallest;
|
|
3040
|
-
}
|
|
3041
|
-
}
|
|
3042
|
-
|
|
3043
|
-
//#endregion
|
|
3044
|
-
//#region src/gpu/vision-executor.ts
|
|
3045
|
-
const MAP_MODE_READ = 1;
|
|
3046
|
-
/** Pack an f32 array to IEEE half (f16) bits (round-to-nearest-even). */
|
|
3047
|
-
const _f16scratch = /* @__PURE__ */ new ArrayBuffer(4);
|
|
3048
|
-
const _f16f32 = new Float32Array(_f16scratch);
|
|
3049
|
-
const _f16u32 = new Uint32Array(_f16scratch);
|
|
3050
|
-
function packF16(src) {
|
|
3051
|
-
const out = new Uint16Array(src.length);
|
|
3052
|
-
for (let i = 0; i < src.length; i++) {
|
|
3053
|
-
_f16f32[0] = src[i];
|
|
3054
|
-
const f = _f16u32[0];
|
|
3055
|
-
const sign = f >>> 16 & 32768;
|
|
3056
|
-
const exp = (f >>> 23 & 255) - 112;
|
|
3057
|
-
const mant = f & 8388607;
|
|
3058
|
-
if (exp <= 0) out[i] = sign;
|
|
3059
|
-
else if (exp >= 31) out[i] = sign | 31744;
|
|
3060
|
-
else {
|
|
3061
|
-
const round = (mant & 4096) !== 0;
|
|
3062
|
-
let half = sign | exp << 10 | mant >>> 13;
|
|
3063
|
-
if (round) half += 1;
|
|
3064
|
-
out[i] = half;
|
|
3192
|
+
function buildT5RelativeBias(biasTable, numBuckets, maxDistance, numHeads, S) {
|
|
3193
|
+
const out = new Float32Array(numHeads * S * S);
|
|
3194
|
+
const buckets = new Int32Array(S * S);
|
|
3195
|
+
for (let q = 0; q < S; q++) for (let k = 0; k < S; k++) buckets[q * S + k] = relativePositionBucket(k - q, numBuckets, maxDistance);
|
|
3196
|
+
for (let h = 0; h < numHeads; h++) {
|
|
3197
|
+
const headBase = h * S * S;
|
|
3198
|
+
for (let q = 0; q < S; q++) for (let k = 0; k < S; k++) {
|
|
3199
|
+
const bucket = buckets[q * S + k];
|
|
3200
|
+
out[headBase + q * S + k] = biasTable[bucket * numHeads + h];
|
|
3065
3201
|
}
|
|
3066
3202
|
}
|
|
3067
3203
|
return out;
|
|
3068
3204
|
}
|
|
3069
3205
|
/**
|
|
3070
|
-
*
|
|
3071
|
-
*
|
|
3072
|
-
*
|
|
3073
|
-
*
|
|
3074
|
-
* the attention grid into windows of this many query rows, each its own
|
|
3075
|
-
* submit + onSubmittedWorkDone() drain, so no single submission runs long enough
|
|
3076
|
-
* to trip the watchdog. Math is unchanged: every query row computes identically
|
|
3077
|
-
* regardless of how the grid is partitioned.
|
|
3206
|
+
* Build the T5 encoder graph for a concrete description length S. The relative
|
|
3207
|
+
* bias `t5_rel_bias` enters as a host-written activation [num_heads, S, S]. Like
|
|
3208
|
+
* Moonshine, the graph also pre-projects encoder_out through every decoder layer's
|
|
3209
|
+
* encoder_attn.k/v_proj → enc_k_layer{i}/enc_v_layer{i} (frozen cross-attn K/V).
|
|
3078
3210
|
*/
|
|
3079
|
-
|
|
3080
|
-
|
|
3081
|
-
|
|
3082
|
-
|
|
3083
|
-
|
|
3084
|
-
|
|
3085
|
-
|
|
3086
|
-
|
|
3087
|
-
|
|
3088
|
-
|
|
3089
|
-
|
|
3090
|
-
|
|
3091
|
-
|
|
3092
|
-
|
|
3093
|
-
|
|
3094
|
-
|
|
3095
|
-
|
|
3096
|
-
|
|
3097
|
-
|
|
3098
|
-
|
|
3099
|
-
|
|
3100
|
-
|
|
3101
|
-
|
|
3102
|
-
|
|
3103
|
-
|
|
3104
|
-
|
|
3105
|
-
|
|
3106
|
-
|
|
3107
|
-
|
|
3108
|
-
|
|
3109
|
-
|
|
3110
|
-
|
|
3111
|
-
|
|
3112
|
-
|
|
3113
|
-
|
|
3114
|
-
|
|
3115
|
-
|
|
3116
|
-
|
|
3117
|
-
|
|
3118
|
-
|
|
3119
|
-
|
|
3211
|
+
function generateParlerEncoderGraph(raw, sDesc) {
|
|
3212
|
+
const c = parseParlerConfig(raw);
|
|
3213
|
+
const H = c.enc_d_model;
|
|
3214
|
+
const innerDim = c.enc_heads * c.enc_d_kv;
|
|
3215
|
+
const tensors = {};
|
|
3216
|
+
const nodes = [];
|
|
3217
|
+
const executionOrder = [];
|
|
3218
|
+
const addTensor = (t) => {
|
|
3219
|
+
tensors[t.name] = t;
|
|
3220
|
+
};
|
|
3221
|
+
const addNode = (n) => {
|
|
3222
|
+
nodes.push(n);
|
|
3223
|
+
executionOrder.push(n.id);
|
|
3224
|
+
};
|
|
3225
|
+
const w = (name, shape) => addTensor({
|
|
3226
|
+
name,
|
|
3227
|
+
shape,
|
|
3228
|
+
dtype: "f32",
|
|
3229
|
+
storage: "constant",
|
|
3230
|
+
safetensorsKey: name
|
|
3231
|
+
});
|
|
3232
|
+
const act = (name, shape) => addTensor({
|
|
3233
|
+
name,
|
|
3234
|
+
shape,
|
|
3235
|
+
dtype: "f32",
|
|
3236
|
+
storage: "activation"
|
|
3237
|
+
});
|
|
3238
|
+
const linear = (id, inp, weight, out, K, N) => {
|
|
3239
|
+
w(weight, [N, K]);
|
|
3240
|
+
addNode({
|
|
3241
|
+
id,
|
|
3242
|
+
opType: "MatMul",
|
|
3243
|
+
inputs: [inp, weight],
|
|
3244
|
+
outputs: [out],
|
|
3245
|
+
attributes: {
|
|
3246
|
+
M_tensor: inp,
|
|
3247
|
+
K,
|
|
3248
|
+
N
|
|
3249
|
+
}
|
|
3250
|
+
});
|
|
3251
|
+
};
|
|
3252
|
+
const t5norm = (id, inp, weight, out) => {
|
|
3253
|
+
w(weight, [H]);
|
|
3254
|
+
act(out, [sDesc, H]);
|
|
3255
|
+
addNode({
|
|
3256
|
+
id,
|
|
3257
|
+
opType: "RMSNorm",
|
|
3258
|
+
inputs: [inp, weight],
|
|
3259
|
+
outputs: [out],
|
|
3260
|
+
attributes: {
|
|
3261
|
+
hidden_size: H,
|
|
3262
|
+
eps: c.ln_eps,
|
|
3263
|
+
seq_len_tensor: inp
|
|
3264
|
+
}
|
|
3265
|
+
});
|
|
3266
|
+
};
|
|
3267
|
+
addTensor({
|
|
3268
|
+
name: "input_ids",
|
|
3269
|
+
shape: ["T"],
|
|
3270
|
+
dtype: "u32",
|
|
3271
|
+
storage: "activation"
|
|
3272
|
+
});
|
|
3273
|
+
addTensor({
|
|
3274
|
+
name: "text_encoder.shared.weight",
|
|
3275
|
+
shape: [c.enc_vocab, H],
|
|
3276
|
+
dtype: "f32",
|
|
3277
|
+
storage: "constant",
|
|
3278
|
+
safetensorsKey: "text_encoder.shared.weight"
|
|
3279
|
+
});
|
|
3280
|
+
act("enc_embed", [sDesc, H]);
|
|
3281
|
+
addNode({
|
|
3282
|
+
id: "enc_embed",
|
|
3283
|
+
opType: "Embedding",
|
|
3284
|
+
inputs: ["input_ids", "text_encoder.shared.weight"],
|
|
3285
|
+
outputs: ["enc_embed"],
|
|
3286
|
+
attributes: {
|
|
3287
|
+
vocab_size: c.enc_vocab,
|
|
3288
|
+
hidden_size: H
|
|
3289
|
+
}
|
|
3290
|
+
});
|
|
3291
|
+
addTensor({
|
|
3292
|
+
name: "t5_rel_bias",
|
|
3293
|
+
shape: [
|
|
3294
|
+
c.enc_heads,
|
|
3295
|
+
sDesc,
|
|
3296
|
+
sDesc
|
|
3297
|
+
],
|
|
3298
|
+
dtype: "f32",
|
|
3299
|
+
storage: "activation"
|
|
3300
|
+
});
|
|
3301
|
+
let prev = "enc_embed";
|
|
3302
|
+
for (let i = 0; i < c.enc_layers; i++) {
|
|
3303
|
+
const p = `enc${i}`;
|
|
3304
|
+
const k = (s) => `text_encoder.encoder.block.${i}.${s}`;
|
|
3305
|
+
t5norm(`${p}_norm0`, prev, k("layer.0.layer_norm.weight"), `${p}_n0`);
|
|
3306
|
+
act(`${p}_q`, [sDesc, innerDim]);
|
|
3307
|
+
act(`${p}_k`, [sDesc, innerDim]);
|
|
3308
|
+
act(`${p}_v`, [sDesc, innerDim]);
|
|
3309
|
+
linear(`${p}_qp`, `${p}_n0`, k("layer.0.SelfAttention.q.weight"), `${p}_q`, H, innerDim);
|
|
3310
|
+
linear(`${p}_kp`, `${p}_n0`, k("layer.0.SelfAttention.k.weight"), `${p}_k`, H, innerDim);
|
|
3311
|
+
linear(`${p}_vp`, `${p}_n0`, k("layer.0.SelfAttention.v.weight"), `${p}_v`, H, innerDim);
|
|
3312
|
+
act(`${p}_attn`, [sDesc, innerDim]);
|
|
3313
|
+
addNode({
|
|
3314
|
+
id: `${p}_sa`,
|
|
3315
|
+
opType: "T5Attention",
|
|
3316
|
+
inputs: [
|
|
3317
|
+
`${p}_q`,
|
|
3318
|
+
`${p}_k`,
|
|
3319
|
+
`${p}_v`,
|
|
3320
|
+
"t5_rel_bias"
|
|
3321
|
+
],
|
|
3322
|
+
outputs: [`${p}_attn`],
|
|
3323
|
+
attributes: {
|
|
3324
|
+
num_heads: c.enc_heads,
|
|
3325
|
+
head_dim: c.enc_d_kv,
|
|
3326
|
+
S: sDesc
|
|
3327
|
+
}
|
|
3328
|
+
});
|
|
3329
|
+
act(`${p}_o`, [sDesc, H]);
|
|
3330
|
+
linear(`${p}_op`, `${p}_attn`, k("layer.0.SelfAttention.o.weight"), `${p}_o`, innerDim, H);
|
|
3331
|
+
act(`${p}_res0`, [sDesc, H]);
|
|
3332
|
+
addNode({
|
|
3333
|
+
id: `${p}_add0`,
|
|
3334
|
+
opType: "Add",
|
|
3335
|
+
inputs: [prev, `${p}_o`],
|
|
3336
|
+
outputs: [`${p}_res0`],
|
|
3337
|
+
attributes: {
|
|
3338
|
+
count_tensor: prev,
|
|
3339
|
+
hidden_size: H
|
|
3340
|
+
}
|
|
3341
|
+
});
|
|
3342
|
+
t5norm(`${p}_norm1`, `${p}_res0`, k("layer.1.layer_norm.weight"), `${p}_n1`);
|
|
3343
|
+
act(`${p}_wi0`, [sDesc, c.enc_d_ff]);
|
|
3344
|
+
act(`${p}_wi1`, [sDesc, c.enc_d_ff]);
|
|
3345
|
+
linear(`${p}_wi0p`, `${p}_n1`, k("layer.1.DenseReluDense.wi_0.weight"), `${p}_wi0`, H, c.enc_d_ff);
|
|
3346
|
+
linear(`${p}_wi1p`, `${p}_n1`, k("layer.1.DenseReluDense.wi_1.weight"), `${p}_wi1`, H, c.enc_d_ff);
|
|
3347
|
+
act(`${p}_gelu`, [sDesc, c.enc_d_ff]);
|
|
3348
|
+
addNode({
|
|
3349
|
+
id: `${p}_act`,
|
|
3350
|
+
opType: "GELU",
|
|
3351
|
+
inputs: [`${p}_wi0`],
|
|
3352
|
+
outputs: [`${p}_gelu`],
|
|
3353
|
+
attributes: { count_tensor: `${p}_wi0` }
|
|
3354
|
+
});
|
|
3355
|
+
act(`${p}_gated`, [sDesc, c.enc_d_ff]);
|
|
3356
|
+
addNode({
|
|
3357
|
+
id: `${p}_gate`,
|
|
3358
|
+
opType: "Mul",
|
|
3359
|
+
inputs: [`${p}_gelu`, `${p}_wi1`],
|
|
3360
|
+
outputs: [`${p}_gated`],
|
|
3361
|
+
attributes: {
|
|
3362
|
+
count_tensor: `${p}_gelu`,
|
|
3363
|
+
hidden_size: c.enc_d_ff
|
|
3364
|
+
}
|
|
3365
|
+
});
|
|
3366
|
+
act(`${p}_wo`, [sDesc, H]);
|
|
3367
|
+
linear(`${p}_wop`, `${p}_gated`, k("layer.1.DenseReluDense.wo.weight"), `${p}_wo`, c.enc_d_ff, H);
|
|
3368
|
+
act(`${p}_res1`, [sDesc, H]);
|
|
3369
|
+
addNode({
|
|
3370
|
+
id: `${p}_add1`,
|
|
3371
|
+
opType: "Add",
|
|
3372
|
+
inputs: [`${p}_res0`, `${p}_wo`],
|
|
3373
|
+
outputs: [`${p}_res1`],
|
|
3374
|
+
attributes: {
|
|
3375
|
+
count_tensor: `${p}_res0`,
|
|
3376
|
+
hidden_size: H
|
|
3377
|
+
}
|
|
3378
|
+
});
|
|
3379
|
+
prev = `${p}_res1`;
|
|
3380
|
+
}
|
|
3381
|
+
w("text_encoder.encoder.final_layer_norm.weight", [H]);
|
|
3382
|
+
act("encoder_out", [sDesc, H]);
|
|
3383
|
+
addNode({
|
|
3384
|
+
id: "enc_final_norm",
|
|
3385
|
+
opType: "RMSNorm",
|
|
3386
|
+
inputs: [prev, "text_encoder.encoder.final_layer_norm.weight"],
|
|
3387
|
+
outputs: ["encoder_out"],
|
|
3388
|
+
attributes: {
|
|
3389
|
+
hidden_size: H,
|
|
3390
|
+
eps: c.ln_eps,
|
|
3391
|
+
seq_len_tensor: prev
|
|
3392
|
+
}
|
|
3393
|
+
});
|
|
3394
|
+
const outputs = ["encoder_out"];
|
|
3395
|
+
for (let i = 0; i < c.dec_layers; i++) {
|
|
3396
|
+
const kw = `decoder.model.decoder.layers.${i}.encoder_attn.k_proj.weight`;
|
|
3397
|
+
const vw = `decoder.model.decoder.layers.${i}.encoder_attn.v_proj.weight`;
|
|
3398
|
+
w(kw, [H, H]);
|
|
3399
|
+
w(vw, [H, H]);
|
|
3400
|
+
const encK = `enc_k_layer${i}`;
|
|
3401
|
+
const encV = `enc_v_layer${i}`;
|
|
3402
|
+
act(encK, [sDesc, H]);
|
|
3403
|
+
act(encV, [sDesc, H]);
|
|
3404
|
+
addNode({
|
|
3405
|
+
id: `enc_kproj${i}`,
|
|
3406
|
+
opType: "MatMul",
|
|
3407
|
+
inputs: ["encoder_out", kw],
|
|
3408
|
+
outputs: [encK],
|
|
3409
|
+
attributes: {
|
|
3410
|
+
M_tensor: "encoder_out",
|
|
3411
|
+
K: H,
|
|
3412
|
+
N: H
|
|
3413
|
+
}
|
|
3414
|
+
});
|
|
3415
|
+
addNode({
|
|
3416
|
+
id: `enc_vproj${i}`,
|
|
3417
|
+
opType: "MatMul",
|
|
3418
|
+
inputs: ["encoder_out", vw],
|
|
3419
|
+
outputs: [encV],
|
|
3420
|
+
attributes: {
|
|
3421
|
+
M_tensor: "encoder_out",
|
|
3422
|
+
K: H,
|
|
3423
|
+
N: H
|
|
3424
|
+
}
|
|
3425
|
+
});
|
|
3426
|
+
outputs.push(encK, encV);
|
|
3427
|
+
}
|
|
3428
|
+
return {
|
|
3429
|
+
architecture: "ParlerTextEncoder",
|
|
3430
|
+
config: baseConfig(c, H, c.enc_heads, c.enc_layers),
|
|
3431
|
+
capabilities: CAPS,
|
|
3432
|
+
tensors,
|
|
3433
|
+
nodes,
|
|
3434
|
+
executionOrder,
|
|
3435
|
+
inputs: ["input_ids", "t5_rel_bias"],
|
|
3436
|
+
outputs
|
|
3437
|
+
};
|
|
3438
|
+
}
|
|
3439
|
+
/**
|
|
3440
|
+
* Build the Parler decoder graph for a single decode step (T=1). The host supplies:
|
|
3441
|
+
* - `dec_input_embed` [1, H]: the SUMMED 9-codebook embedding of the current step's
|
|
3442
|
+
* codes PLUS the sinusoidal position (computed host-side from embed_tokens +
|
|
3443
|
+
* embed_positions, since the 9-way embedding sum is cheap on the host and avoids
|
|
3444
|
+
* 9 Embedding ops + an Add tree per step).
|
|
3445
|
+
* - `enc_k_layer{i}` / `enc_v_layer{i}`: frozen cross-attn K/V from the encoder.
|
|
3446
|
+
* Output: 9 logit rows `logits_cb{c}` [1, audio_vocab] (one per codebook head).
|
|
3447
|
+
*
|
|
3448
|
+
* The prompt prefix (transcript) is handled by prefilling its embeddings through the
|
|
3449
|
+
* SAME graph before audio decode begins (the driver feeds the prompt rows first, with
|
|
3450
|
+
* their sinusoidal positions, so the self-attn KV-cache holds [prompt ++ audio]).
|
|
3451
|
+
*/
|
|
3452
|
+
function generateParlerDecoderGraph(raw, sEnc) {
|
|
3453
|
+
const c = parseParlerConfig(raw);
|
|
3454
|
+
const H = c.dec_hidden;
|
|
3455
|
+
const heads = c.dec_heads;
|
|
3456
|
+
const headDim = Math.floor(H / heads);
|
|
3457
|
+
const tensors = {};
|
|
3458
|
+
const nodes = [];
|
|
3459
|
+
const executionOrder = [];
|
|
3460
|
+
const addTensor = (t) => {
|
|
3461
|
+
tensors[t.name] = t;
|
|
3462
|
+
};
|
|
3463
|
+
const addNode = (n) => {
|
|
3464
|
+
nodes.push(n);
|
|
3465
|
+
executionOrder.push(n.id);
|
|
3466
|
+
};
|
|
3467
|
+
const w = (name, shape) => addTensor({
|
|
3468
|
+
name,
|
|
3469
|
+
shape,
|
|
3470
|
+
dtype: "f32",
|
|
3471
|
+
storage: "constant",
|
|
3472
|
+
safetensorsKey: name
|
|
3473
|
+
});
|
|
3474
|
+
const act = (name, shape) => addTensor({
|
|
3475
|
+
name,
|
|
3476
|
+
shape,
|
|
3477
|
+
dtype: "f32",
|
|
3478
|
+
storage: "activation"
|
|
3479
|
+
});
|
|
3480
|
+
const linear = (id, inp, weight, out, K, N) => {
|
|
3481
|
+
w(weight, [N, K]);
|
|
3482
|
+
addNode({
|
|
3483
|
+
id,
|
|
3484
|
+
opType: "MatMul",
|
|
3485
|
+
inputs: [inp, weight],
|
|
3486
|
+
outputs: [out],
|
|
3487
|
+
attributes: {
|
|
3488
|
+
M_tensor: inp,
|
|
3489
|
+
K,
|
|
3490
|
+
N
|
|
3491
|
+
}
|
|
3492
|
+
});
|
|
3493
|
+
};
|
|
3494
|
+
const layernorm = (id, inp, prefix, out) => {
|
|
3495
|
+
w(`${prefix}.weight`, [H]);
|
|
3496
|
+
w(`${prefix}.bias`, [H]);
|
|
3497
|
+
act(out, ["T", H]);
|
|
3498
|
+
addNode({
|
|
3499
|
+
id,
|
|
3500
|
+
opType: "LayerNorm",
|
|
3501
|
+
inputs: [
|
|
3502
|
+
inp,
|
|
3503
|
+
`${prefix}.weight`,
|
|
3504
|
+
`${prefix}.bias`
|
|
3505
|
+
],
|
|
3506
|
+
outputs: [out],
|
|
3507
|
+
attributes: {
|
|
3508
|
+
hidden_size: H,
|
|
3509
|
+
eps: c.ln_eps,
|
|
3510
|
+
seq_len_tensor: inp,
|
|
3511
|
+
has_bias: true
|
|
3512
|
+
}
|
|
3513
|
+
});
|
|
3514
|
+
};
|
|
3515
|
+
const inputs = ["dec_input_embed"];
|
|
3516
|
+
addTensor({
|
|
3517
|
+
name: "dec_input_embed",
|
|
3518
|
+
shape: ["T", H],
|
|
3519
|
+
dtype: "f32",
|
|
3520
|
+
storage: "activation"
|
|
3521
|
+
});
|
|
3522
|
+
let prev = "dec_input_embed";
|
|
3523
|
+
for (let i = 0; i < c.dec_layers; i++) {
|
|
3524
|
+
const p = `dec${i}`;
|
|
3525
|
+
const lk = `decoder.model.decoder.layers.${i}`;
|
|
3526
|
+
layernorm(`${p}_n1`, prev, `${lk}.self_attn_layer_norm`, `${p}_ln1`);
|
|
3527
|
+
act(`${p}_q`, ["T", H]);
|
|
3528
|
+
act(`${p}_k`, ["T", H]);
|
|
3529
|
+
act(`${p}_v`, ["T", H]);
|
|
3530
|
+
linear(`${p}_qp`, `${p}_ln1`, `${lk}.self_attn.q_proj.weight`, `${p}_q`, H, H);
|
|
3531
|
+
linear(`${p}_kp`, `${p}_ln1`, `${lk}.self_attn.k_proj.weight`, `${p}_k`, H, H);
|
|
3532
|
+
linear(`${p}_vp`, `${p}_ln1`, `${lk}.self_attn.v_proj.weight`, `${p}_v`, H, H);
|
|
3533
|
+
addTensor({
|
|
3534
|
+
name: `${p}_kcache`,
|
|
3535
|
+
shape: ["L_max", H],
|
|
3536
|
+
dtype: "f32",
|
|
3537
|
+
storage: "kv_cache"
|
|
3538
|
+
});
|
|
3539
|
+
addTensor({
|
|
3540
|
+
name: `${p}_vcache`,
|
|
3541
|
+
shape: ["L_max", H],
|
|
3542
|
+
dtype: "f32",
|
|
3543
|
+
storage: "kv_cache"
|
|
3544
|
+
});
|
|
3545
|
+
addNode({
|
|
3546
|
+
id: `${p}_kappend`,
|
|
3547
|
+
opType: "KVCacheAppend",
|
|
3548
|
+
inputs: [`${p}_k`],
|
|
3549
|
+
outputs: [`${p}_kcache`],
|
|
3550
|
+
attributes: {
|
|
3551
|
+
width: H,
|
|
3552
|
+
T_tensor: `${p}_k`
|
|
3553
|
+
}
|
|
3554
|
+
});
|
|
3555
|
+
addNode({
|
|
3556
|
+
id: `${p}_vappend`,
|
|
3557
|
+
opType: "KVCacheAppend",
|
|
3558
|
+
inputs: [`${p}_v`],
|
|
3559
|
+
outputs: [`${p}_vcache`],
|
|
3560
|
+
attributes: {
|
|
3561
|
+
width: H,
|
|
3562
|
+
T_tensor: `${p}_v`
|
|
3563
|
+
}
|
|
3564
|
+
});
|
|
3565
|
+
act(`${p}_sa`, ["T", H]);
|
|
3566
|
+
addNode({
|
|
3567
|
+
id: `${p}_self_attn`,
|
|
3568
|
+
opType: "Attention",
|
|
3569
|
+
inputs: [
|
|
3570
|
+
`${p}_q`,
|
|
3571
|
+
`${p}_kcache`,
|
|
3572
|
+
`${p}_vcache`
|
|
3573
|
+
],
|
|
3574
|
+
outputs: [`${p}_sa`],
|
|
3575
|
+
attributes: {
|
|
3576
|
+
hidden_size: H,
|
|
3577
|
+
num_q_heads: heads,
|
|
3578
|
+
num_kv_heads: c.dec_kv_heads,
|
|
3579
|
+
head_dim: headDim,
|
|
3580
|
+
causal: true,
|
|
3581
|
+
layer_index: i
|
|
3582
|
+
}
|
|
3583
|
+
});
|
|
3584
|
+
act(`${p}_sao`, ["T", H]);
|
|
3585
|
+
linear(`${p}_sap`, `${p}_sa`, `${lk}.self_attn.out_proj.weight`, `${p}_sao`, H, H);
|
|
3586
|
+
act(`${p}_res1`, ["T", H]);
|
|
3587
|
+
addNode({
|
|
3588
|
+
id: `${p}_add1`,
|
|
3589
|
+
opType: "Add",
|
|
3590
|
+
inputs: [prev, `${p}_sao`],
|
|
3591
|
+
outputs: [`${p}_res1`],
|
|
3592
|
+
attributes: {
|
|
3593
|
+
count_tensor: prev,
|
|
3594
|
+
hidden_size: H
|
|
3595
|
+
}
|
|
3596
|
+
});
|
|
3597
|
+
layernorm(`${p}_n2`, `${p}_res1`, `${lk}.encoder_attn_layer_norm`, `${p}_ln2`);
|
|
3598
|
+
act(`${p}_cq`, ["T", H]);
|
|
3599
|
+
linear(`${p}_cqp`, `${p}_ln2`, `${lk}.encoder_attn.q_proj.weight`, `${p}_cq`, H, H);
|
|
3600
|
+
const encK = `enc_k_layer${i}`;
|
|
3601
|
+
const encV = `enc_v_layer${i}`;
|
|
3602
|
+
act(encK, [sEnc, H]);
|
|
3603
|
+
act(encV, [sEnc, H]);
|
|
3604
|
+
inputs.push(encK, encV);
|
|
3605
|
+
act(`${p}_ca`, ["T", H]);
|
|
3606
|
+
addNode({
|
|
3607
|
+
id: `${p}_cross_attn`,
|
|
3608
|
+
opType: "CrossAttention",
|
|
3609
|
+
inputs: [
|
|
3610
|
+
`${p}_cq`,
|
|
3611
|
+
encK,
|
|
3612
|
+
encV
|
|
3613
|
+
],
|
|
3614
|
+
outputs: [`${p}_ca`],
|
|
3615
|
+
attributes: {
|
|
3616
|
+
num_q_heads: heads,
|
|
3617
|
+
num_kv_heads: c.dec_kv_heads,
|
|
3618
|
+
head_dim: headDim
|
|
3619
|
+
}
|
|
3620
|
+
});
|
|
3621
|
+
act(`${p}_cao`, ["T", H]);
|
|
3622
|
+
linear(`${p}_cap`, `${p}_ca`, `${lk}.encoder_attn.out_proj.weight`, `${p}_cao`, H, H);
|
|
3623
|
+
act(`${p}_res2`, ["T", H]);
|
|
3624
|
+
addNode({
|
|
3625
|
+
id: `${p}_add2`,
|
|
3626
|
+
opType: "Add",
|
|
3627
|
+
inputs: [`${p}_res1`, `${p}_cao`],
|
|
3628
|
+
outputs: [`${p}_res2`],
|
|
3629
|
+
attributes: {
|
|
3630
|
+
count_tensor: `${p}_res1`,
|
|
3631
|
+
hidden_size: H
|
|
3632
|
+
}
|
|
3633
|
+
});
|
|
3634
|
+
layernorm(`${p}_n3`, `${p}_res2`, `${lk}.final_layer_norm`, `${p}_ln3`);
|
|
3635
|
+
act(`${p}_fc1`, ["T", c.dec_ffn]);
|
|
3636
|
+
linear(`${p}_fc1p`, `${p}_ln3`, `${lk}.fc1.weight`, `${p}_fc1`, H, c.dec_ffn);
|
|
3637
|
+
act(`${p}_gelu`, ["T", c.dec_ffn]);
|
|
3638
|
+
addNode({
|
|
3639
|
+
id: `${p}_act`,
|
|
3640
|
+
opType: "GeluErf",
|
|
3641
|
+
inputs: [`${p}_fc1`],
|
|
3642
|
+
outputs: [`${p}_gelu`],
|
|
3643
|
+
attributes: { count_tensor: `${p}_fc1` }
|
|
3644
|
+
});
|
|
3645
|
+
act(`${p}_fc2`, ["T", H]);
|
|
3646
|
+
linear(`${p}_fc2p`, `${p}_gelu`, `${lk}.fc2.weight`, `${p}_fc2`, c.dec_ffn, H);
|
|
3647
|
+
act(`${p}_res3`, ["T", H]);
|
|
3648
|
+
addNode({
|
|
3649
|
+
id: `${p}_add3`,
|
|
3650
|
+
opType: "Add",
|
|
3651
|
+
inputs: [`${p}_res2`, `${p}_fc2`],
|
|
3652
|
+
outputs: [`${p}_res3`],
|
|
3653
|
+
attributes: {
|
|
3654
|
+
count_tensor: `${p}_res2`,
|
|
3655
|
+
hidden_size: H
|
|
3656
|
+
}
|
|
3657
|
+
});
|
|
3658
|
+
prev = `${p}_res3`;
|
|
3659
|
+
}
|
|
3660
|
+
layernorm("dec_final_norm", prev, "decoder.model.decoder.layer_norm", "dec_normed");
|
|
3661
|
+
act("dec_last", [1, H]);
|
|
3662
|
+
addNode({
|
|
3663
|
+
id: "slice_last",
|
|
3664
|
+
opType: "SliceLastRow",
|
|
3665
|
+
inputs: ["dec_normed"],
|
|
3666
|
+
outputs: ["dec_last"],
|
|
3667
|
+
attributes: { width: H }
|
|
3668
|
+
});
|
|
3669
|
+
const outputs = [];
|
|
3670
|
+
for (let cb = 0; cb < c.num_codebooks; cb++) {
|
|
3671
|
+
const headW = `decoder.lm_heads.${cb}.weight`;
|
|
3672
|
+
w(headW, [c.audio_vocab, H]);
|
|
3673
|
+
const logit = `logits_cb${cb}`;
|
|
3674
|
+
act(logit, [1, c.audio_vocab]);
|
|
3675
|
+
addNode({
|
|
3676
|
+
id: `lm_head${cb}`,
|
|
3677
|
+
opType: "MatMul",
|
|
3678
|
+
inputs: ["dec_last", headW],
|
|
3679
|
+
outputs: [logit],
|
|
3680
|
+
attributes: {
|
|
3681
|
+
M_tensor: "dec_last",
|
|
3682
|
+
K: H,
|
|
3683
|
+
N: c.audio_vocab
|
|
3684
|
+
}
|
|
3685
|
+
});
|
|
3686
|
+
outputs.push(logit);
|
|
3687
|
+
}
|
|
3688
|
+
return {
|
|
3689
|
+
architecture: "ParlerTTSDecoder",
|
|
3690
|
+
config: baseConfig(c, H, heads, c.dec_layers),
|
|
3691
|
+
capabilities: CAPS,
|
|
3692
|
+
tensors,
|
|
3693
|
+
nodes,
|
|
3694
|
+
executionOrder,
|
|
3695
|
+
inputs,
|
|
3696
|
+
outputs
|
|
3697
|
+
};
|
|
3698
|
+
}
|
|
3699
|
+
/**
|
|
3700
|
+
* Apply the Parler/MusicGen delay pattern to a [numCodebooks, T] code grid produced
|
|
3701
|
+
* by AR decode. During generation, codebook `i` is delayed by `i` steps: at decode
|
|
3702
|
+
* step t the model emits, for each codebook i, the code for ACOUSTIC frame (t − i).
|
|
3703
|
+
* To recover the aligned acoustic grid we simply read codebook i's stream starting
|
|
3704
|
+
* at offset i. This returns the de-delayed [numCodebooks, frames] grid (codebook-major)
|
|
3705
|
+
* ready for the DAC decoder, where frames = T − (numCodebooks − 1).
|
|
3706
|
+
*
|
|
3707
|
+
* `delayed` is codebook-major [numCodebooks, T] (delayed[i*T + t]).
|
|
3708
|
+
*/
|
|
3709
|
+
function revertDelayPattern(delayed, numCodebooks, T) {
|
|
3710
|
+
const frames = T - (numCodebooks - 1);
|
|
3711
|
+
if (frames <= 0) return {
|
|
3712
|
+
codes: new Uint32Array(0),
|
|
3713
|
+
frames: 0
|
|
3714
|
+
};
|
|
3715
|
+
const codes = new Uint32Array(numCodebooks * frames);
|
|
3716
|
+
for (let i = 0; i < numCodebooks; i++) for (let f = 0; f < frames; f++) codes[i * frames + f] = delayed[i * T + (f + i)];
|
|
3717
|
+
return {
|
|
3718
|
+
codes,
|
|
3719
|
+
frames
|
|
3720
|
+
};
|
|
3721
|
+
}
|
|
3722
|
+
|
|
3723
|
+
//#endregion
|
|
3724
|
+
//#region src/gpu/parler-tts.ts
|
|
3725
|
+
/**
|
|
3726
|
+
* ParlerTTS — native describe-your-voice text-to-speech for Gerbil's WebGPU backend.
|
|
3727
|
+
*
|
|
3728
|
+
* Parler-TTS (parler-tts/parler-tts-mini-v1) is an encoder-decoder codec-LM: you
|
|
3729
|
+
* describe the voice in natural language and a Flan-T5 encoder turns that DESCRIPTION
|
|
3730
|
+
* into cross-attention states, while a 24-layer decoder LM cross-attends to them and
|
|
3731
|
+
* autoregressively emits DAC audio codes for the TRANSCRIPT. Same dual-graph shape as
|
|
3732
|
+
* MoonshineSTT (encoder once → frozen cross-attn K/V → AR decode), and the same DAC
|
|
3733
|
+
* codec family as OuteTTS (the shared generateDacSpeechDecoderGraph, here with 9
|
|
3734
|
+
* codebooks + rates [8,8,4,2]).
|
|
3735
|
+
*
|
|
3736
|
+
* Pipeline (engine.speak(text, { describeVoice })):
|
|
3737
|
+
* 1. Tokenize the DESCRIPTION with the T5 tokenizer; build the relative-position
|
|
3738
|
+
* bias; run the T5 encoder once → encoder_out + per-decoder-layer frozen K/V.
|
|
3739
|
+
* 2. Tokenize the TRANSCRIPT (embed_prompts); prefill its embeddings (+ sinusoidal
|
|
3740
|
+
* positions) through the decoder so the self-attn KV-cache holds the prompt.
|
|
3741
|
+
* 3. AR-decode 9 interleaved codebooks under the delay pattern: each step the host
|
|
3742
|
+
* sums the 9 per-codebook embeddings of the current (delayed) codes + the
|
|
3743
|
+
* sinusoidal position, forwards one step, and samples each of the 9 lm_head logit
|
|
3744
|
+
* rows. Stop on EOS in codebook 0 (or the frame cap).
|
|
3745
|
+
* 4. Revert the delay pattern → [9, frames] code grid → shared DAC decode → 44.1 kHz
|
|
3746
|
+
* PCM (tanh applied host-side).
|
|
3747
|
+
*
|
|
3748
|
+
* STATUS — FIRST IMPLEMENTATION. Typechecks + builds. The end-to-end run is GPU-
|
|
3749
|
+
* contended on the dev box, so coherence is pending a final reference check once the
|
|
3750
|
+
* GPU frees up (the T5 encoder cosine + a decoder forward should be diffed vs
|
|
3751
|
+
* transformers ParlerTTSForConditionalGeneration; see scripts/engine/parler-*-ref.py
|
|
3752
|
+
* stubs). The architecture, weight layout, tokenization, delay pattern, and DAC config
|
|
3753
|
+
* are all verified against the live config.json + safetensors header.
|
|
3754
|
+
*/
|
|
3755
|
+
const MAP_MODE_READ$1 = 1;
|
|
3756
|
+
/** Default voice description when the caller doesn't supply one. */
|
|
3757
|
+
const DEFAULT_VOICE_DESCRIPTION = "A clear, neutral voice speaking at a moderate pace with good audio quality.";
|
|
3758
|
+
/** Hard cap on decoded audio frames (~86 fps → ~30 s at 2600). */
|
|
3759
|
+
const DEFAULT_MAX_FRAMES = 2600;
|
|
3760
|
+
/** Build a fresh map of just the constant weights a graph references. */
|
|
3761
|
+
function selectGraphWeights(graph, weights) {
|
|
3762
|
+
const out = /* @__PURE__ */ new Map();
|
|
3763
|
+
for (const [name, desc] of Object.entries(graph.tensors)) {
|
|
3764
|
+
if (desc.storage !== "constant") continue;
|
|
3765
|
+
const w = weights.get(name) ?? (desc.safetensorsKey ? weights.get(desc.safetensorsKey) : void 0);
|
|
3766
|
+
if (w) out.set(name, w);
|
|
3767
|
+
}
|
|
3768
|
+
return out;
|
|
3769
|
+
}
|
|
3770
|
+
/** Argmax over a logit row. */
|
|
3771
|
+
function argmax$1(logits) {
|
|
3772
|
+
let best = 0;
|
|
3773
|
+
let bestV = logits[0];
|
|
3774
|
+
for (let i = 1; i < logits.length; i++) if (logits[i] > bestV) {
|
|
3775
|
+
bestV = logits[i];
|
|
3776
|
+
best = i;
|
|
3777
|
+
}
|
|
3778
|
+
return best;
|
|
3779
|
+
}
|
|
3780
|
+
/** Temperature + optional top-k sample over a logit row. */
|
|
3781
|
+
function sampleLogits(logits, temperature, topK) {
|
|
3782
|
+
const n = logits.length;
|
|
3783
|
+
const scaled = new Float32Array(n);
|
|
3784
|
+
let maxV = Number.NEGATIVE_INFINITY;
|
|
3785
|
+
const invT = 1 / Math.max(temperature, 1e-5);
|
|
3786
|
+
for (let i = 0; i < n; i++) {
|
|
3787
|
+
scaled[i] = logits[i] * invT;
|
|
3788
|
+
if (scaled[i] > maxV) maxV = scaled[i];
|
|
3789
|
+
}
|
|
3790
|
+
let allowed = null;
|
|
3791
|
+
if (topK > 0 && topK < n) {
|
|
3792
|
+
const order = Array.from({ length: n }, (_, i) => i).sort((a, b) => scaled[b] - scaled[a]);
|
|
3793
|
+
allowed = new Int32Array(order.slice(0, topK));
|
|
3794
|
+
}
|
|
3795
|
+
let sum = 0;
|
|
3796
|
+
const probs = new Float32Array(n);
|
|
3797
|
+
if (allowed) for (const idx of allowed) {
|
|
3798
|
+
const e = Math.exp(scaled[idx] - maxV);
|
|
3799
|
+
probs[idx] = e;
|
|
3800
|
+
sum += e;
|
|
3801
|
+
}
|
|
3802
|
+
else for (let i = 0; i < n; i++) {
|
|
3803
|
+
const e = Math.exp(scaled[i] - maxV);
|
|
3804
|
+
probs[i] = e;
|
|
3805
|
+
sum += e;
|
|
3806
|
+
}
|
|
3807
|
+
let pick$1 = Math.random() * sum;
|
|
3808
|
+
if (allowed) {
|
|
3809
|
+
for (const idx of allowed) {
|
|
3810
|
+
pick$1 -= probs[idx];
|
|
3811
|
+
if (pick$1 <= 0) return idx;
|
|
3812
|
+
}
|
|
3813
|
+
return allowed[0];
|
|
3814
|
+
}
|
|
3815
|
+
for (let i = 0; i < n; i++) {
|
|
3816
|
+
pick$1 -= probs[i];
|
|
3817
|
+
if (pick$1 <= 0) return i;
|
|
3818
|
+
}
|
|
3819
|
+
return n - 1;
|
|
3820
|
+
}
|
|
3821
|
+
var ParlerEncoderExecutor = class {
|
|
3822
|
+
ctx;
|
|
3823
|
+
graph;
|
|
3824
|
+
decLayers;
|
|
3825
|
+
hidden;
|
|
3826
|
+
sDesc;
|
|
3827
|
+
weightBuffers = /* @__PURE__ */ new Map();
|
|
3828
|
+
activationBuffers = /* @__PURE__ */ new Map();
|
|
3829
|
+
dispatches = [];
|
|
3830
|
+
constructor(ctx, graph, decLayers, sDesc) {
|
|
3831
|
+
this.ctx = ctx;
|
|
3832
|
+
this.graph = graph;
|
|
3833
|
+
this.decLayers = decLayers;
|
|
3834
|
+
this.hidden = graph.config.hidden_size;
|
|
3835
|
+
this.sDesc = sDesc;
|
|
3836
|
+
for (const [name, desc] of Object.entries(graph.tensors)) {
|
|
3837
|
+
if (desc.storage !== "activation") continue;
|
|
3838
|
+
const count = desc.shape.reduce((a, dim) => a * (typeof dim === "number" ? dim : sDesc), 1);
|
|
3839
|
+
this.activationBuffers.set(name, createStorageBuffer(this.ctx, `penc_${name}`, count * DTYPE_BYTES[desc.dtype]));
|
|
3840
|
+
}
|
|
3841
|
+
}
|
|
3842
|
+
uploadWeights(weights) {
|
|
3843
|
+
for (const [name, desc] of Object.entries(this.graph.tensors)) {
|
|
3844
|
+
if (desc.storage !== "constant") continue;
|
|
3845
|
+
const w = weights.get(name);
|
|
3846
|
+
if (!w) throw new Error(`ParlerEncoderExecutor: missing weight "${name}"`);
|
|
3847
|
+
this.weightBuffers.set(name, createStorageBuffer(this.ctx, `pencw_${name}`, w.data.byteLength, w.data));
|
|
3848
|
+
}
|
|
3849
|
+
}
|
|
3850
|
+
initBindGroups() {
|
|
3851
|
+
const shapes = this.resolveShapes();
|
|
3852
|
+
for (const nodeId of this.graph.executionOrder) {
|
|
3853
|
+
const node = this.graph.nodes.find((n) => n.id === nodeId);
|
|
3854
|
+
const spec = KERNEL_REGISTRY[node.opType];
|
|
3855
|
+
if (!spec) throw new Error(`ParlerEncoderExecutor: no kernel for op "${node.opType}"`);
|
|
3856
|
+
const pipeline = getOrCreatePipeline(this.ctx, `penck_${nodeId}`, spec.shaderCode, spec.entryPoint);
|
|
3857
|
+
const paramsData = spec.buildParams(node, shapes, { seqPos: 0 });
|
|
3858
|
+
const uniformBuffer = createUniformBuffer(this.ctx, `pencu_${nodeId}`, paramsData);
|
|
3859
|
+
const bufferEntries = this.gatherBuffers(spec, node, uniformBuffer);
|
|
3860
|
+
const bindGroup = createBindGroup(this.ctx, pipeline, bufferEntries, `pencbg_${nodeId}`);
|
|
3861
|
+
this.dispatches.push({
|
|
3862
|
+
node,
|
|
3863
|
+
spec,
|
|
3864
|
+
pipeline,
|
|
3865
|
+
bindGroup,
|
|
3866
|
+
uniformBuffer
|
|
3867
|
+
});
|
|
3868
|
+
}
|
|
3869
|
+
}
|
|
3870
|
+
/** Run the encoder. `inputIds` u32 [S], `relBias` f32 [heads,S,S]. */
|
|
3871
|
+
async encode(inputIds, relBias) {
|
|
3872
|
+
const q = this.ctx.device.queue;
|
|
3873
|
+
const inputBuf = this.activationBuffers.get("input_ids");
|
|
3874
|
+
const biasBuf = this.activationBuffers.get("t5_rel_bias");
|
|
3875
|
+
if (!(inputBuf && biasBuf)) throw new Error("ParlerEncoderExecutor: input_ids / t5_rel_bias buffers missing.");
|
|
3876
|
+
q.writeBuffer(inputBuf, 0, inputIds);
|
|
3877
|
+
q.writeBuffer(biasBuf, 0, relBias);
|
|
3878
|
+
const shapes = this.resolveShapes();
|
|
3879
|
+
const sizes = [];
|
|
3880
|
+
for (const d of this.dispatches) {
|
|
3881
|
+
const params = d.spec.buildParams(d.node, shapes, { seqPos: 0 });
|
|
3882
|
+
q.writeBuffer(d.uniformBuffer, 0, params);
|
|
3883
|
+
sizes.push(d.spec.getDispatchSize(d.node, shapes, { seqPos: 0 }));
|
|
3884
|
+
}
|
|
3885
|
+
if (this.ctx.isWebKitWebGPU) for (let i = 0; i < this.dispatches.length; i++) {
|
|
3886
|
+
const enc = this.ctx.device.createCommandEncoder();
|
|
3887
|
+
const p = enc.beginComputePass();
|
|
3888
|
+
p.setPipeline(this.dispatches[i].pipeline);
|
|
3889
|
+
p.setBindGroup(0, this.dispatches[i].bindGroup);
|
|
3890
|
+
p.dispatchWorkgroups(...sizes[i]);
|
|
3891
|
+
p.end();
|
|
3892
|
+
q.submit([enc.finish()]);
|
|
3893
|
+
await q.onSubmittedWorkDone();
|
|
3894
|
+
}
|
|
3895
|
+
else {
|
|
3896
|
+
const enc = this.ctx.device.createCommandEncoder({ label: "parler_encode" });
|
|
3897
|
+
const pass = enc.beginComputePass({ label: "parler_enc_pass" });
|
|
3898
|
+
for (let i = 0; i < this.dispatches.length; i++) {
|
|
3899
|
+
pass.setPipeline(this.dispatches[i].pipeline);
|
|
3900
|
+
pass.setBindGroup(0, this.dispatches[i].bindGroup);
|
|
3901
|
+
pass.dispatchWorkgroups(...sizes[i]);
|
|
3902
|
+
}
|
|
3903
|
+
pass.end();
|
|
3904
|
+
q.submit([enc.finish()]);
|
|
3905
|
+
}
|
|
3906
|
+
const encK = [];
|
|
3907
|
+
const encV = [];
|
|
3908
|
+
for (let i = 0; i < this.decLayers; i++) {
|
|
3909
|
+
encK.push(await this.readBack(`enc_k_layer${i}`, this.sDesc * this.hidden));
|
|
3910
|
+
encV.push(await this.readBack(`enc_v_layer${i}`, this.sDesc * this.hidden));
|
|
3911
|
+
}
|
|
3912
|
+
return {
|
|
3913
|
+
encK,
|
|
3914
|
+
encV
|
|
3915
|
+
};
|
|
3916
|
+
}
|
|
3917
|
+
destroy() {
|
|
3918
|
+
destroyBuffers([...this.weightBuffers.values()]);
|
|
3919
|
+
destroyBuffers([...this.activationBuffers.values()]);
|
|
3920
|
+
for (const d of this.dispatches) d.uniformBuffer.destroy();
|
|
3921
|
+
this.weightBuffers.clear();
|
|
3922
|
+
this.activationBuffers.clear();
|
|
3923
|
+
this.dispatches = [];
|
|
3924
|
+
}
|
|
3925
|
+
resolveShapes() {
|
|
3926
|
+
const resolved = {};
|
|
3927
|
+
for (const [name, desc] of Object.entries(this.graph.tensors)) resolved[name] = desc.shape.map((dim) => typeof dim === "number" ? dim : this.sDesc);
|
|
3928
|
+
return resolved;
|
|
3929
|
+
}
|
|
3930
|
+
gatherBuffers(spec, node, uniformBuffer) {
|
|
3931
|
+
const entries = [];
|
|
3932
|
+
const tensorNames = [...node.inputs, ...node.outputs];
|
|
3933
|
+
let tensorIdx = 0;
|
|
3934
|
+
for (const binding of spec.bindings) if (binding.type === "uniform") entries.push({ buffer: uniformBuffer });
|
|
3935
|
+
else {
|
|
3936
|
+
const tensorName = tensorNames[tensorIdx++];
|
|
3937
|
+
const buffer = this.weightBuffers.get(tensorName) ?? this.activationBuffers.get(tensorName);
|
|
3938
|
+
if (!buffer) throw new Error(`ParlerEncoderExecutor: no buffer for "${tensorName}" in op ${node.id}`);
|
|
3939
|
+
entries.push({ buffer });
|
|
3940
|
+
}
|
|
3941
|
+
return entries;
|
|
3942
|
+
}
|
|
3943
|
+
async readBack(name, elements) {
|
|
3944
|
+
const buffer = this.activationBuffers.get(name);
|
|
3945
|
+
if (!buffer) throw new Error(`ParlerEncoderExecutor: no buffer "${name}"`);
|
|
3946
|
+
const byteLen = Math.min(elements * 4, buffer.size);
|
|
3947
|
+
const readback = this.ctx.device.createBuffer({
|
|
3948
|
+
size: byteLen,
|
|
3949
|
+
usage: 9
|
|
3950
|
+
});
|
|
3951
|
+
const enc = this.ctx.device.createCommandEncoder();
|
|
3952
|
+
enc.copyBufferToBuffer(buffer, 0, readback, 0, byteLen);
|
|
3953
|
+
this.ctx.device.queue.submit([enc.finish()]);
|
|
3954
|
+
await readback.mapAsync(MAP_MODE_READ$1, 0, byteLen);
|
|
3955
|
+
const data = new Float32Array(readback.getMappedRange(0, byteLen).slice(0));
|
|
3956
|
+
readback.unmap();
|
|
3957
|
+
readback.destroy();
|
|
3958
|
+
return data;
|
|
3959
|
+
}
|
|
3960
|
+
};
|
|
3961
|
+
var ParlerTTS = class ParlerTTS {
|
|
3962
|
+
ctx;
|
|
3963
|
+
loaded;
|
|
3964
|
+
cfg;
|
|
3965
|
+
maxSeqLen;
|
|
3966
|
+
_destroyed = false;
|
|
3967
|
+
sampleRate = PARLER_SAMPLE_RATE;
|
|
3968
|
+
constructor(ctx, loaded, maxSeqLen) {
|
|
3969
|
+
this.ctx = ctx;
|
|
3970
|
+
this.loaded = loaded;
|
|
3971
|
+
this.cfg = parseParlerConfig(loaded.rawConfig);
|
|
3972
|
+
this.maxSeqLen = maxSeqLen;
|
|
3973
|
+
}
|
|
3974
|
+
static async create(options = {}) {
|
|
3975
|
+
const ctx = await initGPU();
|
|
3976
|
+
const loaded = await loadParlerTTS({
|
|
3977
|
+
repo: options.repo,
|
|
3978
|
+
revision: options.revision,
|
|
3979
|
+
hfToken: options.hfToken,
|
|
3980
|
+
cacheDir: options.cacheDir,
|
|
3981
|
+
onProgress: options.onProgress
|
|
3982
|
+
});
|
|
3983
|
+
const cfg = parseParlerConfig(loaded.rawConfig);
|
|
3984
|
+
return new ParlerTTS(ctx, loaded, Math.min(options.maxSeqLen ?? 4096, cfg.max_position));
|
|
3985
|
+
}
|
|
3986
|
+
/**
|
|
3987
|
+
* Generate a random natural-language voice description using a host-side template
|
|
3988
|
+
* bank (no model call). Pass the result as `describeVoice` for variety. A loaded
|
|
3989
|
+
* text model can write richer ones — exposed on the Gerbil class as
|
|
3990
|
+
* {@link inventVoiceDescription}.
|
|
3991
|
+
*/
|
|
3992
|
+
static randomVoiceDescription() {
|
|
3993
|
+
const gender = pick(["male", "female"]);
|
|
3994
|
+
return `A ${pick([
|
|
3995
|
+
"warm",
|
|
3996
|
+
"bright",
|
|
3997
|
+
"calm",
|
|
3998
|
+
"energetic",
|
|
3999
|
+
"gentle",
|
|
4000
|
+
"confident",
|
|
4001
|
+
"soft"
|
|
4002
|
+
])} ${gender} voice speaking at a ${pick([
|
|
4003
|
+
"slightly slow",
|
|
4004
|
+
"moderate",
|
|
4005
|
+
"slightly fast"
|
|
4006
|
+
])} pace, in a ${pick([
|
|
4007
|
+
"very clear, studio-quality recording",
|
|
4008
|
+
"clean recording with minimal background noise",
|
|
4009
|
+
"close-sounding, high-quality audio"
|
|
4010
|
+
])}.`;
|
|
4011
|
+
}
|
|
4012
|
+
/** Synthesize `text` in the described voice. Returns 44.1 kHz mono PCM. */
|
|
4013
|
+
async speak(text, opts = {}) {
|
|
4014
|
+
if (this._destroyed) throw new Error("ParlerTTS has been destroyed.");
|
|
4015
|
+
const startTime = performance.now();
|
|
4016
|
+
const description = (opts.describeVoice ?? DEFAULT_VOICE_DESCRIPTION).trim();
|
|
4017
|
+
const descIds = this.loaded.tokenizer.encode(description);
|
|
4018
|
+
const sDesc = descIds.length;
|
|
4019
|
+
const biasTable = this.loaded.modelWeights.get("text_encoder.encoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight");
|
|
4020
|
+
if (!biasTable) throw new Error("ParlerTTS: relative_attention_bias.weight missing from checkpoint.");
|
|
4021
|
+
const relBias = buildT5RelativeBias(biasTable.data, this.cfg.rel_num_buckets, this.cfg.rel_max_distance, this.cfg.enc_heads, sDesc);
|
|
4022
|
+
const encGraph = generateParlerEncoderGraph(this.loaded.rawConfig, sDesc);
|
|
4023
|
+
const encExec = new ParlerEncoderExecutor(this.ctx, encGraph, this.cfg.dec_layers, sDesc);
|
|
4024
|
+
let encK;
|
|
4025
|
+
let encV;
|
|
4026
|
+
try {
|
|
4027
|
+
encExec.uploadWeights(selectGraphWeights(encGraph, this.loaded.modelWeights));
|
|
4028
|
+
encExec.initBindGroups();
|
|
4029
|
+
const enc = await encExec.encode(new Uint32Array(descIds), relBias);
|
|
4030
|
+
encK = enc.encK;
|
|
4031
|
+
encV = enc.encV;
|
|
4032
|
+
} finally {
|
|
4033
|
+
encExec.destroy();
|
|
4034
|
+
}
|
|
4035
|
+
const promptIds = this.loaded.tokenizer.encode(text);
|
|
4036
|
+
const delayed = await this.decodeAudioTokens(promptIds, encK, encV, sDesc, opts);
|
|
4037
|
+
const { codes, frames } = revertDelayPattern(delayed.codes, this.cfg.num_codebooks, delayed.steps);
|
|
4038
|
+
if (frames <= 0) throw new Error(`ParlerTTS produced no audio frames (steps=${delayed.steps}).`);
|
|
4039
|
+
opts.onCodes?.(codes, this.cfg.num_codebooks, frames);
|
|
4040
|
+
const pcm = await this.decodeCodes(codes, frames);
|
|
4041
|
+
const totalTime = (performance.now() - startTime) / 1e3;
|
|
4042
|
+
const audioSeconds = pcm.length / PARLER_SAMPLE_RATE;
|
|
4043
|
+
console.log(`[parler] desc=${sDesc} frames=${frames} audio=${audioSeconds.toFixed(2)}s wall=${totalTime.toFixed(2)}s`);
|
|
4044
|
+
return {
|
|
4045
|
+
pcm,
|
|
4046
|
+
sampleRate: PARLER_SAMPLE_RATE,
|
|
4047
|
+
frames,
|
|
4048
|
+
audioSeconds
|
|
4049
|
+
};
|
|
4050
|
+
}
|
|
4051
|
+
/**
|
|
4052
|
+
* Run the decoder: prefill the transcript prompt rows, then AR-decode the 9-codebook
|
|
4053
|
+
* delay-pattern stream. Returns the DELAYED code grid [9, steps] (codebook-major)
|
|
4054
|
+
* plus the step count, for revertDelayPattern.
|
|
4055
|
+
*
|
|
4056
|
+
* Each step the host sums the 9 per-codebook embeddings of the current delayed codes
|
|
4057
|
+
* + the sinusoidal position, writes that single [1,H] row to `dec_input_embed`,
|
|
4058
|
+
* forwards one step (KV-cache grows over [prompt ++ audio]), and reads back 9 logit
|
|
4059
|
+
* rows — one per codebook lm_head.
|
|
4060
|
+
*/
|
|
4061
|
+
async decodeAudioTokens(promptIds, encK, encV, sEnc, opts) {
|
|
4062
|
+
const cfg = this.cfg;
|
|
4063
|
+
const H = cfg.dec_hidden;
|
|
4064
|
+
const NCB = cfg.num_codebooks;
|
|
4065
|
+
const graph = generateParlerDecoderGraph(this.loaded.rawConfig, sEnc);
|
|
4066
|
+
const maxSteps = Math.min(opts.maxFrames ?? DEFAULT_MAX_FRAMES, this.maxSeqLen - promptIds.length) + (NCB - 1);
|
|
4067
|
+
const exec = new Executor(this.ctx, graph, {
|
|
4068
|
+
maxSeqLen: this.maxSeqLen,
|
|
4069
|
+
kvMode: "f32"
|
|
4070
|
+
});
|
|
4071
|
+
const embedTokens = [];
|
|
4072
|
+
for (let cb = 0; cb < NCB; cb++) {
|
|
4073
|
+
const t = this.loaded.modelWeights.get(`decoder.model.decoder.embed_tokens.${cb}.weight`);
|
|
4074
|
+
if (!t) throw new Error(`ParlerTTS: embed_tokens.${cb} missing.`);
|
|
4075
|
+
embedTokens.push(t.data);
|
|
4076
|
+
}
|
|
4077
|
+
const embedPositions = this.loaded.modelWeights.get("decoder.model.decoder.embed_positions.weights");
|
|
4078
|
+
const embedPrompts = this.loaded.modelWeights.get("embed_prompts.weight");
|
|
4079
|
+
if (!(embedPositions && embedPrompts)) throw new Error("ParlerTTS: embed_positions / embed_prompts missing.");
|
|
4080
|
+
const posTable = embedPositions.data;
|
|
4081
|
+
const promptTable = embedPrompts.data;
|
|
4082
|
+
const grid = new Uint32Array(NCB * maxSteps);
|
|
4083
|
+
let steps = 0;
|
|
4084
|
+
try {
|
|
4085
|
+
exec.uploadWeightsMap(selectGraphWeights(graph, this.loaded.modelWeights));
|
|
4086
|
+
exec.initBindGroups();
|
|
4087
|
+
for (let i = 0; i < cfg.dec_layers; i++) {
|
|
4088
|
+
exec.writeInput(`enc_k_layer${i}`, encK[i]);
|
|
4089
|
+
exec.writeInput(`enc_v_layer${i}`, encV[i]);
|
|
4090
|
+
}
|
|
4091
|
+
exec.reset();
|
|
4092
|
+
let pos = 0;
|
|
4093
|
+
const row = new Float32Array(H);
|
|
4094
|
+
const logitOutputs = graph.outputs.map((name) => ({
|
|
4095
|
+
name,
|
|
4096
|
+
elemCount: cfg.audio_vocab
|
|
4097
|
+
}));
|
|
4098
|
+
for (const tokenId of promptIds) {
|
|
4099
|
+
addRow(row, promptTable, tokenId, H, true);
|
|
4100
|
+
addPosition(row, posTable, pos, H);
|
|
4101
|
+
exec.writeInput("dec_input_embed", row);
|
|
4102
|
+
await exec.runGraphStep([logitOutputs[0]], 1);
|
|
4103
|
+
pos++;
|
|
4104
|
+
}
|
|
4105
|
+
let prevCodes = new Uint32Array(NCB).fill(PARLER_BOS_TOKEN_ID);
|
|
4106
|
+
let frontier = 0;
|
|
4107
|
+
let finished = false;
|
|
4108
|
+
for (let step = 0; step < maxSteps && !finished; step++) {
|
|
4109
|
+
row.fill(0);
|
|
4110
|
+
for (let cb = 0; cb < NCB; cb++) addRow(row, embedTokens[cb], prevCodes[cb], H, false);
|
|
4111
|
+
addPosition(row, posTable, pos, H);
|
|
4112
|
+
exec.writeInput("dec_input_embed", row);
|
|
4113
|
+
const logitRows = await exec.runGraphStep(logitOutputs, 1);
|
|
4114
|
+
pos++;
|
|
4115
|
+
const nextCodes = new Uint32Array(NCB);
|
|
4116
|
+
for (let cb = 0; cb < NCB; cb++) {
|
|
4117
|
+
if (step < cb) {
|
|
4118
|
+
nextCodes[cb] = PARLER_BOS_TOKEN_ID;
|
|
4119
|
+
continue;
|
|
4120
|
+
}
|
|
4121
|
+
if (cb > frontier) logitRows[cb][PARLER_EOS_TOKEN_ID] = Number.NEGATIVE_INFINITY;
|
|
4122
|
+
const code = opts.greedy ? argmax$1(logitRows[cb]) : sampleLogits(logitRows[cb], opts.temperature ?? 1, opts.topK ?? 0);
|
|
4123
|
+
nextCodes[cb] = code;
|
|
4124
|
+
grid[cb * maxSteps + step] = code < PARLER_CODEBOOK_SIZE ? code : 0;
|
|
4125
|
+
}
|
|
4126
|
+
if (frontier < NCB - 1 && nextCodes[frontier] === PARLER_EOS_TOKEN_ID) frontier++;
|
|
4127
|
+
else if (frontier === NCB - 1 && nextCodes[frontier] === PARLER_EOS_TOKEN_ID) finished = true;
|
|
4128
|
+
steps = step + 1;
|
|
4129
|
+
prevCodes = nextCodes;
|
|
4130
|
+
if (pos >= this.maxSeqLen) break;
|
|
4131
|
+
}
|
|
4132
|
+
} finally {
|
|
4133
|
+
exec.destroy();
|
|
4134
|
+
}
|
|
4135
|
+
const tight = new Uint32Array(NCB * steps);
|
|
4136
|
+
for (let cb = 0; cb < NCB; cb++) for (let s = 0; s < steps; s++) tight[cb * steps + s] = grid[cb * maxSteps + s];
|
|
4137
|
+
return {
|
|
4138
|
+
codes: tight,
|
|
4139
|
+
steps
|
|
4140
|
+
};
|
|
4141
|
+
}
|
|
4142
|
+
/** Decode the [9, frames] code grid → 44.1 kHz PCM (tanh host-side). */
|
|
4143
|
+
async decodeCodes(codes, numFrames) {
|
|
4144
|
+
const graph = generateDacSpeechDecoderGraph({
|
|
4145
|
+
numFrames,
|
|
4146
|
+
numCodebooks: this.cfg.num_codebooks,
|
|
4147
|
+
decoderRates: PARLER_DECODER_RATES,
|
|
4148
|
+
decoderDim: PARLER_DAC_DECODER_DIM,
|
|
4149
|
+
latentDim: PARLER_DAC_LATENT_DIM
|
|
4150
|
+
});
|
|
4151
|
+
const pcmLen = dacOutputLength(numFrames, PARLER_DECODER_RATES);
|
|
4152
|
+
const exec = new Executor(this.ctx, graph, {
|
|
4153
|
+
maxSeqLen: numFrames,
|
|
4154
|
+
kvMode: "f32"
|
|
4155
|
+
});
|
|
4156
|
+
try {
|
|
4157
|
+
exec.uploadWeightsMap(selectGraphWeights(graph, this.loaded.codecWeights));
|
|
4158
|
+
exec.initBindGroups();
|
|
4159
|
+
exec.reset();
|
|
4160
|
+
exec.writeInput("audio_codes", codes);
|
|
4161
|
+
const raw = await exec.runGraphOutput(graph.outputs[0], pcmLen);
|
|
4162
|
+
const pcm = new Float32Array(raw.length);
|
|
4163
|
+
for (let i = 0; i < raw.length; i++) pcm[i] = Math.tanh(raw[i]);
|
|
4164
|
+
return pcm;
|
|
4165
|
+
} finally {
|
|
4166
|
+
exec.destroy();
|
|
4167
|
+
}
|
|
4168
|
+
}
|
|
4169
|
+
destroy() {
|
|
4170
|
+
if (this._destroyed) return;
|
|
4171
|
+
this._destroyed = true;
|
|
4172
|
+
this.loaded.modelWeights.clear();
|
|
4173
|
+
this.loaded.codecWeights.clear();
|
|
4174
|
+
}
|
|
4175
|
+
};
|
|
4176
|
+
const PARLER_CODEBOOK_SIZE = 1024;
|
|
4177
|
+
/** Random pick from a list. */
|
|
4178
|
+
function pick(arr) {
|
|
4179
|
+
return arr[Math.floor(Math.random() * arr.length)];
|
|
4180
|
+
}
|
|
4181
|
+
/** Add (or set) the embedding row for `tokenId` into `row`. */
|
|
4182
|
+
function addRow(row, table, tokenId, H, set) {
|
|
4183
|
+
const base = tokenId * H;
|
|
4184
|
+
if (set) for (let i = 0; i < H; i++) row[i] = table[base + i];
|
|
4185
|
+
else for (let i = 0; i < H; i++) row[i] += table[base + i];
|
|
4186
|
+
}
|
|
4187
|
+
/** Add the sinusoidal position row (from the learned embed_positions table). */
|
|
4188
|
+
function addPosition(row, posTable, pos, H) {
|
|
4189
|
+
const base = pos * H;
|
|
4190
|
+
for (let i = 0; i < H; i++) row[i] += posTable[base + i];
|
|
4191
|
+
}
|
|
4192
|
+
|
|
4193
|
+
//#endregion
|
|
4194
|
+
//#region src/gpu/sampler.ts
|
|
4195
|
+
let _heapIndices = null;
|
|
4196
|
+
let _heapValues = null;
|
|
4197
|
+
/**
|
|
4198
|
+
* Sample a token ID from logits.
|
|
4199
|
+
*
|
|
4200
|
+
* Pipeline: repetition penalty → temperature → top-k (min-heap) → softmax → top-p → sample.
|
|
4201
|
+
*/
|
|
4202
|
+
function sampleToken(logits, params = {}, previousTokens) {
|
|
4203
|
+
const temperature = params.temperature ?? .7;
|
|
4204
|
+
const topK = params.topK ?? 50;
|
|
4205
|
+
const topP = params.topP ?? .9;
|
|
4206
|
+
const repetitionPenalty = params.repetitionPenalty ?? 1;
|
|
4207
|
+
if (temperature < 1e-6) return argmax(logits);
|
|
4208
|
+
const N = logits.length;
|
|
4209
|
+
const K = Math.min(topK > 0 ? topK : N, N);
|
|
4210
|
+
if (!_heapIndices || _heapIndices.length < K) {
|
|
4211
|
+
_heapIndices = new Uint32Array(K);
|
|
4212
|
+
_heapValues = new Float32Array(K);
|
|
4213
|
+
}
|
|
4214
|
+
const hIdx = _heapIndices;
|
|
4215
|
+
const hVal = _heapValues;
|
|
4216
|
+
let penaltySet = null;
|
|
4217
|
+
if (repetitionPenalty !== 1 && previousTokens?.length) penaltySet = new Set(previousTokens);
|
|
4218
|
+
let heapSize = 0;
|
|
4219
|
+
for (let i = 0; i < N; i++) {
|
|
4220
|
+
let s = logits[i];
|
|
4221
|
+
if (penaltySet?.has(i)) s = s > 0 ? s / repetitionPenalty : s * repetitionPenalty;
|
|
4222
|
+
s /= temperature;
|
|
4223
|
+
if (heapSize < K) {
|
|
4224
|
+
hIdx[heapSize] = i;
|
|
4225
|
+
hVal[heapSize] = s;
|
|
4226
|
+
heapSize++;
|
|
4227
|
+
if (heapSize === K) for (let j = (K >> 1) - 1; j >= 0; j--) siftDown(hIdx, hVal, j, K);
|
|
4228
|
+
} else if (s > hVal[0]) {
|
|
4229
|
+
hIdx[0] = i;
|
|
4230
|
+
hVal[0] = s;
|
|
4231
|
+
siftDown(hIdx, hVal, 0, K);
|
|
4232
|
+
}
|
|
4233
|
+
}
|
|
4234
|
+
for (let i = 1; i < heapSize; i++) {
|
|
4235
|
+
const vi = hVal[i];
|
|
4236
|
+
const ii = hIdx[i];
|
|
4237
|
+
let j = i - 1;
|
|
4238
|
+
while (j >= 0 && hVal[j] < vi) {
|
|
4239
|
+
hVal[j + 1] = hVal[j];
|
|
4240
|
+
hIdx[j + 1] = hIdx[j];
|
|
4241
|
+
j--;
|
|
4242
|
+
}
|
|
4243
|
+
hVal[j + 1] = vi;
|
|
4244
|
+
hIdx[j + 1] = ii;
|
|
4245
|
+
}
|
|
4246
|
+
const maxScore = hVal[0];
|
|
4247
|
+
let sumExp = 0;
|
|
4248
|
+
for (let i = 0; i < heapSize; i++) {
|
|
4249
|
+
const p = Math.exp(hVal[i] - maxScore);
|
|
4250
|
+
hVal[i] = p;
|
|
4251
|
+
sumExp += p;
|
|
4252
|
+
}
|
|
4253
|
+
const invSum = 1 / sumExp;
|
|
4254
|
+
for (let i = 0; i < heapSize; i++) hVal[i] *= invSum;
|
|
4255
|
+
let candidateCount = heapSize;
|
|
4256
|
+
if (topP < 1) {
|
|
4257
|
+
let cumulative$1 = 0;
|
|
4258
|
+
for (let i = 0; i < heapSize; i++) {
|
|
4259
|
+
cumulative$1 += hVal[i];
|
|
4260
|
+
if (cumulative$1 >= topP) {
|
|
4261
|
+
candidateCount = i + 1;
|
|
4262
|
+
break;
|
|
4263
|
+
}
|
|
4264
|
+
}
|
|
4265
|
+
let sum = 0;
|
|
4266
|
+
for (let i = 0; i < candidateCount; i++) sum += hVal[i];
|
|
4267
|
+
const inv = 1 / sum;
|
|
4268
|
+
for (let i = 0; i < candidateCount; i++) hVal[i] *= inv;
|
|
4269
|
+
}
|
|
4270
|
+
const r = Math.random();
|
|
4271
|
+
let cumulative = 0;
|
|
4272
|
+
for (let i = 0; i < candidateCount; i++) {
|
|
4273
|
+
cumulative += hVal[i];
|
|
4274
|
+
if (r <= cumulative) return hIdx[i];
|
|
4275
|
+
}
|
|
4276
|
+
return hIdx[candidateCount - 1];
|
|
4277
|
+
}
|
|
4278
|
+
/**
|
|
4279
|
+
* Return the index of the maximum value (greedy decoding).
|
|
4280
|
+
*/
|
|
4281
|
+
function argmax(arr) {
|
|
4282
|
+
let maxIdx = 0;
|
|
4283
|
+
let maxVal = arr[0];
|
|
4284
|
+
for (let i = 1; i < arr.length; i++) if (arr[i] > maxVal) {
|
|
4285
|
+
maxVal = arr[i];
|
|
4286
|
+
maxIdx = i;
|
|
4287
|
+
}
|
|
4288
|
+
return maxIdx;
|
|
4289
|
+
}
|
|
4290
|
+
/** Min-heap sift down on parallel index/value typed arrays. */
|
|
4291
|
+
function siftDown(indices, values, i, n) {
|
|
4292
|
+
while (true) {
|
|
4293
|
+
let smallest = i;
|
|
4294
|
+
const left = 2 * i + 1;
|
|
4295
|
+
const right = 2 * i + 2;
|
|
4296
|
+
if (left < n && values[left] < values[smallest]) smallest = left;
|
|
4297
|
+
if (right < n && values[right] < values[smallest]) smallest = right;
|
|
4298
|
+
if (smallest === i) break;
|
|
4299
|
+
const ti = indices[i];
|
|
4300
|
+
indices[i] = indices[smallest];
|
|
4301
|
+
indices[smallest] = ti;
|
|
4302
|
+
const tv = values[i];
|
|
4303
|
+
values[i] = values[smallest];
|
|
4304
|
+
values[smallest] = tv;
|
|
4305
|
+
i = smallest;
|
|
4306
|
+
}
|
|
4307
|
+
}
|
|
4308
|
+
|
|
4309
|
+
//#endregion
|
|
4310
|
+
//#region src/gpu/vision-executor.ts
|
|
4311
|
+
const MAP_MODE_READ = 1;
|
|
4312
|
+
/** Pack an f32 array to IEEE half (f16) bits (round-to-nearest-even). */
|
|
4313
|
+
const _f16scratch = /* @__PURE__ */ new ArrayBuffer(4);
|
|
4314
|
+
const _f16f32 = new Float32Array(_f16scratch);
|
|
4315
|
+
const _f16u32 = new Uint32Array(_f16scratch);
|
|
4316
|
+
function packF16(src) {
|
|
4317
|
+
const out = new Uint16Array(src.length);
|
|
4318
|
+
for (let i = 0; i < src.length; i++) {
|
|
4319
|
+
_f16f32[0] = src[i];
|
|
4320
|
+
const f = _f16u32[0];
|
|
4321
|
+
const sign = f >>> 16 & 32768;
|
|
4322
|
+
const exp = (f >>> 23 & 255) - 112;
|
|
4323
|
+
const mant = f & 8388607;
|
|
4324
|
+
if (exp <= 0) out[i] = sign;
|
|
4325
|
+
else if (exp >= 31) out[i] = sign | 31744;
|
|
4326
|
+
else {
|
|
4327
|
+
const round = (mant & 4096) !== 0;
|
|
4328
|
+
let half = sign | exp << 10 | mant >>> 13;
|
|
4329
|
+
if (round) half += 1;
|
|
4330
|
+
out[i] = half;
|
|
4331
|
+
}
|
|
4332
|
+
}
|
|
4333
|
+
return out;
|
|
4334
|
+
}
|
|
4335
|
+
/**
|
|
4336
|
+
* WebKit only: the ViT bidirectional Attention op is a single dispatch whose
|
|
4337
|
+
* grid is [N query rows, heads], each row looping over all N keys — O(N²) work in
|
|
4338
|
+
* one submission. Past a few hundred rows that can exceed Metal's ~few-second GPU
|
|
4339
|
+
* watchdog and kill the web-content process (indistinguishable from OOM). We split
|
|
4340
|
+
* the attention grid into windows of this many query rows, each its own
|
|
4341
|
+
* submit + onSubmittedWorkDone() drain, so no single submission runs long enough
|
|
4342
|
+
* to trip the watchdog. Math is unchanged: every query row computes identically
|
|
4343
|
+
* regardless of how the grid is partitioned.
|
|
4344
|
+
*/
|
|
4345
|
+
const WEBKIT_ATTN_CHUNK_ROWS = 64;
|
|
4346
|
+
var VisionExecutor = class {
|
|
4347
|
+
ctx;
|
|
4348
|
+
graph;
|
|
4349
|
+
mergeUnit;
|
|
4350
|
+
weightBuffers = /* @__PURE__ */ new Map();
|
|
4351
|
+
activationBuffers = /* @__PURE__ */ new Map();
|
|
4352
|
+
dispatches = [];
|
|
4353
|
+
maxPatches;
|
|
4354
|
+
/** Weight (B) names of MatMulBias nodes stored as f16 (empty without shader-f16). */
|
|
4355
|
+
f16WeightNames = /* @__PURE__ */ new Set();
|
|
4356
|
+
/** Runtime pooled-token count for the Gemma 4 ViT ("Np" dim); 0 for Qwen. */
|
|
4357
|
+
gemma4Np = 0;
|
|
4358
|
+
/** True when this graph is the Gemma 4 vision tower (uses "Np" + PoolMatMul). */
|
|
4359
|
+
isGemma4;
|
|
4360
|
+
constructor(ctx, graph, maxPatches) {
|
|
4361
|
+
this.ctx = ctx;
|
|
4362
|
+
this.graph = graph;
|
|
4363
|
+
this.maxPatches = maxPatches;
|
|
4364
|
+
this.isGemma4 = graph.architecture === "Gemma4VisionModel";
|
|
4365
|
+
const mergedIn = graph.tensors[CANONICAL_KEYS.visMergerFc1W]?.shape[1];
|
|
4366
|
+
const hidden = graph.config.hidden_size;
|
|
4367
|
+
this.mergeUnit = mergedIn && hidden ? mergedIn / hidden : 1;
|
|
4368
|
+
if (ctx.hasF16) {
|
|
4369
|
+
for (const node of graph.nodes) if (node.opType === "MatMulBias" && node.inputs[1]) this.f16WeightNames.add(node.inputs[1]);
|
|
4370
|
+
}
|
|
4371
|
+
this.allocateActivationBuffers();
|
|
4372
|
+
}
|
|
4373
|
+
uploadWeights(weights) {
|
|
4374
|
+
for (const [name, desc] of Object.entries(this.graph.tensors)) {
|
|
4375
|
+
if (desc.storage !== "constant") continue;
|
|
4376
|
+
const w = weights.get(name);
|
|
4377
|
+
if (!w) throw new Error(`VisionExecutor: missing weight "${name}"`);
|
|
4378
|
+
let data = w.data;
|
|
4379
|
+
if (this.f16WeightNames.has(name) && w.data instanceof Float32Array) data = packF16(w.data);
|
|
4380
|
+
const buffer = createStorageBuffer(this.ctx, `vweight_${name}`, data.byteLength, data);
|
|
4381
|
+
this.weightBuffers.set(name, buffer);
|
|
4382
|
+
}
|
|
4383
|
+
}
|
|
4384
|
+
initBindGroups() {
|
|
4385
|
+
const dummyShapes = this.resolveShapes(this.maxPatches);
|
|
3120
4386
|
for (const nodeId of this.graph.executionOrder) {
|
|
3121
4387
|
const node = this.graph.nodes.find((n) => n.id === nodeId);
|
|
3122
4388
|
let spec = KERNEL_REGISTRY[node.opType];
|
|
@@ -4281,6 +5547,8 @@ var WebGPUEngine = class WebGPUEngine {
|
|
|
4281
5547
|
_outeTTS = null;
|
|
4282
5548
|
/** The preset voice the lazily-built OuteTTS engine was created with. */
|
|
4283
5549
|
_outeVoice = null;
|
|
5550
|
+
/** Lazily-created Parler-TTS engine (Flan-T5 encoder + decoder LM + dac_44khz). */
|
|
5551
|
+
_parlerTTS = null;
|
|
4284
5552
|
/**
|
|
4285
5553
|
* WebKit group-size probe state. When true, a candidate group size is being
|
|
4286
5554
|
* tried this page-load and must be promoted (or capped) after the FIRST
|
|
@@ -4945,6 +6213,7 @@ var WebGPUEngine = class WebGPUEngine {
|
|
|
4945
6213
|
async speak(text, options = {}) {
|
|
4946
6214
|
this.checkDestroyed();
|
|
4947
6215
|
const repo = this._createOptions.repo ?? "";
|
|
6216
|
+
if (/parler/i.test(repo) || /parler/i.test(this._architecture)) return this.speakParler(text, options);
|
|
4948
6217
|
if (/outetts/i.test(repo) || /oute/i.test(this._architecture)) return this.speakOute(text, options);
|
|
4949
6218
|
if (!(this._architecture === "KaniTTS2ForCausalLM" || /kani-tts/i.test(repo) || /\btts\b/i.test(repo))) throw new Error(`speak() requires a Kani-TTS or OuteTTS checkpoint (e.g. "${DEFAULT_MODELS.tts}" or "${DEFAULT_MODELS.ttsOute}"), loaded engine is "${this._architecture}" (repo "${repo}").`);
|
|
4950
6219
|
if (!this._kaniTTS) this._kaniTTS = await KaniTTS.create({
|
|
@@ -4965,6 +6234,31 @@ var WebGPUEngine = class WebGPUEngine {
|
|
|
4965
6234
|
});
|
|
4966
6235
|
}
|
|
4967
6236
|
/**
|
|
6237
|
+
* Parler-TTS speak path: lazily build the Parler engine (Flan-T5 description encoder
|
|
6238
|
+
* + cross-attending decoder LM + dac_44khz) and synthesize `text` in the voice given
|
|
6239
|
+
* by `describeVoice`. Single-checkpoint (all weights, incl. the DAC, in one file).
|
|
6240
|
+
*/
|
|
6241
|
+
async speakParler(text, options) {
|
|
6242
|
+
if (!this._parlerTTS) this._parlerTTS = await ParlerTTS.create({
|
|
6243
|
+
repo: this._createOptions.repo,
|
|
6244
|
+
revision: this._createOptions.revision,
|
|
6245
|
+
hfToken: this._createOptions.hfToken,
|
|
6246
|
+
cacheDir: this._createOptions.cacheDir,
|
|
6247
|
+
maxSeqLen: this.maxSeqLen
|
|
6248
|
+
});
|
|
6249
|
+
const speakOpts = {
|
|
6250
|
+
describeVoice: options.describeVoice,
|
|
6251
|
+
temperature: options.temperature,
|
|
6252
|
+
topK: options.topK,
|
|
6253
|
+
maxFrames: options.maxFrames
|
|
6254
|
+
};
|
|
6255
|
+
return this._parlerTTS.speak(text, speakOpts);
|
|
6256
|
+
}
|
|
6257
|
+
/** A random natural-language voice description for Parler-TTS (host template bank). */
|
|
6258
|
+
randomVoiceDescription() {
|
|
6259
|
+
return ParlerTTS.randomVoiceDescription();
|
|
6260
|
+
}
|
|
6261
|
+
/**
|
|
4968
6262
|
* OuteTTS speak path: lazily build the OuteTTS engine (loading the chosen preset
|
|
4969
6263
|
* voice's speaker JSON + the folded DAC codec weights from the mirror), rebuilding
|
|
4970
6264
|
* if the requested voice changed. Mirrors the Kani public shape (temperature/topP/topK).
|
|
@@ -5452,12 +6746,12 @@ var WebGPUEngine = class WebGPUEngine {
|
|
|
5452
6746
|
const f = data instanceof Float32Array ? data : new Float32Array(data.buffer);
|
|
5453
6747
|
let sum = 0;
|
|
5454
6748
|
let maxVal = -Infinity;
|
|
5455
|
-
let argmax$
|
|
6749
|
+
let argmax$3 = 0;
|
|
5456
6750
|
for (let i = 0; i < f.length; i++) {
|
|
5457
6751
|
sum += f[i];
|
|
5458
6752
|
if (f[i] > maxVal) {
|
|
5459
6753
|
maxVal = f[i];
|
|
5460
|
-
argmax$
|
|
6754
|
+
argmax$3 = i;
|
|
5461
6755
|
}
|
|
5462
6756
|
}
|
|
5463
6757
|
const entry = {
|
|
@@ -5470,7 +6764,7 @@ var WebGPUEngine = class WebGPUEngine {
|
|
|
5470
6764
|
Math.round(f[2] * 1e6) / 1e6,
|
|
5471
6765
|
Math.round(f[3] * 1e6) / 1e6
|
|
5472
6766
|
],
|
|
5473
|
-
argmax: argmax$
|
|
6767
|
+
argmax: argmax$3,
|
|
5474
6768
|
maxVal: Math.round(maxVal * 1e4) / 1e4
|
|
5475
6769
|
};
|
|
5476
6770
|
checks.push(entry);
|
|
@@ -5725,6 +7019,7 @@ var WebGPUEngine = class WebGPUEngine {
|
|
|
5725
7019
|
this.executor.destroy();
|
|
5726
7020
|
this._kaniTTS?.destroy();
|
|
5727
7021
|
this._outeTTS?.destroy();
|
|
7022
|
+
this._parlerTTS?.destroy();
|
|
5728
7023
|
this.visionExecutor?.destroy();
|
|
5729
7024
|
clearPipelineCache(this.ctx.device);
|
|
5730
7025
|
this.ctx.device.destroy();
|
|
@@ -5735,5 +7030,5 @@ var WebGPUEngine = class WebGPUEngine {
|
|
|
5735
7030
|
};
|
|
5736
7031
|
|
|
5737
7032
|
//#endregion
|
|
5738
|
-
export {
|
|
5739
|
-
//# sourceMappingURL=gpu-
|
|
7033
|
+
export { dequantizeGemma4VisionProjection as A, audioTokensToDacCodes as C, parseOuteTtsConfig as D, generateOuteTtsBackboneGraph as E, generateGemma4VisionGraph as M, patchGemma4VisionClips as N, KaniTTS as O, resolveGemma4VisionInfo as P, loadOuteSpeaker as S, generateDacSpeechDecoderGraph as T, smartResize as _, buildGemma4PosEmbeds as a, OuteTTS as b, buildMRoPECosSin as c, buildPositionIds as d, buildRotaryCosSin as f, preprocessImageGemma4 as g, preprocessImage as h, buildGemma4PoolMatrix as i, dequantizeMLXProjection as j, generateQwen3_5VisionGraph as k, buildMRoPEPositionIds as l, mropeFreqDims as m, GEMMA4_IMAGE_PROCESSOR as n, buildGemma4RotaryCosSin as o, buildVisionPositionTensors as p, QWEN3_5_IMAGE_PROCESSOR as r, buildGemma4VisionPositionTensors as s, WebGPUEngine as t, buildPosEmbeds as u, VisionExecutor as v, dacOutputLength as w, buildOutePromptString as x, ParlerTTS as y };
|
|
7034
|
+
//# sourceMappingURL=gpu-DbB321DN.mjs.map
|