@tryhamster/gerbil 1.1.3 → 1.2.1

package/dist/gpu/index.mjs

@@ -1,6 +1,6 @@
-import { n as resolveDefaultRepo, t as DEFAULT_MODELS } from "../defaults-9komdrbY.mjs";
-import { C as parseKaniConfig, _ as generateNanoCodecDecoderGraph, a as generateMoonshineEncoderGraph, f as audioTokensToCodes, g as generateKaniTtsGraph, i as generateMoonshineDecoderGraph, o as moonshineEncoderFrames, r as MOONSHINE_REMAINING_WORK, s as parseMoonshineConfig } from "../architectures-wan0swXu.mjs";
-import { C as generateGemma4VisionGraph, S as dequantizeMLXProjection, T as resolveGemma4VisionInfo, _ as smartResize, a as buildGemma4PosEmbeds, b as generateQwen3_5VisionGraph, c as buildMRoPECosSin, d as buildPositionIds, f as buildRotaryCosSin, g as preprocessImageGemma4, h as preprocessImage, i as buildGemma4PoolMatrix, l as buildMRoPEPositionIds, m as mropeFreqDims, n as GEMMA4_IMAGE_PROCESSOR, o as buildGemma4RotaryCosSin, p as buildVisionPositionTensors, r as QWEN3_5_IMAGE_PROCESSOR, s as buildGemma4VisionPositionTensors, t as WebGPUEngine, u as buildPosEmbeds, v as VisionExecutor, w as patchGemma4VisionClips, x as dequantizeGemma4VisionProjection, y as KaniTTS } from "../gpu-BYCQ5aTI.mjs";
-import { a as loadMoonshine, g as initGPU, i as loadModel, n as MoonshineEncoderExecutor, r as loadKaniTTS, s as Executor, t as MoonshineSTT } from "../moonshine-stt-DDUm5WfP.mjs";
+import { n as resolveDefaultRepo, t as DEFAULT_MODELS } from "../defaults-bJcmOaRG.mjs";
+import { a as generateMoonshineEncoderGraph, h as generateNanoCodecDecoderGraph, i as generateMoonshineDecoderGraph, m as generateKaniTtsGraph, o as moonshineEncoderFrames, r as MOONSHINE_REMAINING_WORK, s as parseMoonshineConfig, u as audioTokensToCodes, x as parseKaniConfig } from "../architectures-Ol7yUmKf.mjs";
+import { C as generateGemma4VisionGraph, S as dequantizeMLXProjection, T as resolveGemma4VisionInfo, _ as smartResize, a as buildGemma4PosEmbeds, b as generateQwen3_5VisionGraph, c as buildMRoPECosSin, d as buildPositionIds, f as buildRotaryCosSin, g as preprocessImageGemma4, h as preprocessImage, i as buildGemma4PoolMatrix, l as buildMRoPEPositionIds, m as mropeFreqDims, n as GEMMA4_IMAGE_PROCESSOR, o as buildGemma4RotaryCosSin, p as buildVisionPositionTensors, r as QWEN3_5_IMAGE_PROCESSOR, s as buildGemma4VisionPositionTensors, t as WebGPUEngine, u as buildPosEmbeds, v as VisionExecutor, w as patchGemma4VisionClips, x as dequantizeGemma4VisionProjection, y as KaniTTS } from "../gpu-DR6EG0K-.mjs";
+import { _ as initGPU, a as loadMoonshine, c as Executor, i as loadModel, n as MoonshineEncoderExecutor, o as quantizeKaniBackbone, r as loadKaniTTS, t as MoonshineSTT } from "../moonshine-stt-DSz3VW_0.mjs";
 
