add rnnoise for voice

1 files changed, 169 insertions, 3 deletions

diff --git a/static/mic-worklet.js b/static/mic-worklet.js
index 814bf6a..3276488 100644
--- a/static/mic-worklet.js
+++ b/static/mic-worklet.js
@@ -1,30 +1,196 @@
-// Microphone capture worklet: resamples the context rate to 48 kHz mono and
+// Microphone capture worklet: resamples the context rate to 48 kHz mono,
+// optionally runs RNNoise (WASM) on 10 ms (480-sample) sub-frames, and
 // posts 20 ms frames (960 Float32 samples) together with their RMS level.
 const OUT_RATE = 48000;
 const FRAME_SAMPLES = 960;
+const RNNOISE_FRAME_SAMPLES = 480;
+const PCM16_SCALE = 32768;
+
+function instantiateRnnoise(wasmBinary) {
+  let memory = null;
+  const resizeHeap = (requestedSize) => {
+    if (!memory) return 0;
+    const pages = ((requestedSize >>> 0) - memory.buffer.byteLength + 65535) >>> 16;
+    try {
+      memory.grow(pages);
+      return 1;
+    } catch {
+      return 0;
+    }
+  };
+  const memcpyBig = (dest, src, num) => {
+    if (!memory) return;
+    new Uint8Array(memory.buffer).copyWithin(dest >>> 0, src >>> 0, (src + num) >>> 0);
+  };
+  const imports = {
+    a: { a: resizeHeap, b: memcpyBig },
+    env: {
+      emscripten_resize_heap: resizeHeap,
+      emscripten_memcpy_big: memcpyBig,
+      __assert_fail: () => {
+        throw new Error("RNNoise WASM assertion failed");
+      },
+    },
+  };
+
+  const buildDenoiser = (instance) => {
+    const exp = instance.exports;
+    memory = exp.memory || exp.c;
+    if (typeof exp.emscripten_stack_init === "function") exp.emscripten_stack_init();
+    const initRuntime = exp.__wasm_call_ctors || exp.d;
+    if (typeof initRuntime === "function") initRuntime();
+
+    const rnnoiseInit = exp.rnnoise_init || exp._rnnoise_init || exp.e;
+    const rnnoiseCreate = exp.rnnoise_create || exp._rnnoise_create || exp.f;
+    const malloc = exp.malloc || exp._malloc || exp.g;
+    const rnnoiseDestroy = exp.rnnoise_destroy || exp._rnnoise_destroy || exp.h;
+    const free = exp.free || exp._free || exp.i;
+    const rnnoiseProcessFrame = exp.rnnoise_process_frame || exp._rnnoise_process_frame || exp.j;
+
+    const state = rnnoiseCreate(0);
+    const ptrIn = malloc(RNNOISE_FRAME_SAMPLES * 4);
+    const ptrOut = malloc(RNNOISE_FRAME_SAMPLES * 4);
+    if (!state || !ptrIn || !ptrOut) {
+      if (state) rnnoiseDestroy(state);
+      if (ptrIn) free(ptrIn);
+      if (ptrOut) free(ptrOut);
+      throw new Error("Failed to allocate RNNoise WASM state");
+    }
+
+    let heapF32 = new Float32Array(memory.buffer);
+    let destroyed = false;
+
+    return {
+      processFrame480(frame, offset) {
+        if (destroyed) return 0;
+        if (heapF32.buffer !== memory.buffer) {
+          heapF32 = new Float32Array(memory.buffer);
+        }
+        const inBase = ptrIn >> 2;
+        const outBase = ptrOut >> 2;
+        for (let i = 0; i < RNNOISE_FRAME_SAMPLES; i++) {
+          heapF32[inBase + i] = frame[offset + i] * PCM16_SCALE;
+        }
+        const vad = rnnoiseProcessFrame(state, ptrOut, ptrIn);
+        if (heapF32.buffer !== memory.buffer) {
+          heapF32 = new Float32Array(memory.buffer);
+        }
+        for (let i = 0; i < RNNOISE_FRAME_SAMPLES; i++) {
+          const s = heapF32[outBase + i] / PCM16_SCALE;
+          frame[offset + i] = s < -1 ? -1 : s > 1 ? 1 : s;
+        }
+        return vad;
+      },
+      reset() {
+        if (destroyed) return;
+        if (typeof rnnoiseInit === "function") {
+          rnnoiseInit(state, 0);
+        }
+      },
+      destroy() {
+        if (destroyed) return;
+        destroyed = true;
+        try { rnnoiseDestroy(state); } catch { /* ignore */ }
+        try { free(ptrIn); } catch { /* ignore */ }
+        try { free(ptrOut); } catch { /* ignore */ }
+      },
+    };
+  };
+
+  return WebAssembly.instantiate(wasmBinary, imports).then((res) => buildDenoiser(res.instance || res));
+}
 
 class Ts6MicCapture extends AudioWorkletProcessor {
-  constructor() {
+  constructor(options) {
     super();
     this.ratio = sampleRate / OUT_RATE; // input samples per output sample
     this.t = 0;      // fractional read position; -1 refers to `last`
     this.last = 0;
     this.frame = new Float32Array(FRAME_SAMPLES);
     this.offset = 0;
+
+    const opts = (options && options.processorOptions) || {};
+    this.rnnoiseEnabled = Boolean(opts.rnnoiseEnabled);
+    this.rnnoise = null;
+    this.disposed = false;
+
+    this.port.onmessage = (e) => this.onMessage(e.data);
+    if (opts.wasmBinary) {
+      this.initRnnoise(opts.wasmBinary);
+    }
+  }
+
+  onMessage(msg) {
+    if (!msg || typeof msg !== "object") return;
+    if (msg.type === "rnnoise-init" && msg.wasmBinary) {
+      if (typeof msg.enabled === "boolean") this.rnnoiseEnabled = msg.enabled;
+      this.initRnnoise(msg.wasmBinary);
+    } else if (msg.type === "rnnoise-enable") {
+      const next = Boolean(msg.enabled);
+      if (next && !this.rnnoiseEnabled && this.rnnoise) {
+        this.rnnoise.reset();
+      }
+      this.rnnoiseEnabled = next;
+    } else if (msg.type === "rnnoise-clear") {
+      this.rnnoiseEnabled = false;
+      if (this.rnnoise) {
+        this.rnnoise.destroy();
+        this.rnnoise = null;
+      }
+      this.port.postMessage({ type: "rnnoise-status", ready: false });
+    } else if (msg.type === "dispose") {
+      this.disposed = true;
+      if (this.rnnoise) {
+        this.rnnoise.destroy();
+        this.rnnoise = null;
+      }
+    }
+  }
+
+  initRnnoise(wasmBinary) {
+    instantiateRnnoise(wasmBinary)
+      .then((denoiser) => {
+        if (this.disposed) {
+          denoiser.destroy();
+          return;
+        }
+        if (this.rnnoise) this.rnnoise.destroy();
+        this.rnnoise = denoiser;
+        this.port.postMessage({ type: "rnnoise-status", ready: true });
+      })
+      .catch((err) => {
+        this.port.postMessage({
+          type: "rnnoise-status",
+          ready: false,
+          error: err && err.message ? err.message : String(err),
+        });
+      });
   }
 
   push(sample) {
     this.frame[this.offset++] = sample;
     if (this.offset < FRAME_SAMPLES) return;
+
+    let denoised = false;
+    if (this.rnnoiseEnabled && this.rnnoise) {
+      this.rnnoise.processFrame480(this.frame, 0);
+      this.rnnoise.processFrame480(this.frame, RNNOISE_FRAME_SAMPLES);
+      denoised = true;
+    }
+
     let sum = 0;
     for (let i = 0; i < FRAME_SAMPLES; i++) sum += this.frame[i] * this.frame[i];
     const frame = this.frame;
     this.frame = new Float32Array(FRAME_SAMPLES);
     this.offset = 0;
-    this.port.postMessage({ samples: frame, rms: Math.sqrt(sum / FRAME_SAMPLES) }, [frame.buffer]);
+    this.port.postMessage(
+      { samples: frame, rms: Math.sqrt(sum / FRAME_SAMPLES), denoised },
+      [frame.buffer]
+    );
   }
 
   process(inputs, outputs) {
+    if (this.disposed) return false;
     const output = outputs[0] && outputs[0][0];
     if (output) output.fill(0);
     const channels = inputs[0];