@tryhamster/gerbil 1.2.2 → 1.3.1
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/README.md +6 -6
- package/dist/architectures-Ol7yUmKf.mjs.map +1 -1
- package/dist/browser/index.d.ts +68 -2
- 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 +47 -0
- package/dist/defaults-CudhglS1.mjs.map +1 -0
- package/dist/frameworks/express.d.mts +2 -1
- package/dist/frameworks/express.d.mts.map +1 -1
- package/dist/frameworks/express.mjs +1 -1
- package/dist/frameworks/fastify.d.mts +2 -1
- package/dist/frameworks/fastify.d.mts.map +1 -1
- package/dist/frameworks/fastify.mjs +1 -1
- package/dist/frameworks/hono.d.mts +2 -1
- package/dist/frameworks/hono.d.mts.map +1 -1
- package/dist/frameworks/hono.mjs +1 -1
- package/dist/frameworks/next.d.mts +4 -3
- package/dist/frameworks/next.d.mts.map +1 -1
- package/dist/frameworks/next.mjs +1 -1
- package/dist/frameworks/react.d.mts +2 -1
- package/dist/frameworks/react.d.mts.map +1 -1
- package/dist/frameworks/trpc.d.mts +2 -1
- package/dist/frameworks/trpc.d.mts.map +1 -1
- package/dist/frameworks/trpc.mjs +1 -1
- package/dist/{gerbil-DkilSf1Q.d.mts → gerbil-4Xid85_Z.d.mts} +31 -6
- package/dist/gerbil-4Xid85_Z.d.mts.map +1 -0
- package/dist/gerbil-BikD0xFU.mjs +4 -0
- package/dist/{gerbil-BMZ5wnM_.mjs → gerbil-Du-1M2Wq.mjs} +52 -14
- package/dist/gerbil-Du-1M2Wq.mjs.map +1 -0
- package/dist/gpu/hooks.d.mts +51 -3
- package/dist/gpu/hooks.d.mts.map +1 -1
- package/dist/gpu/hooks.mjs +130 -2
- package/dist/gpu/hooks.mjs.map +1 -1
- package/dist/gpu/index.d.mts +3 -2
- package/dist/gpu/index.mjs +4 -4
- package/dist/{gpu-BZHJDas1.mjs → gpu-BECDib8p.mjs} +2085 -217
- package/dist/gpu-BECDib8p.mjs.map +1 -0
- package/dist/{index-HVS7yeLx.d.mts → index-C4kJrqs1.d.mts} +99 -5
- package/dist/index-C4kJrqs1.d.mts.map +1 -0
- package/dist/{index-Dgmb2kE3.d.mts → index-CRDf-nS4.d.mts} +2 -2
- package/dist/{index-Dgmb2kE3.d.mts.map → index-CRDf-nS4.d.mts.map} +1 -1
- package/dist/index.d.mts +6 -5
- package/dist/index.d.mts.map +1 -1
- package/dist/index.mjs +5 -5
- package/dist/integrations/ai-sdk.d.mts +2 -1
- package/dist/integrations/ai-sdk.d.mts.map +1 -1
- package/dist/integrations/ai-sdk.mjs +1 -1
- package/dist/integrations/langchain.d.mts +2 -1
- package/dist/integrations/langchain.d.mts.map +1 -1
- package/dist/integrations/langchain.mjs +1 -1
- package/dist/integrations/llamaindex.d.mts +2 -1
- package/dist/integrations/llamaindex.d.mts.map +1 -1
- package/dist/integrations/llamaindex.mjs +1 -1
- package/dist/integrations/mcp.d.mts +4 -3
- package/dist/integrations/mcp.d.mts.map +1 -1
- package/dist/integrations/mcp.mjs +4 -4
- package/dist/{mcp-D1laYce-.mjs → mcp-4imtBQLX.mjs} +3 -3
- package/dist/{mcp-D1laYce-.mjs.map → mcp-4imtBQLX.mjs.map} +1 -1
- package/dist/memory/index.d.mts +2 -2
- package/dist/{moonshine-stt-Bj1zJfZg.mjs → moonshine-stt-BO5gZ_D6.mjs} +300 -30
- package/dist/moonshine-stt-BO5gZ_D6.mjs.map +1 -0
- package/dist/moonshine-stt-fVw-UB4s.mjs +4 -0
- package/dist/{one-liner-qLcYdOfH.mjs → one-liner-X5vMbWzs.mjs} +2 -2
- package/dist/{one-liner-qLcYdOfH.mjs.map → one-liner-X5vMbWzs.mjs.map} +1 -1
- package/dist/outetts-DmR0NfSG.d.mts +219 -0
- package/dist/outetts-DmR0NfSG.d.mts.map +1 -0
- package/dist/{repl-RdQ2vKRz.mjs → repl-CGadrKcY.mjs} +3 -3
- package/dist/skills/index.d.mts +6 -5
- package/dist/skills/index.d.mts.map +1 -1
- package/dist/skills/index.mjs +3 -3
- package/dist/{skills-_q7U_ySm.mjs → skills-BFGAXYkm.mjs} +2 -2
- package/dist/{skills-_q7U_ySm.mjs.map → skills-BFGAXYkm.mjs.map} +1 -1
- package/dist/{types-LlyYILII.d.mts → types-Cpf-aOxX.d.mts} +25 -3
- package/dist/types-Cpf-aOxX.d.mts.map +1 -0
- package/dist/{types-DQBe2lFo.d.mts → types-Dqv4m0vo.d.mts} +1 -1
- package/dist/{types-DQBe2lFo.d.mts.map → types-Dqv4m0vo.d.mts.map} +1 -1
- package/package.json +2 -2
- package/dist/defaults-bJcmOaRG.mjs +0 -24
- package/dist/defaults-bJcmOaRG.mjs.map +0 -1
- package/dist/gerbil-BMZ5wnM_.mjs.map +0 -1
- package/dist/gerbil-BWKJtIYj.mjs +0 -4
- package/dist/gerbil-DkilSf1Q.d.mts.map +0 -1
- package/dist/gpu-BZHJDas1.mjs.map +0 -1
- package/dist/index-HVS7yeLx.d.mts.map +0 -1
- package/dist/moonshine-stt-Bj1zJfZg.mjs.map +0 -1
- package/dist/moonshine-stt-Eei7fh3N.mjs +0 -4
- package/dist/types-LlyYILII.d.mts.map +0 -1
|
@@ -1,6 +1,6 @@
|
|
|
1
|
-
import {
|
|
2
|
-
import { E as GEMMA4_VIS_KEYS, 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 { _ as
|
|
1
|
+
import { i as resolveDefaultRepo, n as OUTETTS_ASSETS, r as OUTETTS_PRESET_VOICES, t as DEFAULT_MODELS } from "./defaults-CudhglS1.mjs";
|
|
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 { _ as destroyBuffers, b as verifyGPU, c as quantizeKaniBackbone, d as KERNEL_REGISTRY, f as MATMUL_BIAS_F16C_SPEC, g as createUniformBuffer, h as createStorageBuffer, i as loadModel, l as remapPrunedToken, m as createBindGroup, o as loadOuteTTS, p as clearPipelineCache, r as loadKaniTTS, s as quantizeBackboneInt4, u as Executor, v as getOrCreatePipeline, y as initGPU } from "./moonshine-stt-BO5gZ_D6.mjs";
|
|
4
4
|
|
|
5
5
|
//#region src/gpu/architectures/gemma4_vision.ts
|
|
6
6
|
/**
|
|
@@ -765,9 +765,9 @@ const KANI_END_OF_TEXT = 2;
|
|
|
765
765
|
const KANI_CODEC_CHUNK_FRAMES = 32;
|
|
766
766
|
const KANI_CODEC_LOOKBACK_FRAMES = 32;
|
|
767
767
|
/** Hard cap so a stuck decode cannot loop forever (reference uses 3000). */
|
|
768
|
-
const DEFAULT_MAX_NEW_TOKENS = 3e3;
|
|
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(graph, weights) {
|
|
770
|
+
function selectGraphWeights$1(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;
|
|
@@ -781,7 +781,7 @@ function selectGraphWeights(graph, weights) {
|
|
|
781
781
|
* Sorts by prob desc, keeps the smallest prefix whose cumulative mass ≥ topP,
|
|
782
782
|
* renormalizes, and samples one id.
|
|
783
783
|
*/
|
|
784
|
-
function nucleusSample(probs, ids, topP) {
|
|
784
|
+
function nucleusSample$1(probs, ids, topP) {
|
|
785
785
|
const n = probs.length;
|
|
786
786
|
const order = Array.from({ length: n }, (_, i) => i).sort((a, b) => probs[b] - probs[a]);
|
|
787
787
|
let cum = 0;
|
|
@@ -806,7 +806,7 @@ function nucleusSample(probs, ids, topP) {
|
|
|
806
806
|
* In-place top-k filter over a probability array: keep only the `k` highest, zero
|
|
807
807
|
* the rest, and renormalize. No-op when k <= 0 or k >= length (keep everything).
|
|
808
808
|
*/
|
|
809
|
-
function applyTopK(scores, k) {
|
|
809
|
+
function applyTopK$1(scores, k) {
|
|
810
810
|
const n = scores.length;
|
|
811
811
|
if (k <= 0 || k >= n) return;
|
|
812
812
|
const order = Array.from({ length: n }, (_, i) => i).sort((a, b) => scores[b] - scores[a]);
|
|
@@ -816,7 +816,7 @@ function applyTopK(scores, k) {
|
|
|
816
816
|
if (keepSum > 0) for (let i = 0; i < n; i++) scores[i] /= keepSum;
|
|
817
817
|
}
|
|
818
818
|
/** In-place softmax of `scores`. */
|
|
819
|
-
function softmaxInPlace(scores) {
|
|
819
|
+
function softmaxInPlace$1(scores) {
|
|
820
820
|
let maxV = Number.NEGATIVE_INFINITY;
|
|
821
821
|
for (let i = 0; i < scores.length; i++) if (scores[i] > maxV) maxV = scores[i];
|
|
822
822
|
let sum = 0;
|
|
@@ -834,234 +834,2053 @@ var KaniTTS = class KaniTTS {
|
|
|
834
834
|
cfg;
|
|
835
835
|
rawConfig;
|
|
836
836
|
maxSeqLen;
|
|
837
|
-
/** Backbone executor (built once; reused across speak() calls). */
|
|
837
|
+
/** Backbone executor (built once; reused across speak() calls). */
|
|
838
|
+
backboneExec;
|
|
839
|
+
/** The attention layer indices (carry learnable α) and their α values. */
|
|
840
|
+
attnLayers;
|
|
841
|
+
layerAlpha;
|
|
842
|
+
headDim;
|
|
843
|
+
ropeBase;
|
|
844
|
+
_destroyed = false;
|
|
845
|
+
/** Accent/voice tags the loaded checkpoint advertises (empty = single voice). */
|
|
846
|
+
voiceTags;
|
|
847
|
+
/** The loaded checkpoint's architecture string ("KaniTTS2ForCausalLM" or "Lfm2ForCausalLM"). */
|
|
848
|
+
architecture;
|
|
849
|
+
/** The selected backbone precision ("q4" or "f32"). */
|
|
850
|
+
dtype;
|
|
851
|
+
constructor(ctx, loaded, maxSeqLen, dtype) {
|
|
852
|
+
this.ctx = ctx;
|
|
853
|
+
this.loaded = loaded;
|
|
854
|
+
this.tokenizer = loaded.tokenizer;
|
|
855
|
+
this.rawConfig = loaded.rawConfig;
|
|
856
|
+
this.cfg = parseKaniConfig(loaded.rawConfig);
|
|
857
|
+
this.maxSeqLen = maxSeqLen;
|
|
858
|
+
this.dtype = loaded.preQuantized ? "q4" : dtype;
|
|
859
|
+
this.architecture = (loaded.rawConfig.architectures ?? [])[0] ?? "Lfm2ForCausalLM";
|
|
860
|
+
const langSettings = loaded.rawConfig.language_settings;
|
|
861
|
+
this.voiceTags = Array.isArray(langSettings?.language_tags_list) ? langSettings.language_tags_list.filter((v) => typeof v === "string") : [];
|
|
862
|
+
const hidden = this.rawConfig.hidden_size;
|
|
863
|
+
const heads = this.rawConfig.num_attention_heads;
|
|
864
|
+
this.headDim = this.rawConfig.head_dim ?? Math.floor(hidden / heads);
|
|
865
|
+
this.ropeBase = this.rawConfig.rope_theta ?? 1e6;
|
|
866
|
+
this.attnLayers = kaniAttentionLayerIndices(this.rawConfig);
|
|
867
|
+
this.layerAlpha = /* @__PURE__ */ new Map();
|
|
868
|
+
const useLearnable = this.rawConfig.use_learnable_rope ?? false;
|
|
869
|
+
for (const layer of this.attnLayers) {
|
|
870
|
+
let alpha = 1;
|
|
871
|
+
if (useLearnable) {
|
|
872
|
+
const w = loaded.backboneWeights.get(`learnable_rope_layers.${layer}.alpha_weight`);
|
|
873
|
+
alpha = kaniLayerAlpha(w ? w.data[0] ?? 0 : 0, this.cfg);
|
|
874
|
+
}
|
|
875
|
+
this.layerAlpha.set(layer, alpha);
|
|
876
|
+
}
|
|
877
|
+
this.backboneExec = this.buildBackbone(ctx, loaded, maxSeqLen);
|
|
878
|
+
}
|
|
879
|
+
/**
|
|
880
|
+
* Build the codec-LM backbone graph, quantize it when q4, and upload its weights
|
|
881
|
+
* into a fresh Executor. For a pre-quantized (gerbil-int4) checkpoint the linear
|
|
882
|
+
* projections already carry `.q`/`.scales`/`.zeros`, so quantizeKaniBackbone is an
|
|
883
|
+
* idempotent no-op (it skips any projection whose `.q` is already present) and the
|
|
884
|
+
* int4 weights are uploaded verbatim — no on-device re-quantization.
|
|
885
|
+
*/
|
|
886
|
+
buildBackbone(ctx, loaded, maxSeqLen) {
|
|
887
|
+
const useQ4 = this.dtype === "q4";
|
|
888
|
+
const graph = generateKaniTtsGraph(this.rawConfig, useQ4 ? "q4" : void 0);
|
|
889
|
+
if (useQ4) quantizeKaniBackbone(graph, loaded.backboneWeights);
|
|
890
|
+
const exec = new Executor(ctx, graph, {
|
|
891
|
+
maxSeqLen,
|
|
892
|
+
kvMode: "f32"
|
|
893
|
+
});
|
|
894
|
+
exec.uploadWeightsMap(selectGraphWeights$1(graph, loaded.backboneWeights));
|
|
895
|
+
exec.initBindGroups();
|
|
896
|
+
return exec;
|
|
897
|
+
}
|
|
898
|
+
static async create(options = {}) {
|
|
899
|
+
const ctx = await initGPU();
|
|
900
|
+
const loaded = await loadKaniTTS({
|
|
901
|
+
repo: options.repo,
|
|
902
|
+
codecRepo: options.codecRepo,
|
|
903
|
+
revision: options.revision,
|
|
904
|
+
hfToken: options.hfToken,
|
|
905
|
+
cacheDir: options.cacheDir,
|
|
906
|
+
onProgress: options.onProgress
|
|
907
|
+
});
|
|
908
|
+
const ctxLen = loaded.rawConfig.max_position_embeddings ?? 4096;
|
|
909
|
+
return new KaniTTS(ctx, loaded, Math.min(options.maxSeqLen ?? 2048, ctxLen), options.dtype ?? "q4");
|
|
910
|
+
}
|
|
911
|
+
/** Write per-layer cos/sin for token rows [rowStart, rowStart+positions.length). */
|
|
912
|
+
writeCosSin(positions, rowStart) {
|
|
913
|
+
const rowBytes = this.headDim * 4;
|
|
914
|
+
for (const layer of this.attnLayers) {
|
|
915
|
+
const alpha = this.layerAlpha.get(layer) ?? 1;
|
|
916
|
+
const { cos, sin } = buildKaniLayerCosSin(positions, this.headDim, this.ropeBase, alpha);
|
|
917
|
+
if (rowStart === 0) {
|
|
918
|
+
this.backboneExec.writeInput(kaniCosTensor(layer), cos);
|
|
919
|
+
this.backboneExec.writeInput(kaniSinTensor(layer), sin);
|
|
920
|
+
} else {
|
|
921
|
+
this.backboneExec.writeInputAt(kaniCosTensor(layer), cos, rowStart * rowBytes);
|
|
922
|
+
this.backboneExec.writeInputAt(kaniSinTensor(layer), sin, rowStart * rowBytes);
|
|
923
|
+
}
|
|
924
|
+
}
|
|
925
|
+
}
|
|
926
|
+
/**
|
|
927
|
+
* Accent/voice tags the loaded checkpoint supports. Multi-accent KaniTTS2
|
|
928
|
+
* checkpoints (kani-tts-2-en) return e.g. ["en_us","en_nyork","en_glasg", …];
|
|
929
|
+
* the plain LFM2 family (the Apache default) returns [] (one trained voice,
|
|
930
|
+
* selectable only by fine-tuning the backbone — there is no runtime voice/
|
|
931
|
+
* speaker-prompt knob, and zero-shot cloning is not wired).
|
|
932
|
+
*/
|
|
933
|
+
availableVoices() {
|
|
934
|
+
return [...this.voiceTags];
|
|
935
|
+
}
|
|
936
|
+
/** NanoCodec's fixed output sample rate (Hz). Not selectable — see SpeakResult. */
|
|
937
|
+
get sampleRate() {
|
|
938
|
+
return KANI_SAMPLE_RATE;
|
|
939
|
+
}
|
|
940
|
+
/**
|
|
941
|
+
* Synthesize speech for `text`. Returns 22 kHz mono PCM.
|
|
942
|
+
*
|
|
943
|
+
* Pipeline: build the [SOH]+text+[EOT,EOH] prompt → prefill the codec-LM →
|
|
944
|
+
* AR-decode 4-token frames (per-codebook constrained sampling) until end_of_speech
|
|
945
|
+
* → strip markers → codes → NanoCodec decode → PCM.
|
|
946
|
+
*/
|
|
947
|
+
async speak(text, opts = {}) {
|
|
948
|
+
if (this._destroyed) throw new Error("KaniTTS has been destroyed.");
|
|
949
|
+
const temperature = opts.temperature ?? 1;
|
|
950
|
+
const topP = opts.topP ?? .95;
|
|
951
|
+
const topK = opts.topK ?? 0;
|
|
952
|
+
const repetitionPenalty = opts.repetitionPenalty ?? 1.1;
|
|
953
|
+
const maxNewTokens = Math.min(opts.maxNewTokens ?? DEFAULT_MAX_NEW_TOKENS$1, this.maxSeqLen);
|
|
954
|
+
const promptText = this.voiceTags.length > 0 ? `${(opts.languageTag ?? this.voiceTags[0]).trim()}: ${text}` : text;
|
|
955
|
+
const prompt = [
|
|
956
|
+
KANI_START_OF_HUMAN,
|
|
957
|
+
...this.tokenizer.encode(promptText),
|
|
958
|
+
KANI_END_OF_TEXT,
|
|
959
|
+
KANI_END_OF_HUMAN
|
|
960
|
+
];
|
|
961
|
+
if (prompt.length >= this.maxSeqLen) throw new Error(`Kani prompt is ${prompt.length} tokens >= maxSeqLen ${this.maxSeqLen}.`);
|
|
962
|
+
this.backboneExec.reset();
|
|
963
|
+
const promptPos = computeKaniPositions(prompt, this.cfg);
|
|
964
|
+
this.writeCosSin(promptPos, 0);
|
|
965
|
+
const { logits } = await this.backboneExec.forward(new Uint32Array(prompt));
|
|
966
|
+
const startTime = performance.now();
|
|
967
|
+
const { audioTokens, finished } = await this.runDecodeLoop(prompt, logits, {
|
|
968
|
+
temperature,
|
|
969
|
+
topP,
|
|
970
|
+
topK,
|
|
971
|
+
repetitionPenalty,
|
|
972
|
+
maxNewTokens,
|
|
973
|
+
maxFrames: opts.maxFrames ?? Number.POSITIVE_INFINITY
|
|
974
|
+
});
|
|
975
|
+
const usable = audioTokens.length - audioTokens.length % this.cfg.tokensPerFrame;
|
|
976
|
+
if (usable === 0) throw new Error(`Kani produced no complete audio frames (finished=${finished}, audioTokens=${audioTokens.length}). The codec-LM emitted no speech.`);
|
|
977
|
+
const { codes, numFrames } = audioTokensToCodes(audioTokens.slice(0, usable), this.cfg);
|
|
978
|
+
const pcm = await this.decodeCodes(codes, numFrames);
|
|
979
|
+
const totalTime = (performance.now() - startTime) / 1e3;
|
|
980
|
+
const audioSeconds = numFrames * KANI_HOP / KANI_SAMPLE_RATE;
|
|
981
|
+
console.log(`[kani] frames=${numFrames} audio=${audioSeconds.toFixed(2)}s wall=${totalTime.toFixed(2)}s RTF=${(audioSeconds / totalTime).toFixed(2)}x`);
|
|
982
|
+
return {
|
|
983
|
+
pcm,
|
|
984
|
+
sampleRate: KANI_SAMPLE_RATE,
|
|
985
|
+
frames: numFrames,
|
|
986
|
+
audioSeconds
|
|
987
|
+
};
|
|
988
|
+
}
|
|
989
|
+
/**
|
|
990
|
+
* Autoregressive decode: from the prefill logits, emit one token per step —
|
|
991
|
+
* greedy for the structural markers ([SOA][SOS]) and per-codebook constrained
|
|
992
|
+
* sampling once in speech — collecting the audio tokens between SOS and EOS.
|
|
993
|
+
* Writes each step's per-layer cos/sin row (frame-level logical position) before
|
|
994
|
+
* the forward. Returns the collected audio tokens and whether EOS/cap was hit.
|
|
995
|
+
*/
|
|
996
|
+
async runDecodeLoop(prompt, prefillLogits, p) {
|
|
997
|
+
const seq = [...prompt];
|
|
998
|
+
const audioTokens = [];
|
|
999
|
+
const frameCap = p.maxFrames * this.cfg.tokensPerFrame;
|
|
1000
|
+
let logits = prefillLogits;
|
|
1001
|
+
let inSpeech = false;
|
|
1002
|
+
let finished = false;
|
|
1003
|
+
for (let step = 0; step < p.maxNewTokens; step++) {
|
|
1004
|
+
const nextToken = inSpeech ? this.sampleAudioToken(logits, audioTokens.length % this.cfg.tokensPerFrame, {
|
|
1005
|
+
temperature: p.temperature,
|
|
1006
|
+
topP: p.topP,
|
|
1007
|
+
topK: p.topK,
|
|
1008
|
+
repetitionPenalty: p.repetitionPenalty,
|
|
1009
|
+
previous: seq
|
|
1010
|
+
}) : this.argmax(logits);
|
|
1011
|
+
if (nextToken === this.cfg.startOfSpeech) inSpeech = true;
|
|
1012
|
+
else if (nextToken === this.cfg.endOfSpeech) {
|
|
1013
|
+
seq.push(nextToken);
|
|
1014
|
+
finished = true;
|
|
1015
|
+
break;
|
|
1016
|
+
} else if (inSpeech && nextToken >= this.cfg.audioTokensStart) audioTokens.push(nextToken);
|
|
1017
|
+
seq.push(nextToken);
|
|
1018
|
+
if (inSpeech && audioTokens.length >= frameCap) {
|
|
1019
|
+
finished = true;
|
|
1020
|
+
break;
|
|
1021
|
+
}
|
|
1022
|
+
if (seq.length >= this.maxSeqLen) break;
|
|
1023
|
+
this.writeCosSin(Float32Array.of(this.logicalPositionAt(seq)), this.backboneExec.currentSeqPos);
|
|
1024
|
+
logits = (await this.backboneExec.forward(new Uint32Array([nextToken]))).logits;
|
|
1025
|
+
}
|
|
1026
|
+
return {
|
|
1027
|
+
audioTokens,
|
|
1028
|
+
finished
|
|
1029
|
+
};
|
|
1030
|
+
}
|
|
1031
|
+
/** Logical (frame-level) position of the LAST token in `seq`. */
|
|
1032
|
+
logicalPositionAt(seq) {
|
|
1033
|
+
const positions = computeKaniPositions(seq, this.cfg);
|
|
1034
|
+
return positions[positions.length - 1];
|
|
1035
|
+
}
|
|
1036
|
+
/** Greedy argmax over a logits row. */
|
|
1037
|
+
argmax(logits) {
|
|
1038
|
+
let best = 0;
|
|
1039
|
+
let bestVal = logits[0];
|
|
1040
|
+
for (let i = 1; i < logits.length; i++) if (logits[i] > bestVal) {
|
|
1041
|
+
bestVal = logits[i];
|
|
1042
|
+
best = i;
|
|
1043
|
+
}
|
|
1044
|
+
return best;
|
|
1045
|
+
}
|
|
1046
|
+
/**
|
|
1047
|
+
* Sample one audio token for codebook position `codebook`, constrained to that
|
|
1048
|
+
* codebook's valid window [audio_tokens_start + 4032*c, +4032). Allows end_of_speech
|
|
1049
|
+
* only at codebook 0 (frame boundary). Applies repetition penalty, temperature,
|
|
1050
|
+
* optional top-k, then top-p nucleus sampling.
|
|
1051
|
+
*/
|
|
1052
|
+
sampleAudioToken(logits, codebook, p) {
|
|
1053
|
+
const winStart = this.cfg.audioTokensStart + this.cfg.codebookSize * codebook;
|
|
1054
|
+
const winLen = this.cfg.codebookSize;
|
|
1055
|
+
const allowEos = codebook === 0;
|
|
1056
|
+
const n = winLen + (allowEos ? 1 : 0);
|
|
1057
|
+
const ids = new Int32Array(n);
|
|
1058
|
+
for (let i = 0; i < winLen; i++) ids[i] = winStart + i;
|
|
1059
|
+
if (allowEos) ids[winLen] = this.cfg.endOfSpeech;
|
|
1060
|
+
const prev = p.repetitionPenalty !== 1 ? new Set(p.previous) : null;
|
|
1061
|
+
const scores = new Float32Array(n);
|
|
1062
|
+
for (let i = 0; i < n; i++) {
|
|
1063
|
+
let s = logits[ids[i]];
|
|
1064
|
+
if (prev?.has(ids[i])) s = s > 0 ? s / p.repetitionPenalty : s * p.repetitionPenalty;
|
|
1065
|
+
scores[i] = s / p.temperature;
|
|
1066
|
+
}
|
|
1067
|
+
softmaxInPlace$1(scores);
|
|
1068
|
+
applyTopK$1(scores, p.topK);
|
|
1069
|
+
return nucleusSample$1(scores, ids, p.topP);
|
|
1070
|
+
}
|
|
1071
|
+
/**
|
|
1072
|
+
* Decode NanoCodec codes [groups, T] (group-major) → PCM.
|
|
1073
|
+
*
|
|
1074
|
+
* The decoder graph carries concrete lengths, and the upsampled conv activations
|
|
1075
|
+
* for long clips overflow WebGPU's 65535 per-dimension dispatch cap. The decoder
|
|
1076
|
+
* is fully causal with a small (≤ a few frames) receptive field, so we decode in
|
|
1077
|
+
* frame chunks with a left-context lookback and keep only each chunk's own output
|
|
1078
|
+
* samples — numerically identical to a single decode, but bounded per dispatch.
|
|
1079
|
+
*/
|
|
1080
|
+
async decodeCodes(codes, numFrames) {
|
|
1081
|
+
const groups = 4;
|
|
1082
|
+
const lookback = KANI_CODEC_LOOKBACK_FRAMES;
|
|
1083
|
+
const chunk = KANI_CODEC_CHUNK_FRAMES;
|
|
1084
|
+
const pcm = new Float32Array(numFrames * KANI_HOP);
|
|
1085
|
+
for (let start = 0; start < numFrames; start += chunk) {
|
|
1086
|
+
const ctxStart = Math.max(0, start - lookback);
|
|
1087
|
+
const ctxFrames = start - ctxStart;
|
|
1088
|
+
const end = Math.min(numFrames, start + chunk);
|
|
1089
|
+
const winFrames = end - ctxStart;
|
|
1090
|
+
const winCodes = new Uint32Array(groups * winFrames);
|
|
1091
|
+
for (let g = 0; g < groups; g++) for (let t = 0; t < winFrames; t++) winCodes[g * winFrames + t] = codes[g * numFrames + (ctxStart + t)];
|
|
1092
|
+
const win = await this.decodeCodesWindow(winCodes, winFrames);
|
|
1093
|
+
const dropSamples = ctxFrames * KANI_HOP;
|
|
1094
|
+
const keepSamples = (end - start) * KANI_HOP;
|
|
1095
|
+
pcm.set(win.subarray(dropSamples, dropSamples + keepSamples), start * KANI_HOP);
|
|
1096
|
+
}
|
|
1097
|
+
for (let i = 0; i < pcm.length; i++) pcm[i] = Math.min(1, Math.max(-1, pcm[i]));
|
|
1098
|
+
return pcm;
|
|
1099
|
+
}
|
|
1100
|
+
/** Run the NanoCodec decoder graph for a single (bounded) code window → PCM. */
|
|
1101
|
+
async decodeCodesWindow(winCodes, winFrames) {
|
|
1102
|
+
const graph = generateNanoCodecDecoderGraph({ numFrames: winFrames });
|
|
1103
|
+
const exec = new Executor(this.ctx, graph, {
|
|
1104
|
+
maxSeqLen: winFrames,
|
|
1105
|
+
kvMode: "f32"
|
|
1106
|
+
});
|
|
1107
|
+
try {
|
|
1108
|
+
exec.uploadWeightsMap(selectGraphWeights$1(graph, this.loaded.codecWeights));
|
|
1109
|
+
exec.initBindGroups();
|
|
1110
|
+
exec.reset();
|
|
1111
|
+
exec.writeInput("audio_codes", winCodes);
|
|
1112
|
+
return await exec.runGraphOutput(graph.outputs[0], winFrames * KANI_HOP);
|
|
1113
|
+
} finally {
|
|
1114
|
+
exec.destroy();
|
|
1115
|
+
}
|
|
1116
|
+
}
|
|
1117
|
+
destroy() {
|
|
1118
|
+
if (this._destroyed) return;
|
|
1119
|
+
this._destroyed = true;
|
|
1120
|
+
this.backboneExec.destroy();
|
|
1121
|
+
this.loaded.backboneWeights.clear();
|
|
1122
|
+
this.loaded.codecWeights.clear();
|
|
1123
|
+
}
|
|
1124
|
+
};
|
|
1125
|
+
|
|
1126
|
+
//#endregion
|
|
1127
|
+
//#region src/gpu/architectures/outetts.ts
|
|
1128
|
+
const OUTETTS_C1_BASE = 151669;
|
|
1129
|
+
const OUTETTS_C2_BASE = 152694;
|
|
1130
|
+
const OUTETTS_CODES_PER_CODEBOOK = 1025;
|
|
1131
|
+
const OUTETTS_AUDIO_END = 156729;
|
|
1132
|
+
const OUTETTS_IM_END = 151645;
|
|
1133
|
+
const DAC_SAMPLE_RATE = 24e3;
|
|
1134
|
+
const DAC_HOP = 320;
|
|
1135
|
+
const DAC_NUM_CODEBOOKS = 2;
|
|
1136
|
+
const DAC_CODEBOOK_SIZE = 1024;
|
|
1137
|
+
const DAC_CODEBOOK_DIM = 8;
|
|
1138
|
+
const DAC_LATENT_DIM = 1024;
|
|
1139
|
+
const DAC_DECODER_DIM = 1536;
|
|
1140
|
+
const DAC_DECODER_RATES = [
|
|
1141
|
+
8,
|
|
1142
|
+
5,
|
|
1143
|
+
4,
|
|
1144
|
+
2
|
|
1145
|
+
];
|
|
1146
|
+
const DAC_RESBLOCK_DILATIONS = [
|
|
1147
|
+
1,
|
|
1148
|
+
3,
|
|
1149
|
+
9
|
|
1150
|
+
];
|
|
1151
|
+
const DAC_ENCODER_DIM = 64;
|
|
1152
|
+
const DAC_ENCODER_RATES = [
|
|
1153
|
+
2,
|
|
1154
|
+
4,
|
|
1155
|
+
5,
|
|
1156
|
+
8
|
|
1157
|
+
];
|
|
1158
|
+
function parseOuteTtsConfig(rawConfig) {
|
|
1159
|
+
const hidden_size = rawConfig.hidden_size;
|
|
1160
|
+
const num_heads = rawConfig.num_attention_heads;
|
|
1161
|
+
return {
|
|
1162
|
+
hidden_size,
|
|
1163
|
+
num_layers: rawConfig.num_hidden_layers,
|
|
1164
|
+
num_heads,
|
|
1165
|
+
num_kv_heads: rawConfig.num_key_value_heads ?? num_heads,
|
|
1166
|
+
head_dim: rawConfig.head_dim ?? Math.floor(hidden_size / num_heads),
|
|
1167
|
+
intermediate_size: rawConfig.intermediate_size,
|
|
1168
|
+
vocab_size: rawConfig.vocab_size,
|
|
1169
|
+
rope_base: rawConfig.rope_theta ?? 1e6,
|
|
1170
|
+
rms_norm_eps: rawConfig.rms_norm_eps ?? 1e-6,
|
|
1171
|
+
max_position_embeddings: rawConfig.max_position_embeddings ?? 40960
|
|
1172
|
+
};
|
|
1173
|
+
}
|
|
1174
|
+
/**
|
|
1175
|
+
* Convert the flat generated token stream (the run BETWEEN <|audio_start|> and
|
|
1176
|
+
* <|audio_end|>) into a DAC code grid [NUM_CODEBOOKS, T] (codebook-major [c*T + t]).
|
|
1177
|
+
*
|
|
1178
|
+
* Mirrors outetts v3 `extract_audio_from_tokens`: filter the stream into two ordered
|
|
1179
|
+
* lists by membership of the c1 / c2 ranges, recover the raw code
|
|
1180
|
+
* (id − base), drop the reserved index 1024, and truncate to the shorter list. The
|
|
1181
|
+
* model emits c1,c2 interleaved per frame, but the parse is membership-based (it does
|
|
1182
|
+
* not assume strict pairing) so a stray token cannot desync the grid.
|
|
1183
|
+
*/
|
|
1184
|
+
function audioTokensToDacCodes(tokenIds) {
|
|
1185
|
+
const c1 = [];
|
|
1186
|
+
const c2 = [];
|
|
1187
|
+
for (const id of tokenIds) if (id >= OUTETTS_C1_BASE && id < OUTETTS_C1_BASE + OUTETTS_CODES_PER_CODEBOOK) {
|
|
1188
|
+
const code = id - OUTETTS_C1_BASE;
|
|
1189
|
+
if (code < DAC_CODEBOOK_SIZE) c1.push(code);
|
|
1190
|
+
} else if (id >= OUTETTS_C2_BASE && id < OUTETTS_C2_BASE + OUTETTS_CODES_PER_CODEBOOK) {
|
|
1191
|
+
const code = id - OUTETTS_C2_BASE;
|
|
1192
|
+
if (code < DAC_CODEBOOK_SIZE) c2.push(code);
|
|
1193
|
+
}
|
|
1194
|
+
const numFrames = Math.min(c1.length, c2.length);
|
|
1195
|
+
const codes = new Uint32Array(DAC_NUM_CODEBOOKS * numFrames);
|
|
1196
|
+
for (let t = 0; t < numFrames; t++) {
|
|
1197
|
+
codes[t] = c1[t];
|
|
1198
|
+
codes[numFrames + t] = c2[t];
|
|
1199
|
+
}
|
|
1200
|
+
return {
|
|
1201
|
+
codes,
|
|
1202
|
+
numFrames
|
|
1203
|
+
};
|
|
1204
|
+
}
|
|
1205
|
+
/** ConvTranspose1d output length: (L-1)*stride - 2*pad + (K-1) + output_padding + 1. */
|
|
1206
|
+
function convTransposeLen(L, stride, K, pad, outPad) {
|
|
1207
|
+
return (L - 1) * stride - 2 * pad + (K - 1) + outPad + 1;
|
|
1208
|
+
}
|
|
1209
|
+
/** Canonical DAC decoder weight name. */
|
|
1210
|
+
function dacKey(name) {
|
|
1211
|
+
return `dac.${name}`;
|
|
1212
|
+
}
|
|
1213
|
+
function generateDacSpeechDecoderGraph(opts) {
|
|
1214
|
+
const T = opts.numFrames;
|
|
1215
|
+
const tensors = {};
|
|
1216
|
+
const nodes = [];
|
|
1217
|
+
const executionOrder = [];
|
|
1218
|
+
const addTensor = (desc) => {
|
|
1219
|
+
tensors[desc.name] = desc;
|
|
1220
|
+
};
|
|
1221
|
+
const addNode = (node) => {
|
|
1222
|
+
nodes.push(node);
|
|
1223
|
+
executionOrder.push(node.id);
|
|
1224
|
+
};
|
|
1225
|
+
addTensor({
|
|
1226
|
+
name: "audio_codes",
|
|
1227
|
+
shape: [DAC_NUM_CODEBOOKS, T],
|
|
1228
|
+
dtype: "u32",
|
|
1229
|
+
storage: "activation"
|
|
1230
|
+
});
|
|
1231
|
+
addTensor({
|
|
1232
|
+
name: dacKey("quantizer.codebook"),
|
|
1233
|
+
shape: [
|
|
1234
|
+
DAC_NUM_CODEBOOKS,
|
|
1235
|
+
DAC_CODEBOOK_SIZE,
|
|
1236
|
+
DAC_CODEBOOK_DIM
|
|
1237
|
+
],
|
|
1238
|
+
dtype: "f32",
|
|
1239
|
+
storage: "constant",
|
|
1240
|
+
safetensorsKey: dacKey("quantizer.codebook")
|
|
1241
|
+
});
|
|
1242
|
+
addTensor({
|
|
1243
|
+
name: dacKey("quantizer.out_proj.weight"),
|
|
1244
|
+
shape: [
|
|
1245
|
+
DAC_NUM_CODEBOOKS,
|
|
1246
|
+
DAC_LATENT_DIM,
|
|
1247
|
+
DAC_CODEBOOK_DIM
|
|
1248
|
+
],
|
|
1249
|
+
dtype: "f32",
|
|
1250
|
+
storage: "constant",
|
|
1251
|
+
safetensorsKey: dacKey("quantizer.out_proj.weight")
|
|
1252
|
+
});
|
|
1253
|
+
addTensor({
|
|
1254
|
+
name: dacKey("quantizer.out_proj.bias"),
|
|
1255
|
+
shape: [DAC_NUM_CODEBOOKS, DAC_LATENT_DIM],
|
|
1256
|
+
dtype: "f32",
|
|
1257
|
+
storage: "constant",
|
|
1258
|
+
safetensorsKey: dacKey("quantizer.out_proj.bias")
|
|
1259
|
+
});
|
|
1260
|
+
addTensor({
|
|
1261
|
+
name: "dac_latent",
|
|
1262
|
+
shape: [DAC_LATENT_DIM, T],
|
|
1263
|
+
dtype: "f32",
|
|
1264
|
+
storage: "activation"
|
|
1265
|
+
});
|
|
1266
|
+
addNode({
|
|
1267
|
+
id: "dac_rvq",
|
|
1268
|
+
opType: "RVQDequant",
|
|
1269
|
+
inputs: [
|
|
1270
|
+
"audio_codes",
|
|
1271
|
+
dacKey("quantizer.codebook"),
|
|
1272
|
+
dacKey("quantizer.out_proj.weight"),
|
|
1273
|
+
dacKey("quantizer.out_proj.bias")
|
|
1274
|
+
],
|
|
1275
|
+
outputs: ["dac_latent"],
|
|
1276
|
+
attributes: {
|
|
1277
|
+
NCB: DAC_NUM_CODEBOOKS,
|
|
1278
|
+
Cout: DAC_LATENT_DIM,
|
|
1279
|
+
T,
|
|
1280
|
+
D: DAC_CODEBOOK_DIM,
|
|
1281
|
+
codebook_size: DAC_CODEBOOK_SIZE
|
|
1282
|
+
}
|
|
1283
|
+
});
|
|
1284
|
+
/** Emit a full (non-causal) Conv1dFull with symmetric padding. */
|
|
1285
|
+
function conv1dFull(id, input, weightKey, biasKey, Cin, Cout, K, L, dilation, padding, output) {
|
|
1286
|
+
const Lout = Math.floor((L + 2 * padding - dilation * (K - 1) - 1) / 1) + 1;
|
|
1287
|
+
addTensor({
|
|
1288
|
+
name: weightKey,
|
|
1289
|
+
shape: [
|
|
1290
|
+
Cout,
|
|
1291
|
+
Cin,
|
|
1292
|
+
K
|
|
1293
|
+
],
|
|
1294
|
+
dtype: "f32",
|
|
1295
|
+
storage: "constant",
|
|
1296
|
+
safetensorsKey: weightKey
|
|
1297
|
+
});
|
|
1298
|
+
addTensor({
|
|
1299
|
+
name: biasKey,
|
|
1300
|
+
shape: [Cout],
|
|
1301
|
+
dtype: "f32",
|
|
1302
|
+
storage: "constant",
|
|
1303
|
+
safetensorsKey: biasKey
|
|
1304
|
+
});
|
|
1305
|
+
addTensor({
|
|
1306
|
+
name: output,
|
|
1307
|
+
shape: [Cout, Lout],
|
|
1308
|
+
dtype: "f32",
|
|
1309
|
+
storage: "activation"
|
|
1310
|
+
});
|
|
1311
|
+
addNode({
|
|
1312
|
+
id,
|
|
1313
|
+
opType: "Conv1dFull",
|
|
1314
|
+
inputs: [
|
|
1315
|
+
input,
|
|
1316
|
+
weightKey,
|
|
1317
|
+
biasKey
|
|
1318
|
+
],
|
|
1319
|
+
outputs: [output],
|
|
1320
|
+
attributes: {
|
|
1321
|
+
Cin,
|
|
1322
|
+
Cout,
|
|
1323
|
+
L,
|
|
1324
|
+
Lout,
|
|
1325
|
+
K,
|
|
1326
|
+
stride: 1,
|
|
1327
|
+
padding,
|
|
1328
|
+
dilation
|
|
1329
|
+
}
|
|
1330
|
+
});
|
|
1331
|
+
return Lout;
|
|
1332
|
+
}
|
|
1333
|
+
/** Emit a full-channel Snake1d (alpha is [C], stored [1,C,1] → C). */
|
|
1334
|
+
function snake(id, input, C, L, alphaKey, output) {
|
|
1335
|
+
addTensor({
|
|
1336
|
+
name: alphaKey,
|
|
1337
|
+
shape: [C],
|
|
1338
|
+
dtype: "f32",
|
|
1339
|
+
storage: "constant",
|
|
1340
|
+
safetensorsKey: alphaKey
|
|
1341
|
+
});
|
|
1342
|
+
addTensor({
|
|
1343
|
+
name: output,
|
|
1344
|
+
shape: [C, L],
|
|
1345
|
+
dtype: "f32",
|
|
1346
|
+
storage: "activation"
|
|
1347
|
+
});
|
|
1348
|
+
addNode({
|
|
1349
|
+
id,
|
|
1350
|
+
opType: "Snake1d",
|
|
1351
|
+
inputs: [input, alphaKey],
|
|
1352
|
+
outputs: [output],
|
|
1353
|
+
attributes: {
|
|
1354
|
+
C,
|
|
1355
|
+
L
|
|
1356
|
+
}
|
|
1357
|
+
});
|
|
1358
|
+
}
|
|
1359
|
+
/** Emit a strided ConvTranspose1d (weight [Cin, Cout, K]). */
|
|
1360
|
+
function convTranspose(id, input, weightKey, biasKey, Cin, Cout, K, L, stride, padding, outPad, output) {
|
|
1361
|
+
const Lout = convTransposeLen(L, stride, K, padding, outPad);
|
|
1362
|
+
addTensor({
|
|
1363
|
+
name: weightKey,
|
|
1364
|
+
shape: [
|
|
1365
|
+
Cin,
|
|
1366
|
+
Cout,
|
|
1367
|
+
K
|
|
1368
|
+
],
|
|
1369
|
+
dtype: "f32",
|
|
1370
|
+
storage: "constant",
|
|
1371
|
+
safetensorsKey: weightKey
|
|
1372
|
+
});
|
|
1373
|
+
addTensor({
|
|
1374
|
+
name: biasKey,
|
|
1375
|
+
shape: [Cout],
|
|
1376
|
+
dtype: "f32",
|
|
1377
|
+
storage: "constant",
|
|
1378
|
+
safetensorsKey: biasKey
|
|
1379
|
+
});
|
|
1380
|
+
addTensor({
|
|
1381
|
+
name: output,
|
|
1382
|
+
shape: [Cout, Lout],
|
|
1383
|
+
dtype: "f32",
|
|
1384
|
+
storage: "activation"
|
|
1385
|
+
});
|
|
1386
|
+
addNode({
|
|
1387
|
+
id,
|
|
1388
|
+
opType: "ConvTranspose1d",
|
|
1389
|
+
inputs: [
|
|
1390
|
+
input,
|
|
1391
|
+
weightKey,
|
|
1392
|
+
biasKey
|
|
1393
|
+
],
|
|
1394
|
+
outputs: [output],
|
|
1395
|
+
attributes: {
|
|
1396
|
+
Cin,
|
|
1397
|
+
Cout,
|
|
1398
|
+
L,
|
|
1399
|
+
Lout,
|
|
1400
|
+
K,
|
|
1401
|
+
stride,
|
|
1402
|
+
padding,
|
|
1403
|
+
output_padding: outPad
|
|
1404
|
+
}
|
|
1405
|
+
});
|
|
1406
|
+
return Lout;
|
|
1407
|
+
}
|
|
1408
|
+
/** Element-wise residual add of two equal-shape [C, L] activations. */
|
|
1409
|
+
function add(id, a, b, C, L, output) {
|
|
1410
|
+
addTensor({
|
|
1411
|
+
name: output,
|
|
1412
|
+
shape: [C, L],
|
|
1413
|
+
dtype: "f32",
|
|
1414
|
+
storage: "activation"
|
|
1415
|
+
});
|
|
1416
|
+
addNode({
|
|
1417
|
+
id,
|
|
1418
|
+
opType: "Add",
|
|
1419
|
+
inputs: [a, b],
|
|
1420
|
+
outputs: [output],
|
|
1421
|
+
attributes: {
|
|
1422
|
+
count_tensor: a,
|
|
1423
|
+
hidden_size: C * L
|
|
1424
|
+
}
|
|
1425
|
+
});
|
|
1426
|
+
}
|
|
1427
|
+
let cur = "dac_conv1_out";
|
|
1428
|
+
let length = conv1dFull("dac_conv1", "dac_latent", dacKey("decoder.model.0.weight"), dacKey("decoder.model.0.bias"), DAC_LATENT_DIM, DAC_DECODER_DIM, 7, T, 1, 3, cur);
|
|
1429
|
+
let channels = DAC_DECODER_DIM;
|
|
1430
|
+
for (let stage = 0; stage < DAC_DECODER_RATES.length; stage++) {
|
|
1431
|
+
const m = stage + 1;
|
|
1432
|
+
const rate = DAC_DECODER_RATES[stage];
|
|
1433
|
+
const outCh = channels / 2;
|
|
1434
|
+
const upK = 2 * rate;
|
|
1435
|
+
const upPad = Math.ceil(rate / 2);
|
|
1436
|
+
const snakeOut = `dac_s${stage}_snake`;
|
|
1437
|
+
snake(`dac_s${stage}_snake_op`, cur, channels, length, dacKey(`decoder.model.${m}.block.0.alpha`), snakeOut);
|
|
1438
|
+
const upOut = `dac_s${stage}_up`;
|
|
1439
|
+
length = convTranspose(`dac_s${stage}_convt`, snakeOut, dacKey(`decoder.model.${m}.block.1.weight`), dacKey(`decoder.model.${m}.block.1.bias`), channels, outCh, upK, length, rate, upPad, 0, upOut);
|
|
1440
|
+
channels = outCh;
|
|
1441
|
+
let h = upOut;
|
|
1442
|
+
for (let r = 0; r < DAC_RESBLOCK_DILATIONS.length; r++) {
|
|
1443
|
+
const ru = r + 2;
|
|
1444
|
+
const dil = DAC_RESBLOCK_DILATIONS[r];
|
|
1445
|
+
const p = `decoder.model.${m}.block.${ru}.block`;
|
|
1446
|
+
const pad = Math.floor(6 * dil / 2);
|
|
1447
|
+
const ia = `dac_s${stage}_r${r}_ia`;
|
|
1448
|
+
snake(`${ia}_op`, h, channels, length, dacKey(`${p}.0.alpha`), ia);
|
|
1449
|
+
const ic = `dac_s${stage}_r${r}_ic`;
|
|
1450
|
+
conv1dFull(`${ic}_op`, ia, dacKey(`${p}.1.weight`), dacKey(`${p}.1.bias`), channels, channels, 7, length, dil, pad, ic);
|
|
1451
|
+
const sa = `dac_s${stage}_r${r}_sa`;
|
|
1452
|
+
snake(`${sa}_op`, ic, channels, length, dacKey(`${p}.2.alpha`), sa);
|
|
1453
|
+
const sc = `dac_s${stage}_r${r}_sc`;
|
|
1454
|
+
conv1dFull(`${sc}_op`, sa, dacKey(`${p}.3.weight`), dacKey(`${p}.3.bias`), channels, channels, 1, length, 1, 0, sc);
|
|
1455
|
+
const resAdd = `dac_s${stage}_r${r}_res`;
|
|
1456
|
+
add(`${resAdd}_op`, h, sc, channels, length, resAdd);
|
|
1457
|
+
h = resAdd;
|
|
1458
|
+
}
|
|
1459
|
+
cur = h;
|
|
1460
|
+
}
|
|
1461
|
+
const postSnake = "dac_post_snake";
|
|
1462
|
+
snake("dac_post_snake_op", cur, channels, length, dacKey("decoder.model.5.alpha"), postSnake);
|
|
1463
|
+
const pcm = "dac_pcm";
|
|
1464
|
+
const pcmLen = conv1dFull("dac_conv2", postSnake, dacKey("decoder.model.6.weight"), dacKey("decoder.model.6.bias"), channels, 1, 7, length, 1, 3, pcm);
|
|
1465
|
+
return {
|
|
1466
|
+
architecture: "DacSpeechDecoder",
|
|
1467
|
+
config: {
|
|
1468
|
+
hidden_size: DAC_DECODER_DIM,
|
|
1469
|
+
num_layers: DAC_DECODER_RATES.length,
|
|
1470
|
+
num_heads: 0,
|
|
1471
|
+
num_kv_heads: 0,
|
|
1472
|
+
head_dim: 0,
|
|
1473
|
+
intermediate_size: 0,
|
|
1474
|
+
vocab_size: 0,
|
|
1475
|
+
context_length: pcmLen,
|
|
1476
|
+
rms_norm_eps: 0,
|
|
1477
|
+
norm_type: "rmsnorm",
|
|
1478
|
+
rope_base: 0,
|
|
1479
|
+
rope_dim: 0,
|
|
1480
|
+
kv_layout: "LHSd",
|
|
1481
|
+
is_moe: false,
|
|
1482
|
+
has_vision_tower: false
|
|
1483
|
+
},
|
|
1484
|
+
capabilities: {
|
|
1485
|
+
text: true,
|
|
1486
|
+
vision: false,
|
|
1487
|
+
moe: false
|
|
1488
|
+
},
|
|
1489
|
+
tensors,
|
|
1490
|
+
nodes,
|
|
1491
|
+
executionOrder,
|
|
1492
|
+
inputs: ["audio_codes"],
|
|
1493
|
+
outputs: [pcm]
|
|
1494
|
+
};
|
|
1495
|
+
}
|
|
1496
|
+
/** Final decoded PCM length (samples) for a T-frame code grid, via the DAC stage chain. */
|
|
1497
|
+
function dacOutputLength(numFrames) {
|
|
1498
|
+
let length = numFrames;
|
|
1499
|
+
for (const rate of DAC_DECODER_RATES) length = convTransposeLen(length, rate, 2 * rate, Math.ceil(rate / 2), 0);
|
|
1500
|
+
return length;
|
|
1501
|
+
}
|
|
1502
|
+
/** Strided/dilated Conv1d output length: floor((L + 2*pad - dilation*(K-1) - 1)/stride) + 1. */
|
|
1503
|
+
function conv1dLen(L, K, stride, pad, dilation) {
|
|
1504
|
+
return Math.floor((L + 2 * pad - dilation * (K - 1) - 1) / stride) + 1;
|
|
1505
|
+
}
|
|
1506
|
+
/** Number of latent frames the encoder produces for `numSamples` of padded 24 kHz PCM. */
|
|
1507
|
+
function dacEncoderFrames(numSamples) {
|
|
1508
|
+
let length = numSamples;
|
|
1509
|
+
for (const rate of DAC_ENCODER_RATES) length = conv1dLen(length, 2 * rate, rate, Math.ceil(rate / 2), 1);
|
|
1510
|
+
return length;
|
|
1511
|
+
}
|
|
1512
|
+
function generateDacSpeechEncoderGraph(opts) {
|
|
1513
|
+
const L0 = opts.numSamples;
|
|
1514
|
+
const tensors = {};
|
|
1515
|
+
const nodes = [];
|
|
1516
|
+
const executionOrder = [];
|
|
1517
|
+
const addTensor = (desc) => {
|
|
1518
|
+
tensors[desc.name] = desc;
|
|
1519
|
+
};
|
|
1520
|
+
const addNode = (node) => {
|
|
1521
|
+
nodes.push(node);
|
|
1522
|
+
executionOrder.push(node.id);
|
|
1523
|
+
};
|
|
1524
|
+
addTensor({
|
|
1525
|
+
name: "pcm_in",
|
|
1526
|
+
shape: [1, L0],
|
|
1527
|
+
dtype: "f32",
|
|
1528
|
+
storage: "activation"
|
|
1529
|
+
});
|
|
1530
|
+
/** Emit a Conv1dFull (symmetric padding) and return its output length. */
|
|
1531
|
+
function conv1dFull(id, input, weightKey, biasKey, Cin, Cout, K, L, stride, dilation, padding, output) {
|
|
1532
|
+
const Lout = conv1dLen(L, K, stride, padding, dilation);
|
|
1533
|
+
addTensor({
|
|
1534
|
+
name: weightKey,
|
|
1535
|
+
shape: [
|
|
1536
|
+
Cout,
|
|
1537
|
+
Cin,
|
|
1538
|
+
K
|
|
1539
|
+
],
|
|
1540
|
+
dtype: "f32",
|
|
1541
|
+
storage: "constant",
|
|
1542
|
+
safetensorsKey: weightKey
|
|
1543
|
+
});
|
|
1544
|
+
addTensor({
|
|
1545
|
+
name: biasKey,
|
|
1546
|
+
shape: [Cout],
|
|
1547
|
+
dtype: "f32",
|
|
1548
|
+
storage: "constant",
|
|
1549
|
+
safetensorsKey: biasKey
|
|
1550
|
+
});
|
|
1551
|
+
addTensor({
|
|
1552
|
+
name: output,
|
|
1553
|
+
shape: [Cout, Lout],
|
|
1554
|
+
dtype: "f32",
|
|
1555
|
+
storage: "activation"
|
|
1556
|
+
});
|
|
1557
|
+
addNode({
|
|
1558
|
+
id,
|
|
1559
|
+
opType: "Conv1dFull",
|
|
1560
|
+
inputs: [
|
|
1561
|
+
input,
|
|
1562
|
+
weightKey,
|
|
1563
|
+
biasKey
|
|
1564
|
+
],
|
|
1565
|
+
outputs: [output],
|
|
1566
|
+
attributes: {
|
|
1567
|
+
Cin,
|
|
1568
|
+
Cout,
|
|
1569
|
+
L,
|
|
1570
|
+
Lout,
|
|
1571
|
+
K,
|
|
1572
|
+
stride,
|
|
1573
|
+
padding,
|
|
1574
|
+
dilation
|
|
1575
|
+
}
|
|
1576
|
+
});
|
|
1577
|
+
return Lout;
|
|
1578
|
+
}
|
|
1579
|
+
/** Emit a full-channel Snake1d (alpha is [C], stored [1,C,1] → C). */
|
|
1580
|
+
function snake(id, input, C, L, alphaKey, output) {
|
|
1581
|
+
addTensor({
|
|
1582
|
+
name: alphaKey,
|
|
1583
|
+
shape: [C],
|
|
1584
|
+
dtype: "f32",
|
|
1585
|
+
storage: "constant",
|
|
1586
|
+
safetensorsKey: alphaKey
|
|
1587
|
+
});
|
|
1588
|
+
addTensor({
|
|
1589
|
+
name: output,
|
|
1590
|
+
shape: [C, L],
|
|
1591
|
+
dtype: "f32",
|
|
1592
|
+
storage: "activation"
|
|
1593
|
+
});
|
|
1594
|
+
addNode({
|
|
1595
|
+
id,
|
|
1596
|
+
opType: "Snake1d",
|
|
1597
|
+
inputs: [input, alphaKey],
|
|
1598
|
+
outputs: [output],
|
|
1599
|
+
attributes: {
|
|
1600
|
+
C,
|
|
1601
|
+
L
|
|
1602
|
+
}
|
|
1603
|
+
});
|
|
1604
|
+
}
|
|
1605
|
+
/** Element-wise residual add of two equal-shape [C, L] activations. */
|
|
1606
|
+
function add(id, a, b, C, L, output) {
|
|
1607
|
+
addTensor({
|
|
1608
|
+
name: output,
|
|
1609
|
+
shape: [C, L],
|
|
1610
|
+
dtype: "f32",
|
|
1611
|
+
storage: "activation"
|
|
1612
|
+
});
|
|
1613
|
+
addNode({
|
|
1614
|
+
id,
|
|
1615
|
+
opType: "Add",
|
|
1616
|
+
inputs: [a, b],
|
|
1617
|
+
outputs: [output],
|
|
1618
|
+
attributes: {
|
|
1619
|
+
count_tensor: a,
|
|
1620
|
+
hidden_size: C * L
|
|
1621
|
+
}
|
|
1622
|
+
});
|
|
1623
|
+
}
|
|
1624
|
+
const encKey = (name) => dacKey(`encoder.${name}`);
|
|
1625
|
+
let cur = "dac_enc_conv0";
|
|
1626
|
+
let length = conv1dFull("dac_enc_conv0_op", "pcm_in", encKey("block.0.weight"), encKey("block.0.bias"), 1, DAC_ENCODER_DIM, 7, L0, 1, 1, 3, cur);
|
|
1627
|
+
let channels = DAC_ENCODER_DIM;
|
|
1628
|
+
for (let stage = 0; stage < DAC_ENCODER_RATES.length; stage++) {
|
|
1629
|
+
const m = stage + 1;
|
|
1630
|
+
const rate = DAC_ENCODER_RATES[stage];
|
|
1631
|
+
let h = cur;
|
|
1632
|
+
for (let r = 0; r < DAC_RESBLOCK_DILATIONS.length; r++) {
|
|
1633
|
+
const dil = DAC_RESBLOCK_DILATIONS[r];
|
|
1634
|
+
const p = `block.${m}.block.${r}.block`;
|
|
1635
|
+
const pad = Math.floor(6 * dil / 2);
|
|
1636
|
+
const ia = `dac_enc_s${stage}_r${r}_ia`;
|
|
1637
|
+
snake(`${ia}_op`, h, channels, length, encKey(`${p}.0.alpha`), ia);
|
|
1638
|
+
const ic = `dac_enc_s${stage}_r${r}_ic`;
|
|
1639
|
+
conv1dFull(`${ic}_op`, ia, encKey(`${p}.1.weight`), encKey(`${p}.1.bias`), channels, channels, 7, length, 1, dil, pad, ic);
|
|
1640
|
+
const sa = `dac_enc_s${stage}_r${r}_sa`;
|
|
1641
|
+
snake(`${sa}_op`, ic, channels, length, encKey(`${p}.2.alpha`), sa);
|
|
1642
|
+
const sc = `dac_enc_s${stage}_r${r}_sc`;
|
|
1643
|
+
conv1dFull(`${sc}_op`, sa, encKey(`${p}.3.weight`), encKey(`${p}.3.bias`), channels, channels, 1, length, 1, 1, 0, sc);
|
|
1644
|
+
const resAdd = `dac_enc_s${stage}_r${r}_res`;
|
|
1645
|
+
add(`${resAdd}_op`, h, sc, channels, length, resAdd);
|
|
1646
|
+
h = resAdd;
|
|
1647
|
+
}
|
|
1648
|
+
const snakeOut = `dac_enc_s${stage}_snake`;
|
|
1649
|
+
snake(`${snakeOut}_op`, h, channels, length, encKey(`block.${m}.block.3.alpha`), snakeOut);
|
|
1650
|
+
const outCh = channels * 2;
|
|
1651
|
+
const downOut = `dac_enc_s${stage}_down`;
|
|
1652
|
+
length = conv1dFull(`${downOut}_op`, snakeOut, encKey(`block.${m}.block.4.weight`), encKey(`block.${m}.block.4.bias`), channels, outCh, 2 * rate, length, rate, 1, Math.ceil(rate / 2), downOut);
|
|
1653
|
+
channels = outCh;
|
|
1654
|
+
cur = downOut;
|
|
1655
|
+
}
|
|
1656
|
+
const postSnake = "dac_enc_post_snake";
|
|
1657
|
+
snake("dac_enc_post_snake_op", cur, channels, length, encKey("block.5.alpha"), postSnake);
|
|
1658
|
+
const latent = "dac_enc_latent";
|
|
1659
|
+
const T = conv1dFull("dac_enc_conv6_op", postSnake, encKey("block.6.weight"), encKey("block.6.bias"), channels, DAC_LATENT_DIM, 3, length, 1, 1, 1, latent);
|
|
1660
|
+
addTensor({
|
|
1661
|
+
name: dacKey("quantizer.in_proj.weight"),
|
|
1662
|
+
shape: [
|
|
1663
|
+
DAC_NUM_CODEBOOKS,
|
|
1664
|
+
DAC_CODEBOOK_DIM,
|
|
1665
|
+
DAC_LATENT_DIM
|
|
1666
|
+
],
|
|
1667
|
+
dtype: "f32",
|
|
1668
|
+
storage: "constant",
|
|
1669
|
+
safetensorsKey: dacKey("quantizer.in_proj.weight")
|
|
1670
|
+
});
|
|
1671
|
+
addTensor({
|
|
1672
|
+
name: dacKey("quantizer.in_proj.bias"),
|
|
1673
|
+
shape: [DAC_NUM_CODEBOOKS, DAC_CODEBOOK_DIM],
|
|
1674
|
+
dtype: "f32",
|
|
1675
|
+
storage: "constant",
|
|
1676
|
+
safetensorsKey: dacKey("quantizer.in_proj.bias")
|
|
1677
|
+
});
|
|
1678
|
+
addTensor({
|
|
1679
|
+
name: dacKey("quantizer.codebook"),
|
|
1680
|
+
shape: [
|
|
1681
|
+
DAC_NUM_CODEBOOKS,
|
|
1682
|
+
DAC_CODEBOOK_SIZE,
|
|
1683
|
+
DAC_CODEBOOK_DIM
|
|
1684
|
+
],
|
|
1685
|
+
dtype: "f32",
|
|
1686
|
+
storage: "constant",
|
|
1687
|
+
safetensorsKey: dacKey("quantizer.codebook")
|
|
1688
|
+
});
|
|
1689
|
+
addTensor({
|
|
1690
|
+
name: dacKey("quantizer.out_proj.weight"),
|
|
1691
|
+
shape: [
|
|
1692
|
+
DAC_NUM_CODEBOOKS,
|
|
1693
|
+
DAC_LATENT_DIM,
|
|
1694
|
+
DAC_CODEBOOK_DIM
|
|
1695
|
+
],
|
|
1696
|
+
dtype: "f32",
|
|
1697
|
+
storage: "constant",
|
|
1698
|
+
safetensorsKey: dacKey("quantizer.out_proj.weight")
|
|
1699
|
+
});
|
|
1700
|
+
addTensor({
|
|
1701
|
+
name: dacKey("quantizer.out_proj.bias"),
|
|
1702
|
+
shape: [DAC_NUM_CODEBOOKS, DAC_LATENT_DIM],
|
|
1703
|
+
dtype: "f32",
|
|
1704
|
+
storage: "constant",
|
|
1705
|
+
safetensorsKey: dacKey("quantizer.out_proj.bias")
|
|
1706
|
+
});
|
|
1707
|
+
addTensor({
|
|
1708
|
+
name: "dac_codes",
|
|
1709
|
+
shape: [DAC_NUM_CODEBOOKS, T],
|
|
1710
|
+
dtype: "u32",
|
|
1711
|
+
storage: "activation"
|
|
1712
|
+
});
|
|
1713
|
+
addNode({
|
|
1714
|
+
id: "dac_rvq_encode",
|
|
1715
|
+
opType: "RVQEncode",
|
|
1716
|
+
inputs: [
|
|
1717
|
+
latent,
|
|
1718
|
+
dacKey("quantizer.in_proj.weight"),
|
|
1719
|
+
dacKey("quantizer.in_proj.bias"),
|
|
1720
|
+
dacKey("quantizer.codebook"),
|
|
1721
|
+
dacKey("quantizer.out_proj.weight"),
|
|
1722
|
+
dacKey("quantizer.out_proj.bias")
|
|
1723
|
+
],
|
|
1724
|
+
outputs: ["dac_codes"],
|
|
1725
|
+
attributes: {
|
|
1726
|
+
NCB: DAC_NUM_CODEBOOKS,
|
|
1727
|
+
Cout: DAC_LATENT_DIM,
|
|
1728
|
+
T,
|
|
1729
|
+
D: DAC_CODEBOOK_DIM,
|
|
1730
|
+
codebook_size: DAC_CODEBOOK_SIZE
|
|
1731
|
+
}
|
|
1732
|
+
});
|
|
1733
|
+
return {
|
|
1734
|
+
architecture: "DacSpeechEncoder",
|
|
1735
|
+
config: {
|
|
1736
|
+
hidden_size: DAC_LATENT_DIM,
|
|
1737
|
+
num_layers: DAC_ENCODER_RATES.length,
|
|
1738
|
+
num_heads: 0,
|
|
1739
|
+
num_kv_heads: 0,
|
|
1740
|
+
head_dim: 0,
|
|
1741
|
+
intermediate_size: 0,
|
|
1742
|
+
vocab_size: 0,
|
|
1743
|
+
context_length: T,
|
|
1744
|
+
rms_norm_eps: 0,
|
|
1745
|
+
norm_type: "rmsnorm",
|
|
1746
|
+
rope_base: 0,
|
|
1747
|
+
rope_dim: 0,
|
|
1748
|
+
kv_layout: "LHSd",
|
|
1749
|
+
is_moe: false,
|
|
1750
|
+
has_vision_tower: false
|
|
1751
|
+
},
|
|
1752
|
+
capabilities: {
|
|
1753
|
+
text: true,
|
|
1754
|
+
vision: false,
|
|
1755
|
+
moe: false
|
|
1756
|
+
},
|
|
1757
|
+
tensors,
|
|
1758
|
+
nodes,
|
|
1759
|
+
executionOrder,
|
|
1760
|
+
inputs: ["pcm_in"],
|
|
1761
|
+
outputs: ["dac_codes"]
|
|
1762
|
+
};
|
|
1763
|
+
}
|
|
1764
|
+
function generateOuteTtsBackboneGraph(rawConfig, dtype, groupSize, kvDtype) {
|
|
1765
|
+
const cfg = parseOuteTtsConfig(rawConfig);
|
|
1766
|
+
const useQ4 = dtype === "q4";
|
|
1767
|
+
const groupSz = groupSize ?? DEFAULT_GROUP_SIZE;
|
|
1768
|
+
const hidden_size = cfg.hidden_size;
|
|
1769
|
+
const num_layers = cfg.num_layers;
|
|
1770
|
+
const num_heads = cfg.num_heads;
|
|
1771
|
+
const num_kv_heads = cfg.num_kv_heads;
|
|
1772
|
+
const head_dim = cfg.head_dim;
|
|
1773
|
+
const rope_dim = head_dim;
|
|
1774
|
+
const rope_base = cfg.rope_base;
|
|
1775
|
+
const rms_norm_eps = cfg.rms_norm_eps;
|
|
1776
|
+
const intermediate_size = cfg.intermediate_size;
|
|
1777
|
+
const vocab_size = cfg.vocab_size;
|
|
1778
|
+
const q_dim = num_heads * head_dim;
|
|
1779
|
+
const kv_dim = num_kv_heads * head_dim;
|
|
1780
|
+
const kvCacheDtype = kvDtype ?? "f32";
|
|
1781
|
+
const config = {
|
|
1782
|
+
hidden_size,
|
|
1783
|
+
num_layers,
|
|
1784
|
+
num_heads,
|
|
1785
|
+
num_kv_heads,
|
|
1786
|
+
head_dim,
|
|
1787
|
+
intermediate_size,
|
|
1788
|
+
vocab_size,
|
|
1789
|
+
context_length: cfg.max_position_embeddings,
|
|
1790
|
+
rms_norm_eps,
|
|
1791
|
+
norm_type: "rmsnorm",
|
|
1792
|
+
rope_base,
|
|
1793
|
+
rope_dim,
|
|
1794
|
+
kv_layout: "LHSd",
|
|
1795
|
+
is_moe: false,
|
|
1796
|
+
has_vision_tower: false
|
|
1797
|
+
};
|
|
1798
|
+
const tensors = {};
|
|
1799
|
+
const nodes = [];
|
|
1800
|
+
const executionOrder = [];
|
|
1801
|
+
const addTensor = (desc) => {
|
|
1802
|
+
tensors[desc.name] = desc;
|
|
1803
|
+
};
|
|
1804
|
+
const addNode = (node) => {
|
|
1805
|
+
nodes.push(node);
|
|
1806
|
+
executionOrder.push(node.id);
|
|
1807
|
+
};
|
|
1808
|
+
const addLinear = (id, activationInput, weightName, weightShape, outputName, K, N) => {
|
|
1809
|
+
if (useQ4) {
|
|
1810
|
+
const totalElements = K * N;
|
|
1811
|
+
addTensor({
|
|
1812
|
+
name: `${weightName}.q`,
|
|
1813
|
+
shape: [Math.ceil(totalElements / 8)],
|
|
1814
|
+
dtype: "u32",
|
|
1815
|
+
storage: "constant"
|
|
1816
|
+
});
|
|
1817
|
+
addTensor({
|
|
1818
|
+
name: `${weightName}.scales`,
|
|
1819
|
+
shape: [Math.ceil(totalElements / groupSz)],
|
|
1820
|
+
dtype: "f32",
|
|
1821
|
+
storage: "constant"
|
|
1822
|
+
});
|
|
1823
|
+
addTensor({
|
|
1824
|
+
name: `${weightName}.zeros`,
|
|
1825
|
+
shape: [Math.ceil(totalElements / groupSz)],
|
|
1826
|
+
dtype: "f32",
|
|
1827
|
+
storage: "constant"
|
|
1828
|
+
});
|
|
1829
|
+
addNode({
|
|
1830
|
+
id,
|
|
1831
|
+
opType: "MatMulInt4",
|
|
1832
|
+
inputs: [
|
|
1833
|
+
activationInput,
|
|
1834
|
+
`${weightName}.q`,
|
|
1835
|
+
`${weightName}.scales`,
|
|
1836
|
+
`${weightName}.zeros`
|
|
1837
|
+
],
|
|
1838
|
+
outputs: [outputName],
|
|
1839
|
+
attributes: {
|
|
1840
|
+
M_tensor: activationInput,
|
|
1841
|
+
K,
|
|
1842
|
+
N,
|
|
1843
|
+
group_size: groupSz
|
|
1844
|
+
}
|
|
1845
|
+
});
|
|
1846
|
+
return;
|
|
1847
|
+
}
|
|
1848
|
+
addTensor({
|
|
1849
|
+
name: weightName,
|
|
1850
|
+
shape: weightShape,
|
|
1851
|
+
dtype: "f32",
|
|
1852
|
+
storage: "constant",
|
|
1853
|
+
safetensorsKey: weightName
|
|
1854
|
+
});
|
|
1855
|
+
addNode({
|
|
1856
|
+
id,
|
|
1857
|
+
opType: "MatMul",
|
|
1858
|
+
inputs: [activationInput, weightName],
|
|
1859
|
+
outputs: [outputName],
|
|
1860
|
+
attributes: {
|
|
1861
|
+
M_tensor: activationInput,
|
|
1862
|
+
K,
|
|
1863
|
+
N
|
|
1864
|
+
}
|
|
1865
|
+
});
|
|
1866
|
+
};
|
|
1867
|
+
addTensor({
|
|
1868
|
+
name: "input_ids",
|
|
1869
|
+
shape: ["T"],
|
|
1870
|
+
dtype: "u32",
|
|
1871
|
+
storage: "activation"
|
|
1872
|
+
});
|
|
1873
|
+
addTensor({
|
|
1874
|
+
name: "embed_out",
|
|
1875
|
+
shape: ["T", hidden_size],
|
|
1876
|
+
dtype: "f32",
|
|
1877
|
+
storage: "activation"
|
|
1878
|
+
});
|
|
1879
|
+
addTensor({
|
|
1880
|
+
name: CANONICAL_KEYS.EMBED,
|
|
1881
|
+
shape: [vocab_size, hidden_size],
|
|
1882
|
+
dtype: "f32",
|
|
1883
|
+
storage: "constant",
|
|
1884
|
+
safetensorsKey: CANONICAL_KEYS.EMBED
|
|
1885
|
+
});
|
|
1886
|
+
addNode({
|
|
1887
|
+
id: "embed",
|
|
1888
|
+
opType: "Embedding",
|
|
1889
|
+
inputs: ["input_ids", CANONICAL_KEYS.EMBED],
|
|
1890
|
+
outputs: ["embed_out"],
|
|
1891
|
+
attributes: {
|
|
1892
|
+
vocab_size,
|
|
1893
|
+
hidden_size
|
|
1894
|
+
}
|
|
1895
|
+
});
|
|
1896
|
+
let prevOutput = "embed_out";
|
|
1897
|
+
for (let i = 0; i < num_layers; i++) {
|
|
1898
|
+
const prefix = `layer${i}`;
|
|
1899
|
+
const normW = CANONICAL_KEYS.layerInputNorm(i);
|
|
1900
|
+
addTensor({
|
|
1901
|
+
name: normW,
|
|
1902
|
+
shape: [hidden_size],
|
|
1903
|
+
dtype: "f32",
|
|
1904
|
+
storage: "constant",
|
|
1905
|
+
safetensorsKey: normW
|
|
1906
|
+
});
|
|
1907
|
+
const norm1Out = `${prefix}_norm1_out`;
|
|
1908
|
+
addTensor({
|
|
1909
|
+
name: norm1Out,
|
|
1910
|
+
shape: ["T", hidden_size],
|
|
1911
|
+
dtype: "f32",
|
|
1912
|
+
storage: "activation"
|
|
1913
|
+
});
|
|
1914
|
+
addNode({
|
|
1915
|
+
id: `${prefix}_norm1`,
|
|
1916
|
+
opType: "RMSNorm",
|
|
1917
|
+
inputs: [prevOutput, normW],
|
|
1918
|
+
outputs: [norm1Out],
|
|
1919
|
+
attributes: {
|
|
1920
|
+
hidden_size,
|
|
1921
|
+
eps: rms_norm_eps,
|
|
1922
|
+
seq_len_tensor: prevOutput
|
|
1923
|
+
}
|
|
1924
|
+
});
|
|
1925
|
+
const qOut = `${prefix}_q`;
|
|
1926
|
+
const kOut = `${prefix}_k`;
|
|
1927
|
+
const vOut = `${prefix}_v`;
|
|
1928
|
+
addTensor({
|
|
1929
|
+
name: qOut,
|
|
1930
|
+
shape: ["T", q_dim],
|
|
1931
|
+
dtype: "f32",
|
|
1932
|
+
storage: "activation"
|
|
1933
|
+
});
|
|
1934
|
+
addTensor({
|
|
1935
|
+
name: kOut,
|
|
1936
|
+
shape: ["T", kv_dim],
|
|
1937
|
+
dtype: "f32",
|
|
1938
|
+
storage: "activation"
|
|
1939
|
+
});
|
|
1940
|
+
addTensor({
|
|
1941
|
+
name: vOut,
|
|
1942
|
+
shape: ["T", kv_dim],
|
|
1943
|
+
dtype: "f32",
|
|
1944
|
+
storage: "activation"
|
|
1945
|
+
});
|
|
1946
|
+
addLinear(`${prefix}_q_proj`, norm1Out, CANONICAL_KEYS.qProj(i), [q_dim, hidden_size], qOut, hidden_size, q_dim);
|
|
1947
|
+
addLinear(`${prefix}_k_proj`, norm1Out, CANONICAL_KEYS.kProj(i), [kv_dim, hidden_size], kOut, hidden_size, kv_dim);
|
|
1948
|
+
addLinear(`${prefix}_v_proj`, norm1Out, CANONICAL_KEYS.vProj(i), [kv_dim, hidden_size], vOut, hidden_size, kv_dim);
|
|
1949
|
+
const qNormW = CANONICAL_KEYS.qNorm(i);
|
|
1950
|
+
const kNormW = CANONICAL_KEYS.kNorm(i);
|
|
1951
|
+
addTensor({
|
|
1952
|
+
name: qNormW,
|
|
1953
|
+
shape: [head_dim],
|
|
1954
|
+
dtype: "f32",
|
|
1955
|
+
storage: "constant",
|
|
1956
|
+
safetensorsKey: qNormW
|
|
1957
|
+
});
|
|
1958
|
+
addTensor({
|
|
1959
|
+
name: kNormW,
|
|
1960
|
+
shape: [head_dim],
|
|
1961
|
+
dtype: "f32",
|
|
1962
|
+
storage: "constant",
|
|
1963
|
+
safetensorsKey: kNormW
|
|
1964
|
+
});
|
|
1965
|
+
const qNormed = `${prefix}_q_normed`;
|
|
1966
|
+
const kNormed = `${prefix}_k_normed`;
|
|
1967
|
+
addTensor({
|
|
1968
|
+
name: qNormed,
|
|
1969
|
+
shape: ["T", q_dim],
|
|
1970
|
+
dtype: "f32",
|
|
1971
|
+
storage: "activation"
|
|
1972
|
+
});
|
|
1973
|
+
addTensor({
|
|
1974
|
+
name: kNormed,
|
|
1975
|
+
shape: ["T", kv_dim],
|
|
1976
|
+
dtype: "f32",
|
|
1977
|
+
storage: "activation"
|
|
1978
|
+
});
|
|
1979
|
+
addNode({
|
|
1980
|
+
id: `${prefix}_q_norm`,
|
|
1981
|
+
opType: "RMSNorm",
|
|
1982
|
+
inputs: [qOut, qNormW],
|
|
1983
|
+
outputs: [qNormed],
|
|
1984
|
+
attributes: {
|
|
1985
|
+
hidden_size: head_dim,
|
|
1986
|
+
eps: rms_norm_eps,
|
|
1987
|
+
seq_len_tensor: qOut
|
|
1988
|
+
}
|
|
1989
|
+
});
|
|
1990
|
+
addNode({
|
|
1991
|
+
id: `${prefix}_k_norm`,
|
|
1992
|
+
opType: "RMSNorm",
|
|
1993
|
+
inputs: [kOut, kNormW],
|
|
1994
|
+
outputs: [kNormed],
|
|
1995
|
+
attributes: {
|
|
1996
|
+
hidden_size: head_dim,
|
|
1997
|
+
eps: rms_norm_eps,
|
|
1998
|
+
seq_len_tensor: kOut
|
|
1999
|
+
}
|
|
2000
|
+
});
|
|
2001
|
+
addNode({
|
|
2002
|
+
id: `${prefix}_rope`,
|
|
2003
|
+
opType: "RoPE",
|
|
2004
|
+
inputs: [qNormed, kNormed],
|
|
2005
|
+
outputs: [qNormed, kNormed],
|
|
2006
|
+
attributes: {
|
|
2007
|
+
head_dim,
|
|
2008
|
+
num_q_heads: num_heads,
|
|
2009
|
+
num_kv_heads,
|
|
2010
|
+
rope_base,
|
|
2011
|
+
rope_dim,
|
|
2012
|
+
seq_len_tensor: qNormed
|
|
2013
|
+
}
|
|
2014
|
+
});
|
|
2015
|
+
const kCache = `${prefix}_k_cache`;
|
|
2016
|
+
const vCache = `${prefix}_v_cache`;
|
|
2017
|
+
addTensor({
|
|
2018
|
+
name: kCache,
|
|
2019
|
+
shape: ["L_max", kv_dim],
|
|
2020
|
+
dtype: kvCacheDtype,
|
|
2021
|
+
storage: "kv_cache"
|
|
2022
|
+
});
|
|
2023
|
+
addTensor({
|
|
2024
|
+
name: vCache,
|
|
2025
|
+
shape: ["L_max", kv_dim],
|
|
2026
|
+
dtype: kvCacheDtype,
|
|
2027
|
+
storage: "kv_cache"
|
|
2028
|
+
});
|
|
2029
|
+
addNode({
|
|
2030
|
+
id: `${prefix}_kv_cache_k`,
|
|
2031
|
+
opType: "KVCacheAppend",
|
|
2032
|
+
inputs: [kNormed],
|
|
2033
|
+
outputs: [kCache],
|
|
2034
|
+
attributes: {
|
|
2035
|
+
width: kv_dim,
|
|
2036
|
+
T_tensor: kNormed
|
|
2037
|
+
}
|
|
2038
|
+
});
|
|
2039
|
+
addNode({
|
|
2040
|
+
id: `${prefix}_kv_cache_v`,
|
|
2041
|
+
opType: "KVCacheAppend",
|
|
2042
|
+
inputs: [vOut],
|
|
2043
|
+
outputs: [vCache],
|
|
2044
|
+
attributes: {
|
|
2045
|
+
width: kv_dim,
|
|
2046
|
+
T_tensor: vOut
|
|
2047
|
+
}
|
|
2048
|
+
});
|
|
2049
|
+
const attnOut = `${prefix}_attn_out`;
|
|
2050
|
+
addTensor({
|
|
2051
|
+
name: attnOut,
|
|
2052
|
+
shape: ["T", q_dim],
|
|
2053
|
+
dtype: "f32",
|
|
2054
|
+
storage: "activation"
|
|
2055
|
+
});
|
|
2056
|
+
addNode({
|
|
2057
|
+
id: `${prefix}_attn`,
|
|
2058
|
+
opType: "Attention",
|
|
2059
|
+
inputs: [
|
|
2060
|
+
qNormed,
|
|
2061
|
+
kCache,
|
|
2062
|
+
vCache
|
|
2063
|
+
],
|
|
2064
|
+
outputs: [attnOut],
|
|
2065
|
+
attributes: {
|
|
2066
|
+
hidden_size: q_dim,
|
|
2067
|
+
num_q_heads: num_heads,
|
|
2068
|
+
num_kv_heads,
|
|
2069
|
+
head_dim,
|
|
2070
|
+
causal: true,
|
|
2071
|
+
layer_index: i
|
|
2072
|
+
}
|
|
2073
|
+
});
|
|
2074
|
+
const oProjOut = `${prefix}_o_proj_out`;
|
|
2075
|
+
addTensor({
|
|
2076
|
+
name: oProjOut,
|
|
2077
|
+
shape: ["T", hidden_size],
|
|
2078
|
+
dtype: "f32",
|
|
2079
|
+
storage: "activation"
|
|
2080
|
+
});
|
|
2081
|
+
addLinear(`${prefix}_o_proj`, attnOut, CANONICAL_KEYS.oProj(i), [hidden_size, q_dim], oProjOut, q_dim, hidden_size);
|
|
2082
|
+
const resid1Out = `${prefix}_resid1`;
|
|
2083
|
+
const postNormW = CANONICAL_KEYS.layerPostAttnNorm(i);
|
|
2084
|
+
const norm2Out = `${prefix}_norm2_out`;
|
|
2085
|
+
addTensor({
|
|
2086
|
+
name: resid1Out,
|
|
2087
|
+
shape: ["T", hidden_size],
|
|
2088
|
+
dtype: "f32",
|
|
2089
|
+
storage: "activation"
|
|
2090
|
+
});
|
|
2091
|
+
addTensor({
|
|
2092
|
+
name: postNormW,
|
|
2093
|
+
shape: [hidden_size],
|
|
2094
|
+
dtype: "f32",
|
|
2095
|
+
storage: "constant",
|
|
2096
|
+
safetensorsKey: postNormW
|
|
2097
|
+
});
|
|
2098
|
+
addTensor({
|
|
2099
|
+
name: norm2Out,
|
|
2100
|
+
shape: ["T", hidden_size],
|
|
2101
|
+
dtype: "f32",
|
|
2102
|
+
storage: "activation"
|
|
2103
|
+
});
|
|
2104
|
+
addNode({
|
|
2105
|
+
id: `${prefix}_resid1_norm2`,
|
|
2106
|
+
opType: "ResidualRMSNorm",
|
|
2107
|
+
inputs: [
|
|
2108
|
+
prevOutput,
|
|
2109
|
+
oProjOut,
|
|
2110
|
+
postNormW
|
|
2111
|
+
],
|
|
2112
|
+
outputs: [resid1Out, norm2Out],
|
|
2113
|
+
attributes: {
|
|
2114
|
+
hidden_size,
|
|
2115
|
+
eps: rms_norm_eps,
|
|
2116
|
+
seq_len_tensor: prevOutput
|
|
2117
|
+
}
|
|
2118
|
+
});
|
|
2119
|
+
const gateOut = `${prefix}_gate_out`;
|
|
2120
|
+
const upOut = `${prefix}_up_out`;
|
|
2121
|
+
const swigluOut = `${prefix}_swiglu_out`;
|
|
2122
|
+
const mlpOut = `${prefix}_mlp_out`;
|
|
2123
|
+
addTensor({
|
|
2124
|
+
name: gateOut,
|
|
2125
|
+
shape: ["T", intermediate_size],
|
|
2126
|
+
dtype: "f32",
|
|
2127
|
+
storage: "activation"
|
|
2128
|
+
});
|
|
2129
|
+
addTensor({
|
|
2130
|
+
name: upOut,
|
|
2131
|
+
shape: ["T", intermediate_size],
|
|
2132
|
+
dtype: "f32",
|
|
2133
|
+
storage: "activation"
|
|
2134
|
+
});
|
|
2135
|
+
addTensor({
|
|
2136
|
+
name: swigluOut,
|
|
2137
|
+
shape: ["T", intermediate_size],
|
|
2138
|
+
dtype: "f32",
|
|
2139
|
+
storage: "activation"
|
|
2140
|
+
});
|
|
2141
|
+
addTensor({
|
|
2142
|
+
name: mlpOut,
|
|
2143
|
+
shape: ["T", hidden_size],
|
|
2144
|
+
dtype: "f32",
|
|
2145
|
+
storage: "activation"
|
|
2146
|
+
});
|
|
2147
|
+
addLinear(`${prefix}_gate`, norm2Out, CANONICAL_KEYS.gateProj(i), [intermediate_size, hidden_size], gateOut, hidden_size, intermediate_size);
|
|
2148
|
+
addLinear(`${prefix}_up`, norm2Out, CANONICAL_KEYS.upProj(i), [intermediate_size, hidden_size], upOut, hidden_size, intermediate_size);
|
|
2149
|
+
addNode({
|
|
2150
|
+
id: `${prefix}_swiglu`,
|
|
2151
|
+
opType: "SwiGLU",
|
|
2152
|
+
inputs: [gateOut, upOut],
|
|
2153
|
+
outputs: [swigluOut],
|
|
2154
|
+
attributes: { count_tensor: gateOut }
|
|
2155
|
+
});
|
|
2156
|
+
addLinear(`${prefix}_down`, swigluOut, CANONICAL_KEYS.downProj(i), [hidden_size, intermediate_size], mlpOut, intermediate_size, hidden_size);
|
|
2157
|
+
const resid2Out = `${prefix}_resid2`;
|
|
2158
|
+
addTensor({
|
|
2159
|
+
name: resid2Out,
|
|
2160
|
+
shape: ["T", hidden_size],
|
|
2161
|
+
dtype: "f32",
|
|
2162
|
+
storage: "activation"
|
|
2163
|
+
});
|
|
2164
|
+
addNode({
|
|
2165
|
+
id: `${prefix}_resid2`,
|
|
2166
|
+
opType: "Add",
|
|
2167
|
+
inputs: [resid1Out, mlpOut],
|
|
2168
|
+
outputs: [resid2Out],
|
|
2169
|
+
attributes: {
|
|
2170
|
+
count_tensor: resid1Out,
|
|
2171
|
+
hidden_size
|
|
2172
|
+
}
|
|
2173
|
+
});
|
|
2174
|
+
prevOutput = resid2Out;
|
|
2175
|
+
}
|
|
2176
|
+
addTensor({
|
|
2177
|
+
name: CANONICAL_KEYS.FINAL_NORM,
|
|
2178
|
+
shape: [hidden_size],
|
|
2179
|
+
dtype: "f32",
|
|
2180
|
+
storage: "constant",
|
|
2181
|
+
safetensorsKey: CANONICAL_KEYS.FINAL_NORM
|
|
2182
|
+
});
|
|
2183
|
+
addTensor({
|
|
2184
|
+
name: "final_norm_out",
|
|
2185
|
+
shape: ["T", hidden_size],
|
|
2186
|
+
dtype: "f32",
|
|
2187
|
+
storage: "activation"
|
|
2188
|
+
});
|
|
2189
|
+
addNode({
|
|
2190
|
+
id: "final_norm",
|
|
2191
|
+
opType: "RMSNorm",
|
|
2192
|
+
inputs: [prevOutput, CANONICAL_KEYS.FINAL_NORM],
|
|
2193
|
+
outputs: ["final_norm_out"],
|
|
2194
|
+
attributes: {
|
|
2195
|
+
hidden_size,
|
|
2196
|
+
eps: rms_norm_eps,
|
|
2197
|
+
seq_len_tensor: prevOutput
|
|
2198
|
+
}
|
|
2199
|
+
});
|
|
2200
|
+
addTensor({
|
|
2201
|
+
name: "final_norm_last",
|
|
2202
|
+
shape: [1, hidden_size],
|
|
2203
|
+
dtype: "f32",
|
|
2204
|
+
storage: "activation"
|
|
2205
|
+
});
|
|
2206
|
+
addNode({
|
|
2207
|
+
id: "slice_last_row",
|
|
2208
|
+
opType: "SliceLastRow",
|
|
2209
|
+
inputs: ["final_norm_out"],
|
|
2210
|
+
outputs: ["final_norm_last"],
|
|
2211
|
+
attributes: { width: hidden_size }
|
|
2212
|
+
});
|
|
2213
|
+
addTensor({
|
|
2214
|
+
name: "logits",
|
|
2215
|
+
shape: [1, vocab_size],
|
|
2216
|
+
dtype: "f32",
|
|
2217
|
+
storage: "activation"
|
|
2218
|
+
});
|
|
2219
|
+
addNode({
|
|
2220
|
+
id: "lm_head",
|
|
2221
|
+
opType: "MatMul",
|
|
2222
|
+
inputs: ["final_norm_last", CANONICAL_KEYS.EMBED],
|
|
2223
|
+
outputs: ["logits"],
|
|
2224
|
+
attributes: {
|
|
2225
|
+
M_tensor: "final_norm_last",
|
|
2226
|
+
K: hidden_size,
|
|
2227
|
+
N: vocab_size
|
|
2228
|
+
}
|
|
2229
|
+
});
|
|
2230
|
+
return {
|
|
2231
|
+
architecture: "Qwen3ForCausalLM",
|
|
2232
|
+
config,
|
|
2233
|
+
capabilities: {
|
|
2234
|
+
text: true,
|
|
2235
|
+
vision: false,
|
|
2236
|
+
moe: false
|
|
2237
|
+
},
|
|
2238
|
+
tensors,
|
|
2239
|
+
nodes,
|
|
2240
|
+
executionOrder,
|
|
2241
|
+
inputs: ["input_ids"],
|
|
2242
|
+
outputs: ["logits"]
|
|
2243
|
+
};
|
|
2244
|
+
}
|
|
2245
|
+
|
|
2246
|
+
//#endregion
|
|
2247
|
+
//#region src/gpu/outetts-encode.ts
|
|
2248
|
+
/**
|
|
2249
|
+
* OuteTTS voice-clone helpers — prosody feature extraction + word alignment for
|
|
2250
|
+
* `OuteTTS.encodeSpeaker` (raw audio → OuteSpeaker). These are a faithful port of
|
|
2251
|
+
* outetts v3 `audio_processor.Features` (energy / spectral_centroid / pitch) and the
|
|
2252
|
+
* `create_speaker_from_dict` word loop, so the per-word feature tokens the prompt
|
|
2253
|
+
* builder emits match the reference speaker JSONs the presets ship.
|
|
2254
|
+
*
|
|
2255
|
+
* The DAC encode (audio → c1/c2 codes) is done bit-exact on the GPU
|
|
2256
|
+
* (generateDacSpeechEncoderGraph + RVQEncode, validated 450/450 vs torch). The
|
|
2257
|
+
* ALIGNMENT here is pragmatic: outetts uses Whisper word-level timestamps to slice
|
|
2258
|
+
* the code stream per word; with no on-device Whisper we distribute the encoded
|
|
2259
|
+
* frames across the transcript's words PROPORTIONALLY to word length (chars), then
|
|
2260
|
+
* compute each word's prosody from its audio span. This matches the reference shape
|
|
2261
|
+
* and the global/per-word feature math exactly; only the word→frame boundaries are
|
|
2262
|
+
* approximated (documented on `OuteTTS.encodeSpeaker`).
|
|
2263
|
+
*/
|
|
2264
|
+
const DAC_TPS = 75;
|
|
2265
|
+
const WHITESPACE_RE = /\s+/;
|
|
2266
|
+
const PITCH_MIN_FREQ = 75;
|
|
2267
|
+
const PITCH_MAX_FREQ = 600;
|
|
2268
|
+
const PITCH_FRAME_LEN = 400;
|
|
2269
|
+
const PITCH_HOP = 160;
|
|
2270
|
+
const PITCH_THRESHOLD = .3;
|
|
2271
|
+
const FEATURE_EPS = 1e-10;
|
|
2272
|
+
/** Round a [0,1] value to the [0,100] integer bucket the prompt builder expects. */
|
|
2273
|
+
function scaleValue(value) {
|
|
2274
|
+
return Math.round(Math.min(Math.max(value, 0), 1) * 100);
|
|
2275
|
+
}
|
|
2276
|
+
/** Naive DFT magnitude spectrum of a real signal (length N → ⌊N/2⌋+1 bins). */
|
|
2277
|
+
function rfftMag(x) {
|
|
2278
|
+
const N = x.length;
|
|
2279
|
+
const bins = Math.floor(N / 2) + 1;
|
|
2280
|
+
const mag = new Float32Array(bins);
|
|
2281
|
+
for (let k = 0; k < bins; k++) {
|
|
2282
|
+
let re = 0;
|
|
2283
|
+
let im = 0;
|
|
2284
|
+
const w = -2 * Math.PI * k / N;
|
|
2285
|
+
for (let n = 0; n < N; n++) {
|
|
2286
|
+
const a = w * n;
|
|
2287
|
+
re += x[n] * Math.cos(a);
|
|
2288
|
+
im += x[n] * Math.sin(a);
|
|
2289
|
+
}
|
|
2290
|
+
mag[k] = Math.hypot(re, im);
|
|
2291
|
+
}
|
|
2292
|
+
return mag;
|
|
2293
|
+
}
|
|
2294
|
+
/**
|
|
2295
|
+
* RMS energy of a clip, normalized to [0,1] (the reference clamps via scaleValue).
|
|
2296
|
+
* Matches `Features.extract_audio_features` energy = sqrt(mean(x^2)).
|
|
2297
|
+
*/
|
|
2298
|
+
function rmsEnergy(audio) {
|
|
2299
|
+
let sum = 0;
|
|
2300
|
+
for (let i = 0; i < audio.length; i++) sum += audio[i] * audio[i];
|
|
2301
|
+
return audio.length ? Math.sqrt(sum / audio.length) : 0;
|
|
2302
|
+
}
|
|
2303
|
+
/**
|
|
2304
|
+
* Spectral centroid normalized to [0,1] of (sr/2). Matches the reference:
|
|
2305
|
+
* centroid = Σ f·|rfft(x)| / Σ|rfft(x)|, then / (sr/2).
|
|
2306
|
+
*/
|
|
2307
|
+
function spectralCentroid(audio, sr) {
|
|
2308
|
+
if (audio.length === 0) return 0;
|
|
2309
|
+
const spec = rfftMag(audio);
|
|
2310
|
+
const nyquist = sr / 2;
|
|
2311
|
+
const bins = spec.length;
|
|
2312
|
+
let specSum = FEATURE_EPS;
|
|
2313
|
+
let weighted = 0;
|
|
2314
|
+
for (let k = 0; k < bins; k++) {
|
|
2315
|
+
const freq = bins > 1 ? k / (bins - 1) * nyquist : 0;
|
|
2316
|
+
weighted += freq * spec[k];
|
|
2317
|
+
specSum += spec[k];
|
|
2318
|
+
}
|
|
2319
|
+
return weighted / specSum / nyquist;
|
|
2320
|
+
}
|
|
2321
|
+
/** Hann window of length N (periodic=false, matching torch.hann_window default). */
|
|
2322
|
+
function hannWindow(N) {
|
|
2323
|
+
const w = new Float32Array(N);
|
|
2324
|
+
for (let n = 0; n < N; n++) w[n] = .5 - .5 * Math.cos(2 * Math.PI * n / (N - 1));
|
|
2325
|
+
return w;
|
|
2326
|
+
}
|
|
2327
|
+
/**
|
|
2328
|
+
* Autocorrelation-based average pitch normalized to [0,1] over [min,max] freq.
|
|
2329
|
+
* Faithful port of `calculate_pitch` + `extract_single_pitch_value`: framed Hann-
|
|
2330
|
+
* windowed autocorrelation, per-frame parabolic-interpolated peak in the valid lag
|
|
2331
|
+
* range, voicing threshold, then the mean over frames mapped to [0,1].
|
|
2332
|
+
*/
|
|
2333
|
+
/** Hann-windowed autocorrelation of one frame into `autocorr` (lag 0..frameLen-1). */
|
|
2334
|
+
function frameAutocorr(padded, off, window, frame, autocorr) {
|
|
2335
|
+
for (let i = 0; i < PITCH_FRAME_LEN; i++) frame[i] = padded[off + i] * window[i];
|
|
2336
|
+
for (let lag = 0; lag < PITCH_FRAME_LEN; lag++) {
|
|
2337
|
+
let acc = 0;
|
|
2338
|
+
for (let i = 0; i + lag < PITCH_FRAME_LEN; i++) acc += frame[i] * frame[i + lag];
|
|
2339
|
+
autocorr[lag] = acc;
|
|
2340
|
+
}
|
|
2341
|
+
}
|
|
2342
|
+
/** Voiced pitch (Hz) of one autocorrelation frame, or 0 if unvoiced. */
|
|
2343
|
+
function framePitch(autocorr, minIdx, maxIdx, sr) {
|
|
2344
|
+
let peakIdx = minIdx;
|
|
2345
|
+
let peakVal = Number.NEGATIVE_INFINITY;
|
|
2346
|
+
for (let lag = minIdx; lag < maxIdx; lag++) if (autocorr[lag] > peakVal) {
|
|
2347
|
+
peakVal = autocorr[lag];
|
|
2348
|
+
peakIdx = lag;
|
|
2349
|
+
}
|
|
2350
|
+
const idx = Math.min(Math.max(peakIdx, 1), PITCH_FRAME_LEN - 2);
|
|
2351
|
+
const alpha = autocorr[idx - 1];
|
|
2352
|
+
const beta = autocorr[idx];
|
|
2353
|
+
const gamma = autocorr[idx + 1];
|
|
2354
|
+
let delta = .5 * (alpha - gamma) / (alpha - 2 * beta + gamma + 1e-8);
|
|
2355
|
+
if (!(peakIdx > 0 && peakIdx < PITCH_FRAME_LEN - 1)) delta = 0;
|
|
2356
|
+
const period = (peakIdx + delta) / sr;
|
|
2357
|
+
let pitch = period > 0 ? 1 / period : 0;
|
|
2358
|
+
if (!(peakVal / (autocorr[0] + 1e-8) > PITCH_THRESHOLD)) pitch = 0;
|
|
2359
|
+
return Math.min(Math.max(pitch, PITCH_MIN_FREQ), PITCH_MAX_FREQ);
|
|
2360
|
+
}
|
|
2361
|
+
function averagePitch(audio, sr) {
|
|
2362
|
+
const N = audio.length;
|
|
2363
|
+
if (N < PITCH_FRAME_LEN) return 0;
|
|
2364
|
+
const padLen = (PITCH_FRAME_LEN - N % PITCH_HOP) % PITCH_HOP;
|
|
2365
|
+
const padded = new Float32Array(N + padLen);
|
|
2366
|
+
padded.set(audio);
|
|
2367
|
+
const numFrames = Math.floor((padded.length - PITCH_FRAME_LEN) / PITCH_HOP) + 1;
|
|
2368
|
+
if (numFrames <= 0) return 0;
|
|
2369
|
+
const window = hannWindow(PITCH_FRAME_LEN);
|
|
2370
|
+
const minIdx = Math.max(1, Math.floor(sr / PITCH_MAX_FREQ));
|
|
2371
|
+
const maxIdx = Math.min(PITCH_FRAME_LEN, Math.floor(sr / PITCH_MIN_FREQ));
|
|
2372
|
+
const frame = new Float32Array(PITCH_FRAME_LEN);
|
|
2373
|
+
const autocorr = new Float32Array(PITCH_FRAME_LEN);
|
|
2374
|
+
let pitchSum = 0;
|
|
2375
|
+
for (let f = 0; f < numFrames; f++) {
|
|
2376
|
+
frameAutocorr(padded, f * PITCH_HOP, window, frame, autocorr);
|
|
2377
|
+
pitchSum += framePitch(autocorr, minIdx, maxIdx, sr);
|
|
2378
|
+
}
|
|
2379
|
+
return (pitchSum / numFrames - PITCH_MIN_FREQ) / (PITCH_MAX_FREQ - PITCH_MIN_FREQ);
|
|
2380
|
+
}
|
|
2381
|
+
/**
|
|
2382
|
+
* Port of `Features.extract_audio_features`: energy (RMS), spectral_centroid (FFT
|
|
2383
|
+
* centroid), pitch (autocorrelation), each scaled to the [0,100] integer bucket.
|
|
2384
|
+
* Returns all-zero defaults for an empty/invalid clip (matching the reference).
|
|
2385
|
+
*/
|
|
2386
|
+
function extractAudioFeatures(audio, sr) {
|
|
2387
|
+
if (audio.length === 0) return {
|
|
2388
|
+
energy: 0,
|
|
2389
|
+
spectral_centroid: 0,
|
|
2390
|
+
pitch: 0
|
|
2391
|
+
};
|
|
2392
|
+
return {
|
|
2393
|
+
energy: scaleValue(rmsEnergy(audio)),
|
|
2394
|
+
spectral_centroid: scaleValue(spectralCentroid(audio, sr)),
|
|
2395
|
+
pitch: scaleValue(averagePitch(audio, sr))
|
|
2396
|
+
};
|
|
2397
|
+
}
|
|
2398
|
+
/**
|
|
2399
|
+
* Pragmatic word→frame alignment (no on-device Whisper): split the transcript into
|
|
2400
|
+
* whitespace words and distribute the `numFrames` encoded frames across them
|
|
2401
|
+
* proportionally to each word's character length. Returns contiguous [start,end)
|
|
2402
|
+
* frame spans covering the whole clip. This stands in for outetts' Whisper word-
|
|
2403
|
+
* timestamp slicing — the codes themselves are exact; only the per-word boundaries
|
|
2404
|
+
* are approximate.
|
|
2405
|
+
*/
|
|
2406
|
+
function alignWordsProportional(transcript, numFrames) {
|
|
2407
|
+
const words = transcript.split(WHITESPACE_RE).filter((w) => w.length > 0);
|
|
2408
|
+
if (words.length === 0) return [];
|
|
2409
|
+
const weights = words.map((w) => w.length);
|
|
2410
|
+
const totalWeight = weights.reduce((a, b) => a + b, 0);
|
|
2411
|
+
const out = [];
|
|
2412
|
+
let acc = 0;
|
|
2413
|
+
for (let i = 0; i < words.length; i++) {
|
|
2414
|
+
const startFrame = i === 0 ? 0 : Math.round(acc / totalWeight * numFrames);
|
|
2415
|
+
acc += weights[i];
|
|
2416
|
+
const endFrame = i === words.length - 1 ? numFrames : Math.round(acc / totalWeight * numFrames);
|
|
2417
|
+
out.push({
|
|
2418
|
+
word: words[i],
|
|
2419
|
+
startFrame,
|
|
2420
|
+
endFrame: Math.max(endFrame, startFrame)
|
|
2421
|
+
});
|
|
2422
|
+
}
|
|
2423
|
+
return out;
|
|
2424
|
+
}
|
|
2425
|
+
|
|
2426
|
+
//#endregion
|
|
2427
|
+
//#region src/gpu/outetts.ts
|
|
2428
|
+
/**
|
|
2429
|
+
* OuteTTS — native text-to-speech engine for Gerbil's WebGPU backend.
|
|
2430
|
+
*
|
|
2431
|
+
* OuteTTS-1.0-0.6B (OuteAI/OuteTTS-1.0-0.6B) is a two-stage codec-LM TTS, the same
|
|
2432
|
+
* shape as Kani-TTS but with a plain Qwen3-0.6B backbone and the descript-audio-codec
|
|
2433
|
+
* DAC.speech.v1.0 decoder (24 kHz, 2 codebooks):
|
|
2434
|
+
*
|
|
2435
|
+
* 1. CODEC-LM BACKBONE (Qwen3-0.6B): autoregressively emits interleaved audio codes
|
|
2436
|
+
* <|c1_i|>/<|c2_i|> (one pair per 512-sample frame) into the same vocab as text,
|
|
2437
|
+
* between <|audio_start|> and <|audio_end|>. Standard Qwen3 block math
|
|
2438
|
+
* (RMSNorm/QK-RMSNorm/full RoPE/GQA/SwiGLU) with a KV-cache decode path
|
|
2439
|
+
* (generateOuteTtsBackboneGraph) — simpler than Kani's LFM2 hybrid.
|
|
2440
|
+
* 2. DAC.speech.v1.0 DECODER: RVQ dequant (2× codebook lookup + out_proj) → latent
|
|
2441
|
+
* [1024,T] → DacDecoder (Conv1d + 4 upsample blocks + Snake + tanh) → 24 kHz PCM.
|
|
2442
|
+
* Validated bit-exact vs the torch DAC reference (scripts/engine/test-outetts-codec.mjs,
|
|
2443
|
+
* latent max|err| 2.9e-6, PCM max|err| 1.3e-6, cos 1.0).
|
|
2444
|
+
*
|
|
2445
|
+
* The AR loop runs on the host with full-logit readback; each step samples the next
|
|
2446
|
+
* audio token constrained to the current codebook's valid window (c1 on even
|
|
2447
|
+
* positions, c2 on odd), until <|audio_end|>. The collected tokens are split into a
|
|
2448
|
+
* [2, T] code grid (audioTokensToDacCodes) and decoded once through the DAC graph.
|
|
2449
|
+
*
|
|
2450
|
+
* This is the FOUNDATION: a single preset voice (OuteTTS's bundled en-female-1-neutral
|
|
2451
|
+
* speaker, embedded as in-context reference codes). Cloning, int4, the speak()/useTTS
|
|
2452
|
+
* public surface, R2 mirror, and the playground are a later phase.
|
|
2453
|
+
*/
|
|
2454
|
+
const DEFAULT_MAX_NEW_TOKENS = 4096;
|
|
2455
|
+
const CONTROL_CHARS_RE = /[\x00-\x1F\x7F-\x9F-]/g;
|
|
2456
|
+
const HIRAGANA_RE = /[-ゟ]/;
|
|
2457
|
+
const KATAKANA_RE = /[゠-ヿ]/;
|
|
2458
|
+
const HAN_RE = /[一-鿿]/;
|
|
2459
|
+
const HTTP_URL_RE = /^https?:\/\//;
|
|
2460
|
+
const TRAILING_SLASH_RE = /\/$/;
|
|
2461
|
+
/** Port of PromptProcessor.text_normalizations (the ftfy step is approximated). */
|
|
2462
|
+
function normalizeOuteText(input) {
|
|
2463
|
+
let text = input.normalize("NFKC");
|
|
2464
|
+
text = text.replace(/…/g, "...");
|
|
2465
|
+
text = text.replace(/\.{2,}/g, "...");
|
|
2466
|
+
text = text.replace(/[“”„‟«»]/g, "\"");
|
|
2467
|
+
text = text.replace(/[‘’‛‹›`´]/g, "'");
|
|
2468
|
+
text = text.replace(/[–—―−‐]/g, "-");
|
|
2469
|
+
text = text.replace(/-{2,}/g, "-");
|
|
2470
|
+
text = text.replace(CONTROL_CHARS_RE, "");
|
|
2471
|
+
text = text.replace(/\s+/g, " ");
|
|
2472
|
+
text = text.replace(/\s+([,.?!:;])/g, "$1");
|
|
2473
|
+
text = text.replace(/([,.?!:;])(?=[^\s,.?!:;])/g, "$1 ");
|
|
2474
|
+
text = text.replace(/(\w)\s+'\s*(\w)/g, "$1'$2");
|
|
2475
|
+
text = text.replace(/(\w)\s+'\s*([,.?!:;\s]|$)/g, "$1'$2");
|
|
2476
|
+
text = text.replace(/(\w)\s*'\s*([tsdmreSNTDMRE])\b/gi, "$1'$2");
|
|
2477
|
+
text = text.replace(/([?!])\1+/g, "$1");
|
|
2478
|
+
text = text.replace(/\s+/g, " ");
|
|
2479
|
+
text = text.replace(/"/g, "").trim();
|
|
2480
|
+
return text;
|
|
2481
|
+
}
|
|
2482
|
+
/** Port of PromptProcessor._get_separator (CJK-aware). */
|
|
2483
|
+
function outeSeparator(text) {
|
|
2484
|
+
const hasHiragana = HIRAGANA_RE.test(text);
|
|
2485
|
+
const hasKatakana = KATAKANA_RE.test(text);
|
|
2486
|
+
const hasHan = HAN_RE.test(text);
|
|
2487
|
+
if (hasHiragana || hasKatakana || hasHan) return "。";
|
|
2488
|
+
return ". ";
|
|
2489
|
+
}
|
|
2490
|
+
/**
|
|
2491
|
+
* Port of PromptProcessor.merge_speaker_text: prepend the (normalized) speaker
|
|
2492
|
+
* transcript in front of the target text, joined by the right separator, and return
|
|
2493
|
+
* the trailing-separator string that outetts appends to the speaker's LAST word.
|
|
2494
|
+
*/
|
|
2495
|
+
function mergeSpeakerText(inputText, speakerText) {
|
|
2496
|
+
const sp = speakerText.trim();
|
|
2497
|
+
const separator = outeSeparator(sp);
|
|
2498
|
+
const allowedEnds = separator === "。" ? [
|
|
2499
|
+
"。",
|
|
2500
|
+
"?",
|
|
2501
|
+
"!",
|
|
2502
|
+
"?",
|
|
2503
|
+
"!"
|
|
2504
|
+
] : [
|
|
2505
|
+
".",
|
|
2506
|
+
"?",
|
|
2507
|
+
"!"
|
|
2508
|
+
];
|
|
2509
|
+
let rs = "";
|
|
2510
|
+
if (sp) {
|
|
2511
|
+
const last = sp[sp.length - 1];
|
|
2512
|
+
if (!allowedEnds.includes(last)) rs = separator;
|
|
2513
|
+
else if (separator !== "。") rs = " ";
|
|
2514
|
+
}
|
|
2515
|
+
return {
|
|
2516
|
+
merged: sp + rs + inputText.trim(),
|
|
2517
|
+
sep: rs.trim()
|
|
2518
|
+
};
|
|
2519
|
+
}
|
|
2520
|
+
/** Format one word's feature tokens (energy / spectral_centroid / pitch). */
|
|
2521
|
+
function outeFeatureTokens(f) {
|
|
2522
|
+
return `<|energy_${f?.energy ?? 0}|><|spectral_centroid_${f?.spectral_centroid ?? 0}|><|pitch_${f?.pitch ?? 0}|>`;
|
|
2523
|
+
}
|
|
2524
|
+
/** Port of PromptProcessor.create_codes: one <|word_start|>…<|word_end|> block per word, "\n"-joined. */
|
|
2525
|
+
function createSpeakerCodes(words) {
|
|
2526
|
+
const blocks = [];
|
|
2527
|
+
for (const w of words) {
|
|
2528
|
+
const duration = (w.duration ?? 0).toFixed(2);
|
|
2529
|
+
let block = `${w.word}<|features|><|t_${duration}|>${outeFeatureTokens(w.features)}`;
|
|
2530
|
+
const n = Math.min(w.c1.length, w.c2.length);
|
|
2531
|
+
let pairs = "";
|
|
2532
|
+
for (let i = 0; i < n; i++) pairs += `<|c1_${w.c1[i]}|><|c2_${w.c2[i]}|>`;
|
|
2533
|
+
block += `<|code|>${pairs}`;
|
|
2534
|
+
blocks.push(`<|word_start|>${block}<|word_end|>`);
|
|
2535
|
+
}
|
|
2536
|
+
return blocks.join("\n");
|
|
2537
|
+
}
|
|
2538
|
+
/**
|
|
2539
|
+
* Build the OuteTTS v3 completion prompt STRING (outetts get_completion_prompt). The
|
|
2540
|
+
* whole prompt — text section, speaker word blocks (word + time/energy/spectral/pitch
|
|
2541
|
+
* features + interleaved c1/c2 codes), and the dangling <|word_start|> cue — is
|
|
2542
|
+
* assembled with literal <|...|> markers; the caller tokenizes it in one pass. Exposed
|
|
2543
|
+
* for token-level validation against the Python reference (test-outetts-prompt.mjs).
|
|
2544
|
+
*/
|
|
2545
|
+
function buildOutePromptString(text, speaker) {
|
|
2546
|
+
const normalized = normalizeOuteText(text);
|
|
2547
|
+
let promptText = normalized;
|
|
2548
|
+
let codes = "";
|
|
2549
|
+
if (speaker) {
|
|
2550
|
+
const { merged, sep } = mergeSpeakerText(normalized, speaker.text);
|
|
2551
|
+
promptText = merged;
|
|
2552
|
+
codes = createSpeakerCodes(speaker.words.map((w, i) => i === speaker.words.length - 1 ? {
|
|
2553
|
+
...w,
|
|
2554
|
+
word: w.word + sep
|
|
2555
|
+
} : w));
|
|
2556
|
+
}
|
|
2557
|
+
let prompt = `<|im_start|>\n<|text_start|>${promptText}<|text_end|>\n<|audio_start|>\n`;
|
|
2558
|
+
if (speaker) prompt += `${codes}\n<|word_start|>`;
|
|
2559
|
+
return prompt;
|
|
2560
|
+
}
|
|
2561
|
+
/**
|
|
2562
|
+
* Load an OuteTTS preset speaker (the in-context reference voice). Accepts a full URL
|
|
2563
|
+
* (mirror), a local file path (Node), or a bare preset name resolved against
|
|
2564
|
+
* `baseURL/{name}.json`. The returned object is the OuteSpeaker the prompt builder
|
|
2565
|
+
* embeds (transcript + per-word c1/c2 reference codes + prosody features).
|
|
2566
|
+
*/
|
|
2567
|
+
async function loadOuteSpeaker(voiceOrUrl, baseURL) {
|
|
2568
|
+
if (HTTP_URL_RE.test(voiceOrUrl)) {
|
|
2569
|
+
const resp$1 = await fetch(voiceOrUrl);
|
|
2570
|
+
if (!resp$1.ok) throw new Error(`Failed to fetch OuteTTS speaker ${voiceOrUrl}: ${resp$1.status}`);
|
|
2571
|
+
return await resp$1.json();
|
|
2572
|
+
}
|
|
2573
|
+
if (voiceOrUrl.includes("/") || voiceOrUrl.endsWith(".json")) {
|
|
2574
|
+
const fs = await import("node:fs").catch(() => null);
|
|
2575
|
+
if (fs) return JSON.parse(fs.readFileSync(voiceOrUrl, "utf8"));
|
|
2576
|
+
}
|
|
2577
|
+
if (!baseURL) throw new Error(`loadOuteSpeaker: "${voiceOrUrl}" is a preset name but no baseURL was provided.`);
|
|
2578
|
+
const url = `${baseURL.replace(TRAILING_SLASH_RE, "")}/${voiceOrUrl}.json`;
|
|
2579
|
+
const resp = await fetch(url);
|
|
2580
|
+
if (!resp.ok) throw new Error(`Failed to fetch OuteTTS preset "${voiceOrUrl}" from ${url}: ${resp.status}`);
|
|
2581
|
+
return await resp.json();
|
|
2582
|
+
}
|
|
2583
|
+
/** Build a fresh map of just the constant weights a graph references (see kani-tts). */
|
|
2584
|
+
function selectGraphWeights(graph, weights) {
|
|
2585
|
+
const out = /* @__PURE__ */ new Map();
|
|
2586
|
+
for (const [name, desc] of Object.entries(graph.tensors)) {
|
|
2587
|
+
if (desc.storage !== "constant") continue;
|
|
2588
|
+
const w = weights.get(name) ?? (desc.safetensorsKey ? weights.get(desc.safetensorsKey) : void 0);
|
|
2589
|
+
if (w) out.set(name, w);
|
|
2590
|
+
}
|
|
2591
|
+
return out;
|
|
2592
|
+
}
|
|
2593
|
+
/** Index of the maximum logit (greedy / argmax decode). */
|
|
2594
|
+
function argmax$1(logits) {
|
|
2595
|
+
let best = 0;
|
|
2596
|
+
let bestV = logits[0];
|
|
2597
|
+
for (let i = 1; i < logits.length; i++) if (logits[i] > bestV) {
|
|
2598
|
+
bestV = logits[i];
|
|
2599
|
+
best = i;
|
|
2600
|
+
}
|
|
2601
|
+
return best;
|
|
2602
|
+
}
|
|
2603
|
+
/** In-place softmax of `scores`. */
|
|
2604
|
+
function softmaxInPlace(scores) {
|
|
2605
|
+
let maxV = Number.NEGATIVE_INFINITY;
|
|
2606
|
+
for (const v of scores) if (v > maxV) maxV = v;
|
|
2607
|
+
let sum = 0;
|
|
2608
|
+
for (let i = 0; i < scores.length; i++) {
|
|
2609
|
+
const e = Math.exp(scores[i] - maxV);
|
|
2610
|
+
scores[i] = e;
|
|
2611
|
+
sum += e;
|
|
2612
|
+
}
|
|
2613
|
+
for (let i = 0; i < scores.length; i++) scores[i] /= sum;
|
|
2614
|
+
}
|
|
2615
|
+
/** In-place top-k filter: keep the k highest, zero the rest, renormalize. */
|
|
2616
|
+
function applyTopK(scores, k) {
|
|
2617
|
+
const n = scores.length;
|
|
2618
|
+
if (k <= 0 || k >= n) return;
|
|
2619
|
+
const order = Array.from({ length: n }, (_, i) => i).sort((a, b) => scores[b] - scores[a]);
|
|
2620
|
+
let keepSum = 0;
|
|
2621
|
+
for (let r = 0; r < k; r++) keepSum += scores[order[r]];
|
|
2622
|
+
for (let r = k; r < n; r++) scores[order[r]] = 0;
|
|
2623
|
+
if (keepSum > 0) for (let i = 0; i < n; i++) scores[i] /= keepSum;
|
|
2624
|
+
}
|
|
2625
|
+
/** Top-p nucleus sample from a softmaxed prob array over `ids`. */
|
|
2626
|
+
function nucleusSample(probs, ids, topP) {
|
|
2627
|
+
const n = probs.length;
|
|
2628
|
+
const order = Array.from({ length: n }, (_, i) => i).sort((a, b) => probs[b] - probs[a]);
|
|
2629
|
+
let cum = 0;
|
|
2630
|
+
let cutoff = n;
|
|
2631
|
+
for (let r = 0; r < n; r++) {
|
|
2632
|
+
cum += probs[order[r]];
|
|
2633
|
+
if (cum >= topP) {
|
|
2634
|
+
cutoff = r + 1;
|
|
2635
|
+
break;
|
|
2636
|
+
}
|
|
2637
|
+
}
|
|
2638
|
+
let keepSum = 0;
|
|
2639
|
+
for (let r = 0; r < cutoff; r++) keepSum += probs[order[r]];
|
|
2640
|
+
let pick = Math.random() * keepSum;
|
|
2641
|
+
for (let r = 0; r < cutoff; r++) {
|
|
2642
|
+
pick -= probs[order[r]];
|
|
2643
|
+
if (pick <= 0) return ids[order[r]];
|
|
2644
|
+
}
|
|
2645
|
+
return ids[order[0]];
|
|
2646
|
+
}
|
|
2647
|
+
/**
|
|
2648
|
+
* Resample interleaved-or-mono PCM to mono at `targetRate` via linear interpolation.
|
|
2649
|
+
* Audio is assumed already mono (Float32 [-1,1]); if `srcRate === targetRate` the
|
|
2650
|
+
* input is returned as-is. Linear resampling is sufficient for codec conditioning
|
|
2651
|
+
* (the DAC encoder is robust to the small interpolation error vs torchaudio's
|
|
2652
|
+
* sinc resampler — the cloned codes are stable).
|
|
2653
|
+
*/
|
|
2654
|
+
function resampleMono(audio, srcRate, targetRate) {
|
|
2655
|
+
if (srcRate === targetRate) return audio;
|
|
2656
|
+
const ratio = targetRate / srcRate;
|
|
2657
|
+
const outLen = Math.max(1, Math.round(audio.length * ratio));
|
|
2658
|
+
const out = new Float32Array(outLen);
|
|
2659
|
+
for (let i = 0; i < outLen; i++) {
|
|
2660
|
+
const srcPos = i / ratio;
|
|
2661
|
+
const i0 = Math.floor(srcPos);
|
|
2662
|
+
const i1 = Math.min(i0 + 1, audio.length - 1);
|
|
2663
|
+
const frac = srcPos - i0;
|
|
2664
|
+
out[i] = audio[i0] * (1 - frac) + audio[i1] * frac;
|
|
2665
|
+
}
|
|
2666
|
+
return out;
|
|
2667
|
+
}
|
|
2668
|
+
var OuteTTS = class OuteTTS {
|
|
2669
|
+
ctx;
|
|
2670
|
+
loaded;
|
|
2671
|
+
tokenizer;
|
|
2672
|
+
maxSeqLen;
|
|
838
2673
|
backboneExec;
|
|
839
|
-
|
|
840
|
-
attnLayers;
|
|
841
|
-
layerAlpha;
|
|
842
|
-
headDim;
|
|
843
|
-
ropeBase;
|
|
2674
|
+
speaker;
|
|
844
2675
|
_destroyed = false;
|
|
845
|
-
|
|
846
|
-
voiceTags;
|
|
847
|
-
/** The loaded checkpoint's architecture string ("KaniTTS2ForCausalLM" or "Lfm2ForCausalLM"). */
|
|
848
|
-
architecture;
|
|
2676
|
+
sampleRate = DAC_SAMPLE_RATE;
|
|
849
2677
|
/** The selected backbone precision ("q4" or "f32"). */
|
|
850
2678
|
dtype;
|
|
851
|
-
constructor(ctx, loaded, maxSeqLen, dtype) {
|
|
2679
|
+
constructor(ctx, loaded, maxSeqLen, dtype, speaker) {
|
|
852
2680
|
this.ctx = ctx;
|
|
853
2681
|
this.loaded = loaded;
|
|
854
2682
|
this.tokenizer = loaded.tokenizer;
|
|
855
|
-
this.rawConfig = loaded.rawConfig;
|
|
856
|
-
this.cfg = parseKaniConfig(loaded.rawConfig);
|
|
857
2683
|
this.maxSeqLen = maxSeqLen;
|
|
858
|
-
this.
|
|
859
|
-
this.
|
|
860
|
-
const
|
|
861
|
-
|
|
862
|
-
|
|
863
|
-
const heads = this.rawConfig.num_attention_heads;
|
|
864
|
-
this.headDim = this.rawConfig.head_dim ?? Math.floor(hidden / heads);
|
|
865
|
-
this.ropeBase = this.rawConfig.rope_theta ?? 1e6;
|
|
866
|
-
this.attnLayers = kaniAttentionLayerIndices(this.rawConfig);
|
|
867
|
-
this.layerAlpha = /* @__PURE__ */ new Map();
|
|
868
|
-
const useLearnable = this.rawConfig.use_learnable_rope ?? false;
|
|
869
|
-
for (const layer of this.attnLayers) {
|
|
870
|
-
let alpha = 1;
|
|
871
|
-
if (useLearnable) {
|
|
872
|
-
const w = loaded.backboneWeights.get(`learnable_rope_layers.${layer}.alpha_weight`);
|
|
873
|
-
alpha = kaniLayerAlpha(w ? w.data[0] ?? 0 : 0, this.cfg);
|
|
874
|
-
}
|
|
875
|
-
this.layerAlpha.set(layer, alpha);
|
|
876
|
-
}
|
|
877
|
-
this.backboneExec = this.buildBackbone(ctx, loaded, maxSeqLen);
|
|
878
|
-
}
|
|
879
|
-
/**
|
|
880
|
-
* Build the codec-LM backbone graph, quantize it when q4, and upload its weights
|
|
881
|
-
* into a fresh Executor. For a pre-quantized (gerbil-int4) checkpoint the linear
|
|
882
|
-
* projections already carry `.q`/`.scales`/`.zeros`, so quantizeKaniBackbone is an
|
|
883
|
-
* idempotent no-op (it skips any projection whose `.q` is already present) and the
|
|
884
|
-
* int4 weights are uploaded verbatim — no on-device re-quantization.
|
|
885
|
-
*/
|
|
886
|
-
buildBackbone(ctx, loaded, maxSeqLen) {
|
|
887
|
-
const useQ4 = this.dtype === "q4";
|
|
888
|
-
const graph = generateKaniTtsGraph(this.rawConfig, useQ4 ? "q4" : void 0);
|
|
889
|
-
if (useQ4) quantizeKaniBackbone(graph, loaded.backboneWeights);
|
|
2684
|
+
this.speaker = speaker;
|
|
2685
|
+
this.dtype = dtype;
|
|
2686
|
+
const useQ4 = dtype === "q4";
|
|
2687
|
+
const graph = generateOuteTtsBackboneGraph(loaded.rawConfig, useQ4 ? "q4" : void 0);
|
|
2688
|
+
if (useQ4) quantizeBackboneInt4(graph, loaded.backboneWeights);
|
|
890
2689
|
const exec = new Executor(ctx, graph, {
|
|
891
2690
|
maxSeqLen,
|
|
892
2691
|
kvMode: "f32"
|
|
893
2692
|
});
|
|
894
2693
|
exec.uploadWeightsMap(selectGraphWeights(graph, loaded.backboneWeights));
|
|
895
2694
|
exec.initBindGroups();
|
|
896
|
-
|
|
2695
|
+
this.backboneExec = exec;
|
|
897
2696
|
}
|
|
898
|
-
static async create(options
|
|
2697
|
+
static async create(options) {
|
|
899
2698
|
const ctx = await initGPU();
|
|
900
|
-
const loaded = await
|
|
2699
|
+
const loaded = await loadOuteTTS({
|
|
901
2700
|
repo: options.repo,
|
|
902
|
-
|
|
2701
|
+
codecWeightsPath: options.codecWeightsPath,
|
|
903
2702
|
revision: options.revision,
|
|
904
2703
|
hfToken: options.hfToken,
|
|
905
2704
|
cacheDir: options.cacheDir,
|
|
906
2705
|
onProgress: options.onProgress
|
|
907
2706
|
});
|
|
908
|
-
const ctxLen = loaded.rawConfig.max_position_embeddings ??
|
|
909
|
-
return new
|
|
910
|
-
}
|
|
911
|
-
/** Write per-layer cos/sin for token rows [rowStart, rowStart+positions.length). */
|
|
912
|
-
writeCosSin(positions, rowStart) {
|
|
913
|
-
const rowBytes = this.headDim * 4;
|
|
914
|
-
for (const layer of this.attnLayers) {
|
|
915
|
-
const alpha = this.layerAlpha.get(layer) ?? 1;
|
|
916
|
-
const { cos, sin } = buildKaniLayerCosSin(positions, this.headDim, this.ropeBase, alpha);
|
|
917
|
-
if (rowStart === 0) {
|
|
918
|
-
this.backboneExec.writeInput(kaniCosTensor(layer), cos);
|
|
919
|
-
this.backboneExec.writeInput(kaniSinTensor(layer), sin);
|
|
920
|
-
} else {
|
|
921
|
-
this.backboneExec.writeInputAt(kaniCosTensor(layer), cos, rowStart * rowBytes);
|
|
922
|
-
this.backboneExec.writeInputAt(kaniSinTensor(layer), sin, rowStart * rowBytes);
|
|
923
|
-
}
|
|
924
|
-
}
|
|
2707
|
+
const ctxLen = loaded.rawConfig.max_position_embeddings ?? 40960;
|
|
2708
|
+
return new OuteTTS(ctx, loaded, Math.min(options.maxSeqLen ?? 4096, ctxLen), options.dtype ?? "q4", options.speaker);
|
|
925
2709
|
}
|
|
926
2710
|
/**
|
|
927
|
-
*
|
|
928
|
-
*
|
|
929
|
-
* the
|
|
930
|
-
*
|
|
931
|
-
*
|
|
2711
|
+
* Build the synthesis prompt as a STRING (outetts v3 get_completion_prompt) and
|
|
2712
|
+
* tokenize in one pass. With a speaker, the speaker transcript is merged in front of
|
|
2713
|
+
* the target text (merge_speaker_text), the speaker's reference word blocks
|
|
2714
|
+
* (word + <|features|> + <|t_dur|> + energy/spectral_centroid/pitch + <|code|> +
|
|
2715
|
+
* c1/c2 pairs) are embedded in-context, and the prompt ends with a dangling
|
|
2716
|
+
* <|word_start|> to cue generation. Without a speaker, a plain text-only prompt.
|
|
2717
|
+
*
|
|
2718
|
+
* <|im_start|>\n<|text_start|>{text}<|text_end|>\n<|audio_start|>\n
|
|
2719
|
+
* [ {speaker word blocks "\n"-joined} \n <|word_start|> ]
|
|
2720
|
+
*
|
|
2721
|
+
* This reproduces the reference token layout EXACTLY (validated to the byte vs
|
|
2722
|
+
* outetts.PromptProcessor — see scripts/engine/outetts-ref-prompt.py and
|
|
2723
|
+
* test-outetts-prompt.mjs). The per-word time/feature tokens were the gating fix:
|
|
2724
|
+
* the earlier foundation dropped them and the backbone produced incoherent audio.
|
|
932
2725
|
*/
|
|
933
|
-
|
|
934
|
-
return
|
|
935
|
-
}
|
|
936
|
-
/** NanoCodec's fixed output sample rate (Hz). Not selectable — see SpeakResult. */
|
|
937
|
-
get sampleRate() {
|
|
938
|
-
return KANI_SAMPLE_RATE;
|
|
2726
|
+
buildPrompt(text, speaker) {
|
|
2727
|
+
return this.tokenizer.encode(buildOutePromptString(text, speaker ?? this.speaker));
|
|
939
2728
|
}
|
|
940
2729
|
/**
|
|
941
|
-
* Synthesize speech for `text`. Returns
|
|
2730
|
+
* Synthesize speech for `text`. Returns 24 kHz mono PCM.
|
|
942
2731
|
*
|
|
943
|
-
* Pipeline: build
|
|
944
|
-
*
|
|
945
|
-
* →
|
|
2732
|
+
* Pipeline: build prompt → prefill backbone → AR-decode interleaved c1/c2 audio
|
|
2733
|
+
* tokens until <|audio_end|> → split to a [2, T] DAC code grid → DAC decode → tanh
|
|
2734
|
+
* → PCM.
|
|
946
2735
|
*/
|
|
947
2736
|
async speak(text, opts = {}) {
|
|
948
|
-
if (this._destroyed) throw new Error("
|
|
949
|
-
const temperature = opts.temperature ?? 1;
|
|
950
|
-
const topP = opts.topP ?? .95;
|
|
951
|
-
const topK = opts.topK ?? 0;
|
|
952
|
-
const repetitionPenalty = opts.repetitionPenalty ?? 1.1;
|
|
953
|
-
const maxNewTokens = Math.min(opts.maxNewTokens ?? DEFAULT_MAX_NEW_TOKENS, this.maxSeqLen);
|
|
954
|
-
const promptText = this.voiceTags.length > 0 ? `${(opts.languageTag ?? this.voiceTags[0]).trim()}: ${text}` : text;
|
|
955
|
-
const prompt = [
|
|
956
|
-
KANI_START_OF_HUMAN,
|
|
957
|
-
...this.tokenizer.encode(promptText),
|
|
958
|
-
KANI_END_OF_TEXT,
|
|
959
|
-
KANI_END_OF_HUMAN
|
|
960
|
-
];
|
|
961
|
-
if (prompt.length >= this.maxSeqLen) throw new Error(`Kani prompt is ${prompt.length} tokens >= maxSeqLen ${this.maxSeqLen}.`);
|
|
962
|
-
this.backboneExec.reset();
|
|
963
|
-
const promptPos = computeKaniPositions(prompt, this.cfg);
|
|
964
|
-
this.writeCosSin(promptPos, 0);
|
|
965
|
-
const { logits } = await this.backboneExec.forward(new Uint32Array(prompt));
|
|
2737
|
+
if (this._destroyed) throw new Error("OuteTTS has been destroyed.");
|
|
966
2738
|
const startTime = performance.now();
|
|
967
|
-
const
|
|
968
|
-
|
|
969
|
-
|
|
970
|
-
topK,
|
|
971
|
-
repetitionPenalty,
|
|
972
|
-
maxNewTokens,
|
|
973
|
-
maxFrames: opts.maxFrames ?? Number.POSITIVE_INFINITY
|
|
974
|
-
});
|
|
975
|
-
const usable = audioTokens.length - audioTokens.length % this.cfg.tokensPerFrame;
|
|
976
|
-
if (usable === 0) throw new Error(`Kani produced no complete audio frames (finished=${finished}, audioTokens=${audioTokens.length}). The codec-LM emitted no speech.`);
|
|
977
|
-
const { codes, numFrames } = audioTokensToCodes(audioTokens.slice(0, usable), this.cfg);
|
|
2739
|
+
const audioTokens = await this.generateTokens(text, opts);
|
|
2740
|
+
const { codes, numFrames } = audioTokensToDacCodes(audioTokens);
|
|
2741
|
+
if (numFrames === 0) throw new Error(`OuteTTS produced no complete audio frames (audioTokens=${audioTokens.length}).`);
|
|
978
2742
|
const pcm = await this.decodeCodes(codes, numFrames);
|
|
979
2743
|
const totalTime = (performance.now() - startTime) / 1e3;
|
|
980
|
-
const audioSeconds =
|
|
981
|
-
console.log(`[
|
|
2744
|
+
const audioSeconds = pcm.length / DAC_SAMPLE_RATE;
|
|
2745
|
+
console.log(`[outetts] frames=${numFrames} audio=${audioSeconds.toFixed(2)}s wall=${totalTime.toFixed(2)}s RTF=${(audioSeconds / totalTime).toFixed(2)}x`);
|
|
982
2746
|
return {
|
|
983
2747
|
pcm,
|
|
984
|
-
sampleRate:
|
|
2748
|
+
sampleRate: DAC_SAMPLE_RATE,
|
|
985
2749
|
frames: numFrames,
|
|
986
2750
|
audioSeconds
|
|
987
2751
|
};
|
|
988
2752
|
}
|
|
989
2753
|
/**
|
|
990
|
-
*
|
|
991
|
-
*
|
|
992
|
-
*
|
|
993
|
-
*
|
|
994
|
-
*
|
|
2754
|
+
* Build an in-context speaker reference for voice cloning.
|
|
2755
|
+
*
|
|
2756
|
+
* Two inputs are supported:
|
|
2757
|
+
* 1. A pre-encoded {@link OuteSpeaker} (transcript + per-word c1/c2 DAC codes +
|
|
2758
|
+
* prosody features — the preset-JSON shape). Returned verbatim.
|
|
2759
|
+
* 2. Raw audio (Float32 PCM + transcript). The clip is resampled to 24 kHz mono,
|
|
2760
|
+
* hop-padded, run through the DAC.speech ENCODER + RVQ-encode on the GPU
|
|
2761
|
+
* (generateDacSpeechEncoderGraph; encoder latent validated bit-exact and codes
|
|
2762
|
+
* 450/450 vs torch DAC.encode — scripts/engine/test-outetts-encode.mjs), then
|
|
2763
|
+
* split into per-word {@link OuteSpeakerWord} blocks with prosody features.
|
|
2764
|
+
*
|
|
2765
|
+
* Fidelity vs outetts `create_speaker`: the DAC-encode (audio → c1/c2 codes) and the
|
|
2766
|
+
* energy/spectral_centroid/pitch feature math are a faithful port. Two documented
|
|
2767
|
+
* approximations: (a) word→frame ALIGNMENT — outetts uses Whisper word timestamps to
|
|
2768
|
+
* slice the code stream; with no on-device Whisper we distribute the encoded frames
|
|
2769
|
+
* across the transcript words proportionally to word length (see
|
|
2770
|
+
* {@link alignWordsProportional}); (b) we skip outetts' pyloudnorm pre-normalization
|
|
2771
|
+
* (it mainly rescales absolute level / the energy bucket). The cloned voice timbre
|
|
2772
|
+
* and prosody come through; the per-word boundaries are heuristic.
|
|
2773
|
+
*/
|
|
2774
|
+
async encodeSpeaker(ref) {
|
|
2775
|
+
if ("words" in ref && "text" in ref) return ref;
|
|
2776
|
+
if (this._destroyed) throw new Error("OuteTTS has been destroyed.");
|
|
2777
|
+
const { audio, sampleRate, transcript } = ref;
|
|
2778
|
+
if (!audio || audio.length === 0) throw new Error("OuteTTS.encodeSpeaker: empty audio.");
|
|
2779
|
+
if (!transcript || transcript.trim().length === 0) throw new Error("OuteTTS.encodeSpeaker: a transcript of the clip is required.");
|
|
2780
|
+
const mono = resampleMono(audio, sampleRate, DAC_SAMPLE_RATE);
|
|
2781
|
+
const paddedLen = Math.ceil(mono.length / DAC_HOP) * DAC_HOP;
|
|
2782
|
+
const pcm = new Float32Array(paddedLen);
|
|
2783
|
+
pcm.set(mono);
|
|
2784
|
+
const numFrames = dacEncoderFrames(paddedLen);
|
|
2785
|
+
const codes = await this.encodeAudioToCodes(pcm, numFrames);
|
|
2786
|
+
const c1All = codes.subarray(0, numFrames);
|
|
2787
|
+
const c2All = codes.subarray(numFrames, 2 * numFrames);
|
|
2788
|
+
const words = alignWordsProportional(transcript, numFrames).map((a) => {
|
|
2789
|
+
const startSample = Math.min(a.startFrame * DAC_HOP, mono.length);
|
|
2790
|
+
const endSample = Math.min(a.endFrame * DAC_HOP, mono.length);
|
|
2791
|
+
const features = extractAudioFeatures(mono.subarray(startSample, endSample), DAC_SAMPLE_RATE);
|
|
2792
|
+
return {
|
|
2793
|
+
word: a.word,
|
|
2794
|
+
duration: Math.round((a.endFrame - a.startFrame) / DAC_TPS * 100) / 100,
|
|
2795
|
+
c1: Array.from(c1All.subarray(a.startFrame, a.endFrame)),
|
|
2796
|
+
c2: Array.from(c2All.subarray(a.startFrame, a.endFrame)),
|
|
2797
|
+
features
|
|
2798
|
+
};
|
|
2799
|
+
});
|
|
2800
|
+
return {
|
|
2801
|
+
text: transcript.trim(),
|
|
2802
|
+
words
|
|
2803
|
+
};
|
|
2804
|
+
}
|
|
2805
|
+
/** Run the DAC.speech encoder + RVQ-encode graph → codebook-major codes[NCB·T] (u32). */
|
|
2806
|
+
async encodeAudioToCodes(pcm, numFrames) {
|
|
2807
|
+
const graph = generateDacSpeechEncoderGraph({ numSamples: pcm.length });
|
|
2808
|
+
const exec = new Executor(this.ctx, graph, {
|
|
2809
|
+
maxSeqLen: numFrames,
|
|
2810
|
+
kvMode: "f32"
|
|
2811
|
+
});
|
|
2812
|
+
try {
|
|
2813
|
+
exec.uploadWeightsMap(selectGraphWeights(graph, this.loaded.codecWeights));
|
|
2814
|
+
exec.initBindGroups();
|
|
2815
|
+
exec.reset();
|
|
2816
|
+
exec.writeInput("pcm_in", pcm);
|
|
2817
|
+
const raw = await exec.runGraphOutput("dac_codes", DAC_NUM_CODEBOOKS * numFrames);
|
|
2818
|
+
return new Uint32Array(raw.buffer, raw.byteOffset, DAC_NUM_CODEBOOKS * numFrames).slice();
|
|
2819
|
+
} finally {
|
|
2820
|
+
exec.destroy();
|
|
2821
|
+
}
|
|
2822
|
+
}
|
|
2823
|
+
/** Build the prompt, prefill, and AR-decode the raw generated token stream. */
|
|
2824
|
+
async generateTokens(text, opts) {
|
|
2825
|
+
const prompt = this.buildPrompt(text, opts.speaker);
|
|
2826
|
+
if (prompt.length >= this.maxSeqLen) throw new Error(`OuteTTS prompt is ${prompt.length} tokens >= maxSeqLen ${this.maxSeqLen}.`);
|
|
2827
|
+
this.backboneExec.reset();
|
|
2828
|
+
const { logits } = await this.backboneExec.forward(new Uint32Array(prompt));
|
|
2829
|
+
return this.runDecodeLoop(prompt, logits, {
|
|
2830
|
+
temperature: opts.temperature ?? .4,
|
|
2831
|
+
topP: opts.topP ?? .9,
|
|
2832
|
+
topK: opts.topK ?? 40,
|
|
2833
|
+
repetitionPenalty: opts.repetitionPenalty ?? 1.1,
|
|
2834
|
+
repetitionRange: opts.repetitionRange ?? 64,
|
|
2835
|
+
maxNewTokens: Math.min(opts.maxNewTokens ?? DEFAULT_MAX_NEW_TOKENS, this.maxSeqLen),
|
|
2836
|
+
greedy: opts.greedy ?? false
|
|
2837
|
+
});
|
|
2838
|
+
}
|
|
2839
|
+
/** Test hook: return the raw generated token stream (no DAC decode). */
|
|
2840
|
+
async generateTokensForTest(text, opts = {}) {
|
|
2841
|
+
return this.generateTokens(text, opts);
|
|
2842
|
+
}
|
|
2843
|
+
/**
|
|
2844
|
+
* Autoregressive decode: from the prefill logits, sample one token per step over the
|
|
2845
|
+
* full vocab (greedy argmax, or temperature/top-k/top-p with repetition penalty).
|
|
2846
|
+
* Stops at <|audio_end|>/<|im_end|> or the token cap. Returns the raw token stream
|
|
2847
|
+
* (word-block structure + interleaved c1/c2 codes); audioTokensToDacCodes extracts
|
|
2848
|
+
* the codes.
|
|
995
2849
|
*/
|
|
996
2850
|
async runDecodeLoop(prompt, prefillLogits, p) {
|
|
997
2851
|
const seq = [...prompt];
|
|
998
|
-
const
|
|
999
|
-
const frameCap = p.maxFrames * this.cfg.tokensPerFrame;
|
|
2852
|
+
const generated = [];
|
|
1000
2853
|
let logits = prefillLogits;
|
|
1001
|
-
let inSpeech = false;
|
|
1002
|
-
let finished = false;
|
|
1003
2854
|
for (let step = 0; step < p.maxNewTokens; step++) {
|
|
1004
|
-
const
|
|
2855
|
+
const next = p.greedy ? argmax$1(logits) : this.sampleToken(logits, {
|
|
1005
2856
|
temperature: p.temperature,
|
|
1006
2857
|
topP: p.topP,
|
|
1007
2858
|
topK: p.topK,
|
|
1008
2859
|
repetitionPenalty: p.repetitionPenalty,
|
|
1009
|
-
previous: seq
|
|
1010
|
-
})
|
|
1011
|
-
if (
|
|
1012
|
-
|
|
1013
|
-
|
|
1014
|
-
finished = true;
|
|
1015
|
-
break;
|
|
1016
|
-
} else if (inSpeech && nextToken >= this.cfg.audioTokensStart) audioTokens.push(nextToken);
|
|
1017
|
-
seq.push(nextToken);
|
|
1018
|
-
if (inSpeech && audioTokens.length >= frameCap) {
|
|
1019
|
-
finished = true;
|
|
1020
|
-
break;
|
|
1021
|
-
}
|
|
2860
|
+
previous: seq.slice(-p.repetitionRange)
|
|
2861
|
+
});
|
|
2862
|
+
if (next === OUTETTS_AUDIO_END || next === OUTETTS_IM_END) break;
|
|
2863
|
+
generated.push(next);
|
|
2864
|
+
seq.push(next);
|
|
1022
2865
|
if (seq.length >= this.maxSeqLen) break;
|
|
1023
|
-
this.
|
|
1024
|
-
logits = (await this.backboneExec.forward(new Uint32Array([nextToken]))).logits;
|
|
1025
|
-
}
|
|
1026
|
-
return {
|
|
1027
|
-
audioTokens,
|
|
1028
|
-
finished
|
|
1029
|
-
};
|
|
1030
|
-
}
|
|
1031
|
-
/** Logical (frame-level) position of the LAST token in `seq`. */
|
|
1032
|
-
logicalPositionAt(seq) {
|
|
1033
|
-
const positions = computeKaniPositions(seq, this.cfg);
|
|
1034
|
-
return positions[positions.length - 1];
|
|
1035
|
-
}
|
|
1036
|
-
/** Greedy argmax over a logits row. */
|
|
1037
|
-
argmax(logits) {
|
|
1038
|
-
let best = 0;
|
|
1039
|
-
let bestVal = logits[0];
|
|
1040
|
-
for (let i = 1; i < logits.length; i++) if (logits[i] > bestVal) {
|
|
1041
|
-
bestVal = logits[i];
|
|
1042
|
-
best = i;
|
|
2866
|
+
logits = (await this.backboneExec.forward(new Uint32Array([next]))).logits;
|
|
1043
2867
|
}
|
|
1044
|
-
return
|
|
2868
|
+
return generated;
|
|
1045
2869
|
}
|
|
1046
2870
|
/**
|
|
1047
|
-
* Sample one
|
|
1048
|
-
*
|
|
1049
|
-
*
|
|
1050
|
-
* optional top-k, then top-p nucleus sampling.
|
|
2871
|
+
* Sample one token over the FULL vocab: repetition penalty over the recent window,
|
|
2872
|
+
* then temperature → top-k → top-p nucleus. (OuteTTS emits structural word-block
|
|
2873
|
+
* tokens interleaved with audio codes, so the sampler must not be window-restricted.)
|
|
1051
2874
|
*/
|
|
1052
|
-
|
|
1053
|
-
const
|
|
1054
|
-
const winLen = this.cfg.codebookSize;
|
|
1055
|
-
const allowEos = codebook === 0;
|
|
1056
|
-
const n = winLen + (allowEos ? 1 : 0);
|
|
2875
|
+
sampleToken(logits, p) {
|
|
2876
|
+
const n = logits.length;
|
|
1057
2877
|
const ids = new Int32Array(n);
|
|
1058
|
-
for (let i = 0; i < winLen; i++) ids[i] = winStart + i;
|
|
1059
|
-
if (allowEos) ids[winLen] = this.cfg.endOfSpeech;
|
|
1060
|
-
const prev = p.repetitionPenalty !== 1 ? new Set(p.previous) : null;
|
|
1061
2878
|
const scores = new Float32Array(n);
|
|
2879
|
+
const prev = p.repetitionPenalty !== 1 ? new Set(p.previous) : null;
|
|
1062
2880
|
for (let i = 0; i < n; i++) {
|
|
1063
|
-
|
|
1064
|
-
|
|
2881
|
+
ids[i] = i;
|
|
2882
|
+
let s = logits[i];
|
|
2883
|
+
if (prev?.has(i)) s = s > 0 ? s / p.repetitionPenalty : s * p.repetitionPenalty;
|
|
1065
2884
|
scores[i] = s / p.temperature;
|
|
1066
2885
|
}
|
|
1067
2886
|
softmaxInPlace(scores);
|
|
@@ -1069,47 +2888,29 @@ var KaniTTS = class KaniTTS {
|
|
|
1069
2888
|
return nucleusSample(scores, ids, p.topP);
|
|
1070
2889
|
}
|
|
1071
2890
|
/**
|
|
1072
|
-
* Decode
|
|
1073
|
-
*
|
|
1074
|
-
*
|
|
1075
|
-
*
|
|
1076
|
-
*
|
|
1077
|
-
*
|
|
1078
|
-
* samples — numerically identical to a single decode, but bounded per dispatch.
|
|
2891
|
+
* Decode the [2, T] DAC code grid (codebook-major) → PCM, applying the final tanh.
|
|
2892
|
+
* The DAC decoder is non-causal (symmetric padding) and its odd rate-5 transpose
|
|
2893
|
+
* makes interior windows position-dependent, so the whole clip is decoded in ONE
|
|
2894
|
+
* pass (no chunking). The codec conv/Snake/Add kernels use 2D dispatch (the
|
|
2895
|
+
* dispatch2D helper) so even multi-second clips stay under WebGPU's 65535 per-dim
|
|
2896
|
+
* workgroup cap.
|
|
1079
2897
|
*/
|
|
1080
2898
|
async decodeCodes(codes, numFrames) {
|
|
1081
|
-
const
|
|
1082
|
-
const
|
|
1083
|
-
const chunk = KANI_CODEC_CHUNK_FRAMES;
|
|
1084
|
-
const pcm = new Float32Array(numFrames * KANI_HOP);
|
|
1085
|
-
for (let start = 0; start < numFrames; start += chunk) {
|
|
1086
|
-
const ctxStart = Math.max(0, start - lookback);
|
|
1087
|
-
const ctxFrames = start - ctxStart;
|
|
1088
|
-
const end = Math.min(numFrames, start + chunk);
|
|
1089
|
-
const winFrames = end - ctxStart;
|
|
1090
|
-
const winCodes = new Uint32Array(groups * winFrames);
|
|
1091
|
-
for (let g = 0; g < groups; g++) for (let t = 0; t < winFrames; t++) winCodes[g * winFrames + t] = codes[g * numFrames + (ctxStart + t)];
|
|
1092
|
-
const win = await this.decodeCodesWindow(winCodes, winFrames);
|
|
1093
|
-
const dropSamples = ctxFrames * KANI_HOP;
|
|
1094
|
-
const keepSamples = (end - start) * KANI_HOP;
|
|
1095
|
-
pcm.set(win.subarray(dropSamples, dropSamples + keepSamples), start * KANI_HOP);
|
|
1096
|
-
}
|
|
1097
|
-
for (let i = 0; i < pcm.length; i++) pcm[i] = Math.min(1, Math.max(-1, pcm[i]));
|
|
1098
|
-
return pcm;
|
|
1099
|
-
}
|
|
1100
|
-
/** Run the NanoCodec decoder graph for a single (bounded) code window → PCM. */
|
|
1101
|
-
async decodeCodesWindow(winCodes, winFrames) {
|
|
1102
|
-
const graph = generateNanoCodecDecoderGraph({ numFrames: winFrames });
|
|
2899
|
+
const graph = generateDacSpeechDecoderGraph({ numFrames });
|
|
2900
|
+
const pcmLen = dacOutputLength(numFrames);
|
|
1103
2901
|
const exec = new Executor(this.ctx, graph, {
|
|
1104
|
-
maxSeqLen:
|
|
2902
|
+
maxSeqLen: numFrames,
|
|
1105
2903
|
kvMode: "f32"
|
|
1106
2904
|
});
|
|
1107
2905
|
try {
|
|
1108
2906
|
exec.uploadWeightsMap(selectGraphWeights(graph, this.loaded.codecWeights));
|
|
1109
2907
|
exec.initBindGroups();
|
|
1110
2908
|
exec.reset();
|
|
1111
|
-
exec.writeInput("audio_codes",
|
|
1112
|
-
|
|
2909
|
+
exec.writeInput("audio_codes", codes);
|
|
2910
|
+
const raw = await exec.runGraphOutput(graph.outputs[0], pcmLen);
|
|
2911
|
+
const pcm = new Float32Array(raw.length);
|
|
2912
|
+
for (let i = 0; i < raw.length; i++) pcm[i] = Math.tanh(raw[i]);
|
|
2913
|
+
return pcm;
|
|
1113
2914
|
} finally {
|
|
1114
2915
|
exec.destroy();
|
|
1115
2916
|
}
|
|
@@ -2476,6 +4277,10 @@ var WebGPUEngine = class WebGPUEngine {
|
|
|
2476
4277
|
_createOptions;
|
|
2477
4278
|
/** Lazily-created Kani-TTS engine (codec-LM + NanoCodec) backing speak(). */
|
|
2478
4279
|
_kaniTTS = null;
|
|
4280
|
+
/** Lazily-created OuteTTS engine (Qwen3-0.6B codec-LM + DAC.speech) backing speak(). */
|
|
4281
|
+
_outeTTS = null;
|
|
4282
|
+
/** The preset voice the lazily-built OuteTTS engine was created with. */
|
|
4283
|
+
_outeVoice = null;
|
|
2479
4284
|
/**
|
|
2480
4285
|
* WebKit group-size probe state. When true, a candidate group size is being
|
|
2481
4286
|
* tried this page-load and must be promoted (or capped) after the FIRST
|
|
@@ -3140,7 +4945,8 @@ var WebGPUEngine = class WebGPUEngine {
|
|
|
3140
4945
|
async speak(text, options = {}) {
|
|
3141
4946
|
this.checkDestroyed();
|
|
3142
4947
|
const repo = this._createOptions.repo ?? "";
|
|
3143
|
-
if (
|
|
4948
|
+
if (/outetts/i.test(repo) || /oute/i.test(this._architecture)) return this.speakOute(text, options);
|
|
4949
|
+
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}").`);
|
|
3144
4950
|
if (!this._kaniTTS) this._kaniTTS = await KaniTTS.create({
|
|
3145
4951
|
repo: this._createOptions.repo,
|
|
3146
4952
|
revision: this._createOptions.revision,
|
|
@@ -3158,8 +4964,69 @@ var WebGPUEngine = class WebGPUEngine {
|
|
|
3158
4964
|
maxFrames: options.maxFrames
|
|
3159
4965
|
});
|
|
3160
4966
|
}
|
|
4967
|
+
/**
|
|
4968
|
+
* OuteTTS speak path: lazily build the OuteTTS engine (loading the chosen preset
|
|
4969
|
+
* voice's speaker JSON + the folded DAC codec weights from the mirror), rebuilding
|
|
4970
|
+
* if the requested voice changed. Mirrors the Kani public shape (temperature/topP/topK).
|
|
4971
|
+
*/
|
|
4972
|
+
async speakOute(text, options) {
|
|
4973
|
+
await this.ensureOute(options);
|
|
4974
|
+
return this._outeTTS.speak(text, {
|
|
4975
|
+
temperature: options.temperature,
|
|
4976
|
+
topP: options.topP,
|
|
4977
|
+
topK: options.topK,
|
|
4978
|
+
repetitionPenalty: options.repetitionPenalty,
|
|
4979
|
+
speaker: options.cloneSpeaker
|
|
4980
|
+
});
|
|
4981
|
+
}
|
|
4982
|
+
/**
|
|
4983
|
+
* Lazily build (or reuse) the OuteTTS engine for the requested preset voice,
|
|
4984
|
+
* rebuilding if the voice changed. Loads the preset speaker JSON + folded DAC
|
|
4985
|
+
* codec weights from the mirror (or a locally-cached `{cacheDir}/{voice}.json`).
|
|
4986
|
+
* Shared by {@link speakOute} and {@link encodeSpeaker}.
|
|
4987
|
+
*/
|
|
4988
|
+
async ensureOute(options) {
|
|
4989
|
+
const voice = options.voice ?? OUTETTS_ASSETS.defaultVoice;
|
|
4990
|
+
if (this._outeTTS && this._outeVoice !== voice) {
|
|
4991
|
+
this._outeTTS.destroy();
|
|
4992
|
+
this._outeTTS = null;
|
|
4993
|
+
}
|
|
4994
|
+
if (this._outeTTS) return;
|
|
4995
|
+
const cacheDir = this._createOptions.cacheDir;
|
|
4996
|
+
let speakerRef = voice;
|
|
4997
|
+
if (cacheDir) {
|
|
4998
|
+
const local = `${cacheDir.replace(/\/$/, "")}/${voice}.json`;
|
|
4999
|
+
if ((await import("node:fs").catch(() => null))?.existsSync(local)) speakerRef = local;
|
|
5000
|
+
}
|
|
5001
|
+
const speaker = await loadOuteSpeaker(speakerRef, OUTETTS_ASSETS.speakersBaseURL);
|
|
5002
|
+
this._outeTTS = await OuteTTS.create({
|
|
5003
|
+
repo: this._createOptions.repo,
|
|
5004
|
+
revision: this._createOptions.revision,
|
|
5005
|
+
hfToken: this._createOptions.hfToken,
|
|
5006
|
+
cacheDir: this._createOptions.cacheDir,
|
|
5007
|
+
maxSeqLen: this.maxSeqLen,
|
|
5008
|
+
dtype: options.dtype ?? "q4",
|
|
5009
|
+
codecWeightsPath: options.codecWeightsPath ?? OUTETTS_ASSETS.codecWeightsURL,
|
|
5010
|
+
speaker
|
|
5011
|
+
});
|
|
5012
|
+
this._outeVoice = voice;
|
|
5013
|
+
}
|
|
5014
|
+
/**
|
|
5015
|
+
* OuteTTS voice cloning: encode a raw reference clip (mono PCM + its transcript)
|
|
5016
|
+
* into a reusable speaker (DAC.speech encode → per-word c1/c2 codes + prosody).
|
|
5017
|
+
* Pass the result to `speak(text, { cloneSpeaker })` to synthesize new text in that
|
|
5018
|
+
* voice. Requires an OuteTTS checkpoint (the folded codec carries the encoder).
|
|
5019
|
+
*/
|
|
5020
|
+
async encodeSpeaker(ref) {
|
|
5021
|
+
this.checkDestroyed();
|
|
5022
|
+
const repo = this._createOptions.repo ?? "";
|
|
5023
|
+
if (!(/outetts/i.test(repo) || /oute/i.test(this._architecture))) throw new Error(`encodeSpeaker() requires an OuteTTS checkpoint (e.g. "${DEFAULT_MODELS.ttsOute}"), loaded engine is "${this._architecture}" (repo "${repo}").`);
|
|
5024
|
+
await this.ensureOute({});
|
|
5025
|
+
return this._outeTTS.encodeSpeaker(ref);
|
|
5026
|
+
}
|
|
3161
5027
|
/** Accent/voice tags the lazily-built TTS engine supports ([] = single voice). */
|
|
3162
5028
|
ttsVoices() {
|
|
5029
|
+
if (/outetts/i.test(this._createOptions.repo ?? "") || /oute/i.test(this._architecture)) return [...OUTETTS_PRESET_VOICES];
|
|
3163
5030
|
return this._kaniTTS?.availableVoices() ?? [];
|
|
3164
5031
|
}
|
|
3165
5032
|
/**
|
|
@@ -3585,12 +5452,12 @@ var WebGPUEngine = class WebGPUEngine {
|
|
|
3585
5452
|
const f = data instanceof Float32Array ? data : new Float32Array(data.buffer);
|
|
3586
5453
|
let sum = 0;
|
|
3587
5454
|
let maxVal = -Infinity;
|
|
3588
|
-
let argmax$
|
|
5455
|
+
let argmax$2 = 0;
|
|
3589
5456
|
for (let i = 0; i < f.length; i++) {
|
|
3590
5457
|
sum += f[i];
|
|
3591
5458
|
if (f[i] > maxVal) {
|
|
3592
5459
|
maxVal = f[i];
|
|
3593
|
-
argmax$
|
|
5460
|
+
argmax$2 = i;
|
|
3594
5461
|
}
|
|
3595
5462
|
}
|
|
3596
5463
|
const entry = {
|
|
@@ -3603,7 +5470,7 @@ var WebGPUEngine = class WebGPUEngine {
|
|
|
3603
5470
|
Math.round(f[2] * 1e6) / 1e6,
|
|
3604
5471
|
Math.round(f[3] * 1e6) / 1e6
|
|
3605
5472
|
],
|
|
3606
|
-
argmax: argmax$
|
|
5473
|
+
argmax: argmax$2,
|
|
3607
5474
|
maxVal: Math.round(maxVal * 1e4) / 1e4
|
|
3608
5475
|
};
|
|
3609
5476
|
checks.push(entry);
|
|
@@ -3857,6 +5724,7 @@ var WebGPUEngine = class WebGPUEngine {
|
|
|
3857
5724
|
this._destroyed = true;
|
|
3858
5725
|
this.executor.destroy();
|
|
3859
5726
|
this._kaniTTS?.destroy();
|
|
5727
|
+
this._outeTTS?.destroy();
|
|
3860
5728
|
this.visionExecutor?.destroy();
|
|
3861
5729
|
clearPipelineCache(this.ctx.device);
|
|
3862
5730
|
this.ctx.device.destroy();
|
|
@@ -3867,5 +5735,5 @@ var WebGPUEngine = class WebGPUEngine {
|
|
|
3867
5735
|
};
|
|
3868
5736
|
|
|
3869
5737
|
//#endregion
|
|
3870
|
-
export {
|
|
3871
|
-
//# sourceMappingURL=gpu-
|
|
5738
|
+
export { dequantizeMLXProjection as A, dacOutputLength as C, KaniTTS as D, parseOuteTtsConfig as E, patchGemma4VisionClips as M, resolveGemma4VisionInfo as N, generateQwen3_5VisionGraph as O, audioTokensToDacCodes as S, generateOuteTtsBackboneGraph as T, smartResize as _, buildGemma4PosEmbeds as a, buildOutePromptString as b, buildMRoPECosSin as c, buildPositionIds as d, buildRotaryCosSin as f, preprocessImageGemma4 as g, preprocessImage as h, buildGemma4PoolMatrix as i, generateGemma4VisionGraph as j, dequantizeGemma4VisionProjection 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, generateDacSpeechDecoderGraph as w, loadOuteSpeaker as x, OuteTTS as y };
|
|
5739
|
+
//# sourceMappingURL=gpu-BECDib8p.mjs.map
|