-export { DEFAULT_MODELS, Executor, GEMMA4_IMAGE_PROCESSOR, KaniTTS, MOONSHINE_REMAINING_WORK, MoonshineEncoderExecutor, MoonshineSTT, QWEN3_5_IMAGE_PROCESSOR, VisionExecutor, WebGPUEngine, audioTokensToCodes, buildGemma4PoolMatrix, buildGemma4PosEmbeds, buildGemma4RotaryCosSin, buildGemma4VisionPositionTensors, buildMRoPECosSin, buildMRoPEPositionIds, buildPosEmbeds, buildPositionIds, buildRotaryCosSin, buildVisionPositionTensors, dequantizeGemma4VisionProjection, dequantizeMLXProjection, generateGemma4VisionGraph, generateKaniTtsGraph, generateMoonshineDecoderGraph, generateMoonshineEncoderGraph, generateNanoCodecDecoderGraph, generateQwen3_5VisionGraph, initGPU, loadKaniTTS, loadModel, loadMoonshine, moonshineEncoderFrames, mropeFreqDims, parseKaniConfig, parseMoonshineConfig, patchGemma4VisionClips, preprocessImage, preprocessImageGemma4, resolveDefaultRepo, resolveGemma4VisionInfo, smartResize };
+export { DEFAULT_MODELS, Executor, GEMMA4_IMAGE_PROCESSOR, KaniTTS, MOONSHINE_REMAINING_WORK, MoonshineEncoderExecutor, MoonshineSTT, QWEN3_5_IMAGE_PROCESSOR, VisionExecutor, WebGPUEngine, audioTokensToCodes, buildGemma4PoolMatrix, buildGemma4PosEmbeds, buildGemma4RotaryCosSin, buildGemma4VisionPositionTensors, buildMRoPECosSin, buildMRoPEPositionIds, buildPosEmbeds, buildPositionIds, buildRotaryCosSin, buildVisionPositionTensors, dequantizeGemma4VisionProjection, dequantizeMLXProjection, generateGemma4VisionGraph, generateKaniTtsGraph, generateMoonshineDecoderGraph, generateMoonshineEncoderGraph, generateNanoCodecDecoderGraph, generateQwen3_5VisionGraph, initGPU, loadKaniTTS, loadModel, loadMoonshine, moonshineEncoderFrames, mropeFreqDims, parseKaniConfig, parseMoonshineConfig, patchGemma4VisionClips, preprocessImage, preprocessImageGemma4, quantizeKaniBackbone, resolveDefaultRepo, resolveGemma4VisionInfo, smartResize };

package/dist/{gpu-BYCQ5aTI.mjs → gpu-DR6EG0K-.mjs}

