@tryhamster/gerbil 1.2.2 → 1.3.1

package/dist/{moonshine-stt-Bj1zJfZg.mjs → moonshine-stt-BO5gZ_D6.mjs}

@@ -1088,6 +1088,28 @@ fn main(
 function cdiv(a, b) {
 	return Math.ceil(a / b);
 }
+/**
+* Spread a flat workgroup count over a 2D grid so neither dimension exceeds WebGPU's
+* 65535 per-dimension cap. Used by the audio codec-decoder kernels (Conv1dFull /
+* ConvTranspose1d / Snake1d / RVQDequant) whose outputs can run to hundreds of
+* thousands of workgroups for multi-second clips. The kernels recover the flat index
+* as `gid.y * (num_workgroups.x * 64) + gid.x`, so the x dimension must stay a
+* multiple of the natural grid width — we keep x = min(count, 65535) and tile y.
+*/
+const MAX_WORKGROUPS_PER_DIM = 65535;
+function dispatch2D(workgroups) {
+	if (workgroups <= MAX_WORKGROUPS_PER_DIM) return [
+		workgroups,
+		1,
+		1
+	];
+	const x = MAX_WORKGROUPS_PER_DIM;
+	return [
+		x,
+		cdiv(workgroups, x),
+		1
+	];
+}
 /** Reinterpret an f32 value as its u32 bit pattern. */
 function f32BitsToU32(value) {
 	const f = new Float32Array(1);
@@ -3298,9 +3320,12 @@ struct Params {
 @group(0) @binding(2) var<storage, read_write> output: array<f32>;
 @group(0) @binding(3) var<storage, read> params: Params;
 
+// 2D-dispatch safe: for very large tensors (e.g. DAC vocoder residual adds on
+// multi-second clips) the workgroup count exceeds the 65535 per-dim cap, so the
+// dispatch spreads over y. Flat index = gid.y * (num_workgroups.x * 256) + gid.x.
 @compute @workgroup_size(256)
-fn main(@builtin(global_invocation_id) gid: vec3u) {
-  let idx = gid.x;
+fn main(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) ng: vec3u) {
+  let idx = gid.y * (ng.x * 256u) + gid.x;
   if (idx >= params.count) { return; }
   output[idx] = a[idx] + b[idx];
 }
@@ -5471,11 +5496,7 @@ const addSpec = {
 		{ type: "uniform" }
 	],
 	getDispatchSize(op, resolvedShapes) {
-		return [
-			cdiv(outputElementCount(op, resolvedShapes), 256),
-			1,
-			1
-		];
+		return dispatch2D(cdiv(outputElementCount(op, resolvedShapes), 256));
 	},
 	buildParams(op, resolvedShapes) {
 		return buildUniformBuffer([outputElementCount(op, resolvedShapes)]);
@@ -7046,9 +7067,12 @@ struct Params {
 @group(0) @binding(3) var<storage, read_write> y: array<f32>;
 @group(0) @binding(4) var<storage, read> params: Params;
 
+// 2D dispatch: outputs can exceed 65535 workgroups in x for long audio, so the
+// dispatch spreads across y. Flat index = gid.y * (num_workgroups.x * 64) + gid.x
+// (num_workgroups is the builtin grid size; no extra param needed).
 @compute @workgroup_size(64)
-fn main(@builtin(global_invocation_id) gid: vec3u) {
-  let idx = gid.x;
+fn main(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) ng: vec3u) {
+  let idx = gid.y * (ng.x * 64u) + gid.x;
   let total = params.Cout * params.Lout;
   if (idx >= total) { return; }
   let oc = idx / params.Lout;
@@ -7085,8 +7109,8 @@ struct Params {
 @group(0) @binding(4) var<storage, read> params: Params;
 
 @compute @workgroup_size(64)
-fn main(@builtin(global_invocation_id) gid: vec3u) {
-  let idx = gid.x;
+fn main(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) ng: vec3u) {
+  let idx = gid.y * (ng.x * 64u) + gid.x;
   let total = params.Cout * params.Lout;
   if (idx >= total) { return; }
   let oc = idx / params.Lout;
@@ -7118,8 +7142,8 @@ struct Params { C: u32, L: u32 }
 @group(0) @binding(3) var<storage, read> params: Params;
 
 @compute @workgroup_size(64)
-fn main(@builtin(global_invocation_id) gid: vec3u) {
-  let idx = gid.x;
+fn main(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) ng: vec3u) {
+  let idx = gid.y * (ng.x * 64u) + gid.x;
   let total = params.C * params.L;
   if (idx >= total) { return; }
   let c = idx / params.L;
@@ -7128,6 +7152,145 @@ fn main(@builtin(global_invocation_id) gid: vec3u) {
   y[idx] = x[idx] + (1.0 / (a + 1e-9)) * s * s;
 }
 `;
+const WGSL_RVQ_DEQUANT = `\
+// Residual-Vector-Quantizer dequant for DAC.speech (descript-audio-codec lineage).
+// Reconstruct the continuous latent z[Cout, T] from codes[NCB, T] (u32 indices):
+//   z[oc, t] = sum_c ( out_proj_bias[c][oc]
+//                      + sum_d codebook[c][ codes[c,t], d ] * out_proj_w[c][oc, d] )
+// where each quantizer c (0..NCB-1) has:
+//   codebook  : [codebook_size, codebook_dim]   (nn.Embedding)
+//   out_proj  : Conv1d(codebook_dim -> Cout, kernel=1) -> weight [Cout, codebook_dim], bias [Cout]
+// All NCB quantizers share the same codebook_dim (D) and Cout. Tensors are packed
+// codebook-major: codebook[c] starts at c*codebook_size*D, out_proj_w[c] at c*Cout*D,
+// out_proj_bias[c] at c*Cout, codes[c] row at c*T. One thread per (oc, t) output.
+// Mobile-safe: flat workgroup, no select(), no shared memory, no exp().
+struct Params { NCB: u32, Cout: u32, T: u32, D: u32, codebook_size: u32, _p0: u32, _p1: u32, _p2: u32 }
+@group(0) @binding(0) var<storage, read> codes: array<u32>;      // [NCB, T]
+@group(0) @binding(1) var<storage, read> codebook: array<f32>;   // [NCB, codebook_size, D]
+@group(0) @binding(2) var<storage, read> proj_w: array<f32>;     // [NCB, Cout, D]
+@group(0) @binding(3) var<storage, read> proj_b: array<f32>;     // [NCB, Cout]
+@group(0) @binding(4) var<storage, read_write> z: array<f32>;    // [Cout, T]
+@group(0) @binding(5) var<storage, read> params: Params;
+
+@compute @workgroup_size(64)
+fn main(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) ng: vec3u) {
+  let flat = gid.y * (ng.x * 64u) + gid.x;
+  let total = params.Cout * params.T;
+  if (flat >= total) { return; }
+  let oc = flat / params.T;
+  let t = flat % params.T;
+  let D = params.D;
+  var acc: f32 = 0.0;
+  for (var c: u32 = 0u; c < params.NCB; c = c + 1u) {
+    let idx = codes[c * params.T + t];
+    let cb_row = (c * params.codebook_size + idx) * D;   // codebook[c][idx, :]
+    let w_row = (c * params.Cout + oc) * D;              // out_proj_w[c][oc, :]
+    var dot: f32 = 0.0;
+    for (var d: u32 = 0u; d < D; d = d + 1u) {
+      dot = dot + codebook[cb_row + d] * proj_w[w_row + d];
+    }
+    acc = acc + dot + proj_b[c * params.Cout + oc];
+  }
+  z[oc * params.T + t] = acc;
+}
+`;
+const WGSL_RVQ_ENCODE = `\
+// Residual-Vector-Quantizer ENCODE for DAC.speech (descript-audio-codec lineage),
+// the voice-cloning mirror of RVQDequant. Maps the encoder latent z[Cout, T] to
+// discrete codes[NCB, T] exactly like torch quantizer.encode:
+//   residual = z
+//   for c in 0..NCB:
+//     z_e = in_proj_c(residual)                      # Conv1d(Cout->D, k1): [D, T]
+//     idx = argmin_n || normalize(z_e) - normalize(codebook_c[n]) ||^2   (ViT-VQGAN)
+//     codes[c, t] = idx
+//     z_q = out_proj_c(codebook_c[idx])              # Conv1d(D->Cout, k1): [Cout]
+//     residual = residual - z_q
+// The L2-normalized euclidean argmin equals the max cosine-similarity entry. One
+// thread per frame t; the codebook loop and the (small, recomputed) residual run
+// inside the thread so no Cout-wide scratch is needed. NCB is tiny (2) and the
+// residual removal is recomputed per step from the already-chosen codes, which is
+// O(NCB^2 * Cout * D) per frame — negligible for DAC.speech (NCB=2, Cout=1024, D=8).
+// Mobile-safe: flat workgroup_size(64), no shared memory, no select(), no exp().
+//
+// Layout (codebook-major): z[Cout,T]; in_proj_w[NCB,D,Cout], in_proj_b[NCB,D];
+// codebook[NCB,codebook_size,D]; out_proj_w[NCB,Cout,D], out_proj_b[NCB,Cout].
+const MAX_NCB: u32 = 8u;     // DAC.speech uses 2; cap registers conservatively.
+const MAX_D: u32 = 16u;      // codebook_dim 8.
+struct Params { NCB: u32, Cout: u32, T: u32, D: u32, codebook_size: u32, _p0: u32, _p1: u32, _p2: u32 }
+@group(0) @binding(0) var<storage, read> z: array<f32>;          // [Cout, T]
+@group(0) @binding(1) var<storage, read> in_w: array<f32>;       // [NCB, D, Cout]
+@group(0) @binding(2) var<storage, read> in_b: array<f32>;       // [NCB, D]
+@group(0) @binding(3) var<storage, read> codebook: array<f32>;   // [NCB, codebook_size, D]
+@group(0) @binding(4) var<storage, read> out_w: array<f32>;      // [NCB, Cout, D]
+@group(0) @binding(5) var<storage, read> out_b: array<f32>;      // [NCB, Cout]
+@group(0) @binding(6) var<storage, read_write> codes: array<u32>; // [NCB, T]
+@group(0) @binding(7) var<storage, read> params: Params;
+
+@compute @workgroup_size(64)
+fn main(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) ng: vec3u) {
+  let t = gid.y * (ng.x * 64u) + gid.x;
+  if (t >= params.T) { return; }
+  let Cout = params.Cout;
+  let D = params.D;
+  let cbsize = params.codebook_size;
+
+  var chosen: array<u32, 8>;   // chosen[c] = code idx for codebook c so far (<= MAX_NCB)
+
+  for (var c: u32 = 0u; c < params.NCB; c = c + 1u) {
+    // ── z_e[d] = in_proj_c(residual)[d], residual = z - sum_{c'<c} out_proj_{c'}(cb_{c'}[chosen[c']]) ──
+    // Project the residual without materializing it: in_proj is linear, so
+    //   z_e[d] = in_b[c][d] + sum_oc in_w[c][d][oc] * residual[oc]
+    // and residual[oc] = z[oc] - sum_{c'<c} ( out_b[c'][oc] + sum_k out_w[c'][oc][k]*cb_sel_{c'}[k] ).
+    var ze: array<f32, 16>;
+    for (var d: u32 = 0u; d < D; d = d + 1u) {
+      ze[d] = in_b[c * D + d];
+    }
+    for (var oc: u32 = 0u; oc < Cout; oc = oc + 1u) {
+      // residual at this output channel
+      var res_oc: f32 = z[oc * params.T + t];
+      for (var cp: u32 = 0u; cp < c; cp = cp + 1u) {
+        let sel = chosen[cp];
+        let cb_row = (cp * cbsize + sel) * D;
+        let ow_row = (cp * Cout + oc) * D;
+        var zq: f32 = out_b[cp * Cout + oc];
+        for (var k: u32 = 0u; k < D; k = k + 1u) {
+          zq = zq + out_w[ow_row + k] * codebook[cb_row + k];
+        }
+        res_oc = res_oc - zq;
+      }
+      for (var d: u32 = 0u; d < D; d = d + 1u) {
+        ze[d] = ze[d] + in_w[(c * D + d) * Cout + oc] * res_oc;
+      }
+    }
+
+    // ── L2-normalize z_e over its D dims (ViT-VQGAN); torch F.normalize uses eps=1e-12. ──
+    var ze_norm2: f32 = 0.0;
+    for (var d: u32 = 0u; d < D; d = d + 1u) { ze_norm2 = ze_norm2 + ze[d] * ze[d]; }
+    let ze_inv = 1.0 / max(sqrt(ze_norm2), 1e-12);
+    var zen: array<f32, 16>;
+    for (var d: u32 = 0u; d < D; d = d + 1u) { zen[d] = ze[d] * ze_inv; }
+
+    // ── argmax cosine sim == argmin L2 over the (also L2-normalized) codebook. ──
+    var best: u32 = 0u;
+    var best_sim: f32 = -3.4e38;
+    let cb_base = c * cbsize * D;
+    for (var n: u32 = 0u; n < cbsize; n = n + 1u) {
+      let row = cb_base + n * D;
+      var nn2: f32 = 0.0;
+      var dot: f32 = 0.0;
+      for (var d: u32 = 0u; d < D; d = d + 1u) {
+        let cv = codebook[row + d];
+        nn2 = nn2 + cv * cv;
+        dot = dot + zen[d] * cv;
+      }
+      let sim = dot / max(sqrt(nn2), 1e-12);
+      if (sim > best_sim) { best_sim = sim; best = n; }
+    }
+    chosen[c] = best;
+    codes[c * params.T + t] = best;
+  }
+}
+`;
 const WGSL_TANH = `\
 // Elementwise tanh (Moonshine conv1 activation). Clamp the argument: on Metal/Dawn
 // tanh() of a large magnitude returns NaN (internal exp overflows). |arg|>=10
@@ -7417,11 +7580,7 @@ const conv1dFullSpec = {
 	getDispatchSize(op) {
 		const Cout = op.attributes.Cout;
 		const Lout = op.attributes.Lout;
-		return [
-			cdiv(Cout * Lout, 64),
-			1,
-			1
-		];
+		return dispatch2D(cdiv(Cout * Lout, 64));
 	},
 	buildParams(op) {
 		return buildUniformBuffer([
@@ -7449,11 +7608,7 @@ const convTranspose1dSpec = {
 	getDispatchSize(op) {
 		const Cout = op.attributes.Cout;
 		const Lout = op.attributes.Lout;
-		return [
-			cdiv(Cout * Lout, 64),
-			1,
-			1
-		];
+		return dispatch2D(cdiv(Cout * Lout, 64));
 	},
 	buildParams(op) {
 		return buildUniformBuffer([
@@ -7480,16 +7635,71 @@ const snake1dSpec = {
 	getDispatchSize(op) {
 		const C = op.attributes.C;
 		const L = op.attributes.L;
-		return [
-			cdiv(C * L, 64),
-			1,
-			1
-		];
+		return dispatch2D(cdiv(C * L, 64));
 	},
 	buildParams(op) {
 		return buildUniformBuffer([op.attributes.C, op.attributes.L]);
 	}
 };
+const rvqDequantSpec = {
+	shaderCode: WGSL_RVQ_DEQUANT,
+	entryPoint: "main",
+	bindings: [
+		{ type: "storage-read" },
+		{ type: "storage-read" },
+		{ type: "storage-read" },
+		{ type: "storage-read" },
+		{ type: "storage-read-write" },
+		{ type: "uniform" }
+	],
+	getDispatchSize(op) {
+		const Cout = op.attributes.Cout;
+		const T = op.attributes.T;
+		return dispatch2D(cdiv(Cout * T, 64));
+	},
+	buildParams(op) {
+		return buildUniformBuffer([
+			op.attributes.NCB,
+			op.attributes.Cout,
+			op.attributes.T,
+			op.attributes.D,
+			op.attributes.codebook_size,
+			0,
+			0,
+			0
+		]);
+	}
+};
+const rvqEncodeSpec = {
+	shaderCode: WGSL_RVQ_ENCODE,
+	entryPoint: "main",
+	bindings: [
+		{ type: "storage-read" },
+		{ type: "storage-read" },
+		{ type: "storage-read" },
+		{ type: "storage-read" },
+		{ type: "storage-read" },
+		{ type: "storage-read" },
+		{ type: "storage-read-write" },
+		{ type: "uniform" }
+	],
+	getDispatchSize(op) {
+		const T = op.attributes.T;
+		return dispatch2D(cdiv(T, 64));
+	},
+	buildParams(op) {
+		return buildUniformBuffer([
+			op.attributes.NCB,
+			op.attributes.Cout,
+			op.attributes.T,
+			op.attributes.D,
+			op.attributes.codebook_size,
+			0,
+			0,
+			0
+		]);
+	}
+};
 const WGSL_FSQ_DEQUANT = `\
 // Finite-scalar-quantization decode. codes[groups,T] (u32) -> latent[groups*dims,T]
 // (f32). Channel ch = g*dims + d; per-dim level/base passed as small constants.
@@ -7724,6 +7934,8 @@ const KERNEL_REGISTRY = {
 	Conv1dFull: conv1dFullSpec,
 	ConvTranspose1d: convTranspose1dSpec,
 	Snake1d: snake1dSpec,
+	RVQDequant: rvqDequantSpec,
+	RVQEncode: rvqEncodeSpec,
 	FSQDequant: fsqDequantSpec,
 	HalfSnake1d: halfSnake1dSpec,
 	ConvTranspose1dDepthwise: convTranspose1dDepthwiseSpec,
@@ -12036,6 +12248,64 @@ function quantizeKaniBackbone(graph, weights) {
 		weights.delete(sourceName);
 	}
 }
+/**
+* Generic INT4 backbone quantizer (architecture-neutral alias of quantizeKaniBackbone).
+* Walks a built graph's MatMulInt4 nodes and quantizes each linear's f32 source weight
+* into `.q`/`.scales`/`.zeros`, freeing the f32. Used by the OuteTTS (Qwen3) backbone;
+* the embedding/LM-head are left untouched (the graph keeps them f32). Idempotent.
+*/
+const quantizeBackboneInt4 = quantizeKaniBackbone;
+const OUTETTS_DEFAULT_REPO = "OuteAI/OuteTTS-1.0-0.6B";
+/** Parse a folded DAC safetensors buffer into a canonical-name weight map (F32 only). */
+function parseFoldedDacSafetensors(buf) {
+	const file = parseSafetensorsHeader(buf);
+	const out = /* @__PURE__ */ new Map();
+	for (const entry of file.entries) out.set(entry.name, {
+		data: toF32(entry, buf, file),
+		shape: entry.shape
+	});
+	return out;
+}
+async function loadOuteTTS(options) {
+	const revision = options.revision ?? "main";
+	const baseURL = await resolveModelBaseURL(options.repo ?? OUTETTS_DEFAULT_REPO, revision);
+	await initFs();
+	options.onProgress?.(0, 100, "Fetching config + tokenizer…");
+	const [rawConfig, tokenizerJSON] = await Promise.all([fetchJSON(baseURL, "config.json", options.hfToken, options.cacheDir), fetchJSON(baseURL, "tokenizer.json", options.hfToken, options.cacheDir)]);
+	const tokenizerConfig = await fetchJSON(baseURL, "tokenizer_config.json", options.hfToken, options.cacheDir).catch(() => null);
+	const tokenizer = Tokenizer.fromJSON(tokenizerJSON, tokenizerConfig);
+	options.onProgress?.(8, 100, "Downloading OuteTTS backbone…");
+	const backboneBuf = await fetchBinary(baseURL, "model.safetensors", options.hfToken, (loaded, total) => options.onProgress?.(8 + loaded / (total || 1) * 75, 100, "backbone weights"), options.cacheDir);
+	const backboneFile = parseSafetensorsHeader(backboneBuf);
+	const backboneWeights = /* @__PURE__ */ new Map();
+	for (const entry of backboneFile.entries) {
+		const key = entry.name.startsWith("model.") ? entry.name.slice(6) : entry.name;
+		backboneWeights.set(key, {
+			data: toF32(entry, backboneBuf, backboneFile),
+			shape: entry.shape
+		});
+	}
+	options.onProgress?.(85, 100, "Loading DAC.speech decoder…");
+	const codecPath = options.codecWeightsPath;
+	let codecBuf;
+	if (/^https?:\/\//.test(codecPath)) {
+		const resp = await fetch(codecPath);
+		if (!resp.ok) throw new Error(`Failed to fetch DAC codec weights from ${codecPath}: ${resp.status}`);
+		codecBuf = await resp.arrayBuffer();
+	} else {
+		if (!_fs) throw new Error("loadOuteTTS: codecWeightsPath is a local path but Node fs is unavailable. Pass a mirror URL in the browser.");
+		const fileBuf = _fs.readFileSync(codecPath);
+		codecBuf = fileBuf.buffer.slice(fileBuf.byteOffset, fileBuf.byteOffset + fileBuf.byteLength);
+	}
+	const codecWeights = parseFoldedDacSafetensors(codecBuf);
+	options.onProgress?.(100, 100, "OuteTTS loaded.");
+	return {
+		backboneWeights,
+		codecWeights,
+		tokenizer,
+		rawConfig
+	};
+}
 
 //#endregion
 //#region src/gpu/moonshine-executor.ts
@@ -12401,5 +12671,5 @@ function selectGraphWeights(graph, weights) {
 }
 
 //#endregion
-export { initGPU as _, loadMoonshine as a, Executor as c, clearPipelineCache as d, createBindGroup as f, getOrCreatePipeline as g, destroyBuffers as h, loadModel as i, KERNEL_REGISTRY as l, createUniformBuffer as m, MoonshineEncoderExecutor as n, quantizeKaniBackbone as o, createStorageBuffer as p, loadKaniTTS as r, remapPrunedToken as s, MoonshineSTT as t, MATMUL_BIAS_F16C_SPEC as u, verifyGPU as v };
-//# sourceMappingURL=moonshine-stt-Bj1zJfZg.mjs.map
+export { destroyBuffers as _, loadMoonshine as a, verifyGPU as b, quantizeKaniBackbone as c, KERNEL_REGISTRY as d, MATMUL_BIAS_F16C_SPEC as f, createUniformBuffer as g, createStorageBuffer as h, loadModel as i, remapPrunedToken as l, createBindGroup as m, MoonshineEncoderExecutor as n, loadOuteTTS as o, clearPipelineCache as p, loadKaniTTS as r, quantizeBackboneInt4 as s, MoonshineSTT as t, Executor as u, getOrCreatePipeline as v, initGPU as y };
+//# sourceMappingURL=moonshine-stt-BO5gZ_D6.mjs.map