@@ -1,6 +1,6 @@
-import { n as resolveDefaultRepo } from "./defaults-9komdrbY.mjs";
-import { C as parseKaniConfig, E as GEMMA4_VIS_KEYS, T as DTYPE_BYTES, _ as generateNanoCodecDecoderGraph, b as kaniLayerAlpha, d as KANI_START_OF_HUMAN, f as audioTokensToCodes, g as generateKaniTtsGraph, m as computeKaniPositions, p as buildKaniLayerCosSin, u as KANI_END_OF_HUMAN, v as kaniAttentionLayerIndices, w as CANONICAL_KEYS, x as kaniSinTensor, y as kaniCosTensor } from "./architectures-wan0swXu.mjs";
-import { _ as verifyGPU, c as KERNEL_REGISTRY, d as createBindGroup, f as createStorageBuffer, g as initGPU, h as getOrCreatePipeline, i as loadModel, l as MATMUL_BIAS_F16C_SPEC, m as destroyBuffers, o as remapPrunedToken, p as createUniformBuffer, r as loadKaniTTS, s as Executor, u as clearPipelineCache } from "./moonshine-stt-DDUm5WfP.mjs";
+import { n as resolveDefaultRepo, t as DEFAULT_MODELS } from "./defaults-bJcmOaRG.mjs";
+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";
+import { _ as initGPU, c as Executor, d as clearPipelineCache, f as createBindGroup, g as getOrCreatePipeline, h as destroyBuffers, i as loadModel, l as KERNEL_REGISTRY, m as createUniformBuffer, o as quantizeKaniBackbone, p as createStorageBuffer, r as loadKaniTTS, s as remapPrunedToken, u as MATMUL_BIAS_F16C_SPEC, v as verifyGPU } from "./moonshine-stt-DSz3VW_0.mjs";
 
 //#region src/gpu/architectures/gemma4_vision.ts
 /**
@@ -802,6 +802,19 @@ function nucleusSample(probs, ids, topP) {
 	}
 	return ids[order[0]];
 }
+/**
+* In-place top-k filter over a probability array: keep only the `k` highest, zero
+* the rest, and renormalize. No-op when k <= 0 or k >= length (keep everything).
+*/
+function applyTopK(scores, k) {
+	const n = scores.length;
+	if (k <= 0 || k >= n) return;
+	const order = Array.from({ length: n }, (_, i) => i).sort((a, b) => scores[b] - scores[a]);
+	let keepSum = 0;
+	for (let r = 0; r < k; r++) keepSum += scores[order[r]];
+	for (let r = k; r < n; r++) scores[order[r]] = 0;
+	if (keepSum > 0) for (let i = 0; i < n; i++) scores[i] /= keepSum;
+}
 /** In-place softmax of `scores`. */
 function softmaxInPlace(scores) {
 	let maxV = Number.NEGATIVE_INFINITY;
@@ -829,14 +842,23 @@ var KaniTTS = class KaniTTS {
 	headDim;
 	ropeBase;
 	_destroyed = false;
-	architecture = "KaniTTS2ForCausalLM";
-	constructor(ctx, loaded, maxSeqLen) {
+	/** Accent/voice tags the loaded checkpoint advertises (empty = single voice). */
+	voiceTags;
+	/** The loaded checkpoint's architecture string ("KaniTTS2ForCausalLM" or "Lfm2ForCausalLM"). */
+	architecture;
+	/** The selected backbone precision ("q4" or "f32"). */
+	dtype;
+	constructor(ctx, loaded, maxSeqLen, dtype) {
 		this.ctx = ctx;
 		this.loaded = loaded;
 		this.tokenizer = loaded.tokenizer;
 		this.rawConfig = loaded.rawConfig;
 		this.cfg = parseKaniConfig(loaded.rawConfig);
 		this.maxSeqLen = maxSeqLen;
+		this.dtype = loaded.preQuantized ? "q4" : dtype;
+		this.architecture = (loaded.rawConfig.architectures ?? [])[0] ?? "Lfm2ForCausalLM";
+		const langSettings = loaded.rawConfig.language_settings;
+		this.voiceTags = Array.isArray(langSettings?.language_tags_list) ? langSettings.language_tags_list.filter((v) => typeof v === "string") : [];
 		const hidden = this.rawConfig.hidden_size;
 		const heads = this.rawConfig.num_attention_heads;
 		this.headDim = this.rawConfig.head_dim ?? Math.floor(hidden / heads);
@@ -852,13 +874,26 @@ var KaniTTS = class KaniTTS {
 			}
 			this.layerAlpha.set(layer, alpha);
 		}
-		const graph = generateKaniTtsGraph(this.rawConfig);
-		this.backboneExec = new Executor(ctx, graph, {
+		this.backboneExec = this.buildBackbone(ctx, loaded, maxSeqLen);
+	}
+	/**
+	* Build the codec-LM backbone graph, quantize it when q4, and upload its weights
+	* into a fresh Executor. For a pre-quantized (gerbil-int4) checkpoint the linear
+	* projections already carry `.q`/`.scales`/`.zeros`, so quantizeKaniBackbone is an
+	* idempotent no-op (it skips any projection whose `.q` is already present) and the
+	* int4 weights are uploaded verbatim — no on-device re-quantization.
+	*/
+	buildBackbone(ctx, loaded, maxSeqLen) {
+		const useQ4 = this.dtype === "q4";
+		const graph = generateKaniTtsGraph(this.rawConfig, useQ4 ? "q4" : void 0);
+		if (useQ4) quantizeKaniBackbone(graph, loaded.backboneWeights);
+		const exec = new Executor(ctx, graph, {
 			maxSeqLen,
 			kvMode: "f32"
 		});
-		this.backboneExec.uploadWeightsMap(selectGraphWeights(graph, loaded.backboneWeights));
-		this.backboneExec.initBindGroups();
+		exec.uploadWeightsMap(selectGraphWeights(graph, loaded.backboneWeights));
+		exec.initBindGroups();
+		return exec;
 	}
 	static async create(options = {}) {
 		const ctx = await initGPU();
@@ -871,7 +906,7 @@ var KaniTTS = class KaniTTS {
 			onProgress: options.onProgress
 		});
 		const ctxLen = loaded.rawConfig.max_position_embeddings ?? 4096;
-		return new KaniTTS(ctx, loaded, Math.min(options.maxSeqLen ?? 2048, ctxLen));
+		return new KaniTTS(ctx, loaded, Math.min(options.maxSeqLen ?? 2048, ctxLen), options.dtype ?? "q4");
 	}
 	/** Write per-layer cos/sin for token rows [rowStart, rowStart+positions.length). */
 	writeCosSin(positions, rowStart) {
@@ -889,6 +924,20 @@ var KaniTTS = class KaniTTS {
 		}
 	}
 	/**
+	* Accent/voice tags the loaded checkpoint supports. Multi-accent KaniTTS2
+	* checkpoints (kani-tts-2-en) return e.g. ["en_us","en_nyork","en_glasg", …];
+	* the plain LFM2 family (the Apache default) returns [] (one trained voice,
+	* selectable only by fine-tuning the backbone — there is no runtime voice/
+	* speaker-prompt knob, and zero-shot cloning is not wired).
+	*/
+	availableVoices() {
+		return [...this.voiceTags];
+	}
+	/** NanoCodec's fixed output sample rate (Hz). Not selectable — see SpeakResult. */
+	get sampleRate() {
+		return KANI_SAMPLE_RATE;
+	}
+	/**
 	* Synthesize speech for `text`. Returns 22 kHz mono PCM.
 	*
 	* Pipeline: build the [SOH]+text+[EOT,EOH] prompt → prefill the codec-LM →
@@ -899,12 +948,13 @@ var KaniTTS = class KaniTTS {
 		if (this._destroyed) throw new Error("KaniTTS has been destroyed.");
 		const temperature = opts.temperature ?? 1;
 		const topP = opts.topP ?? .95;
+		const topK = opts.topK ?? 0;
 		const repetitionPenalty = opts.repetitionPenalty ?? 1.1;
-		const tag = opts.languageTag ?? "en_us";
 		const maxNewTokens = Math.min(opts.maxNewTokens ?? DEFAULT_MAX_NEW_TOKENS, this.maxSeqLen);
+		const promptText = this.voiceTags.length > 0 ? `${(opts.languageTag ?? this.voiceTags[0]).trim()}: ${text}` : text;
 		const prompt = [
 			KANI_START_OF_HUMAN,
-			...this.tokenizer.encode(`${tag.trim()}: ${text}`),
+			...this.tokenizer.encode(promptText),
 			KANI_END_OF_TEXT,
 			KANI_END_OF_HUMAN
 		];
@@ -917,6 +967,7 @@ var KaniTTS = class KaniTTS {
 		const { audioTokens, finished } = await this.runDecodeLoop(prompt, logits, {
 			temperature,
 			topP,
+			topK,
 			repetitionPenalty,
 			maxNewTokens,
 			maxFrames: opts.maxFrames ?? Number.POSITIVE_INFINITY
@@ -953,6 +1004,7 @@ var KaniTTS = class KaniTTS {
 			const nextToken = inSpeech ? this.sampleAudioToken(logits, audioTokens.length % this.cfg.tokensPerFrame, {
 				temperature: p.temperature,
 				topP: p.topP,
+				topK: p.topK,
 				repetitionPenalty: p.repetitionPenalty,
 				previous: seq
 			}) : this.argmax(logits);
@@ -994,7 +1046,8 @@ var KaniTTS = class KaniTTS {
 	/**
 	* Sample one audio token for codebook position `codebook`, constrained to that
 	* codebook's valid window [audio_tokens_start + 4032*c, +4032). Allows end_of_speech
-	* only at codebook 0 (frame boundary). Applies temperature, top-p, rep-penalty.
+	* only at codebook 0 (frame boundary). Applies repetition penalty, temperature,
+	* optional top-k, then top-p nucleus sampling.
 	*/
 	sampleAudioToken(logits, codebook, p) {
 		const winStart = this.cfg.audioTokensStart + this.cfg.codebookSize * codebook;
@@ -1012,6 +1065,7 @@ var KaniTTS = class KaniTTS {
 			scores[i] = s / p.temperature;
 		}
 		softmaxInPlace(scores);
+		applyTopK(scores, p.topK);
 		return nucleusSample(scores, ids, p.topP);
 	}
 	/**
@@ -3069,29 +3123,44 @@ var WebGPUEngine = class WebGPUEngine {
 		throw new Error(`generateObject failed after ${attemptsMax} attempt(s): ${lastError}. Last output: ${JSON.stringify(lastText.slice(0, 200))}`);
 	}
 	/**
-	* Text-to-speech: text → 22 kHz PCM via Kani-TTS-2 (LFM2-350M codec-LM + NVIDIA
+	* Text-to-speech: text → 22 kHz PCM via KaniTTS (LFM2-350M codec-LM + NVIDIA
 	* NeMo NanoCodec). Returns `{ pcm: Float32Array, sampleRate: 22050 }`.
 	*
 	* Runs the full pipeline: the codec-LM backbone autoregressively emits NanoCodec
-	* audio tokens (4 per frame, frame-level positions + learnable per-layer RoPE),
-	* then the bit-exact NanoCodec decoder (FSQ + causal HiFi-GAN) turns the codes
-	* into PCM. The heavy lifting lives in {@link KaniTTS} (src/gpu/kani-tts.ts); this
-	* lazily constructs that engine on first use (downloading the NanoCodec codec
-	* checkpoint alongside the backbone).
+	* audio tokens (4 per frame), then the bit-exact NanoCodec decoder (FSQ + causal
+	* HiFi-GAN) turns the codes into PCM. The heavy lifting lives in {@link KaniTTS}
+	* (src/gpu/kani-tts.ts); this lazily constructs that engine on first use
+	* (downloading the NanoCodec codec checkpoint alongside the backbone).
 	*
-	* Requires a Kani-TTS-2 checkpoint (architecture "KaniTTS2ForCausalLM").
+	* The default checkpoint is the Apache-2.0 `nineninesix/kani-tts-450m-0.2-ft`
+	* (plain `Lfm2ForCausalLM` backbone); the earlier `KaniTTS2ForCausalLM`
+	* checkpoints (multi-accent) also work. The codec-LM backbone defaults to INT4
+	* (`dtype: "q4"`) — pass `dtype: "f32"` for full precision.
 	*/
 	async speak(text, options = {}) {
 		this.checkDestroyed();
-		if (this._architecture !== "KaniTTS2ForCausalLM") throw new Error(`speak() requires a Kani-TTS-2 model (architecture "KaniTTS2ForCausalLM"), loaded engine is "${this._architecture}".`);
+		const repo = this._createOptions.repo ?? "";
+		if (!(this._architecture === "KaniTTS2ForCausalLM" || /kani-tts/i.test(repo) || /\btts\b/i.test(repo))) throw new Error(`speak() requires a Kani-TTS checkpoint (e.g. "${DEFAULT_MODELS.tts}"), loaded engine is "${this._architecture}" (repo "${repo}").`);
 		if (!this._kaniTTS) this._kaniTTS = await KaniTTS.create({
 			repo: this._createOptions.repo,
 			revision: this._createOptions.revision,
 			hfToken: this._createOptions.hfToken,
 			cacheDir: this._createOptions.cacheDir,
-			maxSeqLen: this.maxSeqLen
+			maxSeqLen: this.maxSeqLen,
+			dtype: options.dtype
+		});
+		return this._kaniTTS.speak(text, {
+			languageTag: options.voice ?? options.languageTag,
+			temperature: options.temperature,
+			topP: options.topP,
+			topK: options.topK,
+			repetitionPenalty: options.repetitionPenalty,
+			maxFrames: options.maxFrames
 		});
-		return this._kaniTTS.speak(text, options);
+	}
+	/** Accent/voice tags the lazily-built TTS engine supports ([] = single voice). */
+	ttsVoices() {
+		return this._kaniTTS?.availableVoices() ?? [];
 	}
 	/**
 	* Describe an image: image-in → text-out. Runs the vision encoder, splices the
@@ -3799,4 +3868,4 @@ var WebGPUEngine = class WebGPUEngine {
 
 //#endregion
 export { generateGemma4VisionGraph as C, dequantizeMLXProjection as S, resolveGemma4VisionInfo as T, smartResize as _, buildGemma4PosEmbeds as a, generateQwen3_5VisionGraph as b, buildMRoPECosSin as c, buildPositionIds as d, buildRotaryCosSin as f, preprocessImageGemma4 as g, preprocessImage as h, buildGemma4PoolMatrix as i, 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, patchGemma4VisionClips as w, dequantizeGemma4VisionProjection as x, KaniTTS as y };
-//# sourceMappingURL=gpu-BYCQ5aTI.mjs.map
+//# sourceMappingURL=gpu-DR6EG0K-.mjs.map