Use .npy model file instead of .bin file

rnnoise/index.html

@@ -66,6 +66,7 @@
     <script src="https://stackpath.bootstrapcdn.com/bootstrap/4.5.2/js/bootstrap.min.js" integrity="sha384-B4gt1jrGC7Jh4AgTPSdUtOBvfO8shuf57BaghqFfPlYxofvL8/KUEfYiJOMMV+rV" crossorigin="anonymous"></script>
     <script src="https://webmachinelearning.github.io/webnn-polyfill/dist/webnn-polyfill.js"></script>
     <script src="signal/signal.js"></script>
+    <script src="utils/numpy.js"></script>
     <script type="module" src="./rnnoise.js"></script>
     <script type="module" src="./processer.js"></script>
     <script type="module" src="./main.js"></script>

rnnoise/main.js

@@ -22,7 +22,7 @@ const DenoiseInfo = document.getElementById('denoise-info');
 const fileInput = document.getElementById('file-input');
 const originalAudio = document.getElementById('original-audio');
 const denoisedAudio = document.getElementById('denoised-audio');
-const recorderWorker = new Worker('./recorderWorker.js');
+const recorderWorker = new Worker('./utils/recorderWorker.js');
 
 recorderWorker.postMessage({
   command: 'init',
@@ -65,7 +65,10 @@ async function denoise() {
   let audioData = [];
   let audioContext = new AudioContext({sampleRate: 48000});
   let sampleRate = audioContext.sampleRate;
-  let steps = 40000;
+  let steps = 48000;
+  let vadInitialHiddenStateBuffer = new Float32Array(1 * batchSize * 24).fill(0);
+  let noiseInitialHiddenStateBuffer = new Float32Array(1 * batchSize * 48).fill(0);
+  let denoiseInitialHiddenStateBuffer = new Float32Array(1 * batchSize * 96).fill(0);
 
   if(audioContext.state != "running") {
     audioContext.resume().then(function() {
@@ -91,13 +94,20 @@ async function denoise() {
     let start = performance.now();
     let features = analyser.preProcessing(framePCM);
     const preProcessingTime = (performance.now() - start).toFixed(2);
-    let inputTensor = new Float32Array(features);
+    let inputBuffer = new Float32Array(features);
     start = performance.now();
-    let outputTensor = await rnnoise.compute(inputTensor);
+    let outputs = await rnnoise.compute(
+      inputBuffer, vadInitialHiddenStateBuffer, 
+      noiseInitialHiddenStateBuffer, denoiseInitialHiddenStateBuffer
+    );
     const executionTime = (performance.now() - start).toFixed(2);
     // rnnoise.dispose();
+    vadInitialHiddenStateBuffer = outputs.vadGruYH.buffer;
+    noiseInitialHiddenStateBuffer = outputs.noiseGruYH.buffer;
+    denoiseInitialHiddenStateBuffer = outputs.denoiseGruYH.buffer;
+
     start = performance.now();
-    let output = analyser.postProcessing(outputTensor.buffer);
+    let output = analyser.postProcessing(outputs.denoiseOutput.buffer);
     const postProcessingTime = (performance.now() - start).toFixed(2);
     if (i == 0 ) {
       audioData.push(...output);
@@ -162,7 +172,7 @@ fileInput.addEventListener('input', (event) => {
 });
 
 window.onload = async function () {
-  log(modelInfo, `Creating NSNet2 with input shape ` +
+  log(modelInfo, `Creating RNNoise with input shape ` +
     `[${batchSize} (batch_size) x 100 (frames) x 42].`, true);
   log(modelInfo, '- Loading model...');
   let start = performance.now();

rnnoise/rnnoise.js

@@ -2,100 +2,134 @@
 
 const nn = navigator.ml.getNeuralNetworkContext();
 
+function sizeOfShape(shape) {
+  return shape.reduce((a, b) => {
+    return a * b;
+  });
+}
+
 export class RNNoise {
   constructor(url, batchSize) {
     this.url_ = url;
     this.batchSize_ = batchSize;
     this.frames_ = 100;
     this.model_ = null;
     this.compilation_ = null;
+    this.builder = null;
   }
 
-  async fetchData(fileName) {
+  async buildConstantByNpy(fileName) {
+    const dataTypeMap = new Map([
+      ['f2', {type: 'float16', array: Uint16Array}],
+      ['f4', {type: 'float32', array: Float32Array}],
+      ['f8', {type: 'float64', array: Float64Array}],
+      ['i1', {type: 'int8', array: Int8Array}],
+      ['i2', {type: 'int16', array: Int16Array}],
+      ['i4', {type: 'int32', array: Int32Array}],
+      ['i8', {type: 'int64', array: BigInt64Array}],
+      ['u1', {type: 'uint8', array: Uint8Array}],
+      ['u2', {type: 'uint16', array: Uint16Array}],
+      ['u4', {type: 'uint32', array: Uint32Array}],
+      ['u8', {type: 'uint64', array: BigUint64Array}],
+    ]);
     const response = await fetch(this.url_ + fileName);
-    return new Float32Array(await response.arrayBuffer());
+    const buffer = await response.arrayBuffer();
+    const npArray = new numpy.Array(new Uint8Array(buffer));
+    if (!dataTypeMap.has(npArray.dataType)) {
+      throw new Error(`Data type ${npArray.dataType} is not supported.`);
+    }
+    const dimensions = npArray.shape;
+    const type = dataTypeMap.get(npArray.dataType).type;
+    const TypedArrayConstructor = dataTypeMap.get(npArray.dataType).array;
+    const typedArray = new TypedArrayConstructor(sizeOfShape(dimensions));
+    const dataView = new DataView(npArray.data.buffer);
+    const littleEndian = npArray.byteOrder === '<';
+    for (let i = 0; i < sizeOfShape(dimensions); ++i) {
+      typedArray[i] = dataView[`get` + type[0].toUpperCase() + type.substr(1)](
+          i * TypedArrayConstructor.BYTES_PER_ELEMENT, littleEndian);
+    }
+    return this.builder.constant({type, dimensions}, typedArray);
   }
 
   async load() {
-    const inputDenseKernel0Data = await this.fetchData('input_dense_kernel_0.bin');
-    const inputDenseBias0Data = await this.fetchData('input_dense_bias_0.bin');
-    const vadGruWData = await this.fetchData('vad_gru_W.bin');
-    const vadGruRData = await this.fetchData('vad_gru_R.bin');
-    const vadGruBData = await this.fetchData('vad_gru_B.bin');
-    const noiseGruWData = await this.fetchData('noise_gru_W.bin');
-    const noiseGruRData = await this.fetchData('noise_gru_R.bin');
-    const noiseGruBData = await this.fetchData('noise_gru_B.bin');
-    const denoiseGruWData = await this.fetchData('denoise_gru_W.bin');
-    const denoiseGruRData = await this.fetchData('denoise_gru_R.bin');
-    const denoiseGruBData = await this.fetchData('denoise_gru_B.bin');
-    const denoiseOutputKernel0Data = await this.fetchData('denoise_output_kernel_0.bin');
-    const denoiseOutputBias0Data = await this.fetchData('denoise_output_bias_0.bin');
-
-    const builder = nn.createModelBuilder();
-
-    const input = builder.input('input', {type: 'float32', dimensions: [this.batchSize_, 100, 42]});
-    const inputDenseKernel0 = builder.constant({type: 'float32', dimensions: [42, 24]}, inputDenseKernel0Data);
-    const inputDense0 = builder.matmul(input, inputDenseKernel0);
-
-    const inputDenseBias = builder.constant({type: 'float32', dimensions: [24]}, inputDenseBias0Data);
-    const biasedTensorName2 = builder.add(inputDense0, inputDenseBias);
-
-    const inputDenseTanh0 = builder.tanh(biasedTensorName2);
-
-    const vadGruX = builder.transpose(inputDenseTanh0, {permutation: [1, 0, 2]});
-    const vadGruW = builder.constant({type: 'float32', dimensions: [1, 72, 24]}, vadGruWData);
-    const vadGruR = builder.constant({type: 'float32', dimensions: [1, 72, 24]}, vadGruRData);
-    const vadGruB = builder.constant({type: 'float32', dimensions: [1, 72]}, vadGruBData.subarray(0, 72));
-    const vadGruRB = builder.constant({type: 'float32', dimensions: [1, 72]}, vadGruBData.subarray(72, 144));
-    const [, vadGruY] = builder.gru(vadGruX, vadGruW, vadGruR, this.frames_, 24, {bias: vadGruB, recurrentBias: vadGruRB, returnSequence: true, resetAfter: false, activations: ["sigmoid", "relu"]});
-
-    const vadGruYTransposed = builder.transpose(vadGruY, {permutation: [2, 0, 1, 3]});
-
-    const vadGruTranspose1 = builder.reshape(vadGruYTransposed, [-1, 100, 24]);
-
-    const concatenate1 = builder.concat([inputDenseTanh0, vadGruTranspose1, input], 2);
-
-    const noiseGruX = builder.transpose(concatenate1, {permutation: [1, 0, 2]});
-    const noiseGruW = builder.constant({type: 'float32', dimensions: [1, 144, 90]}, noiseGruWData);
-    const noiseGruR = builder.constant({type: 'float32', dimensions: [1, 144, 48]}, noiseGruRData);
-    const noiseGruB = builder.constant({type: 'float32', dimensions: [1, 144]}, noiseGruBData.subarray(0, 144));
-    const noiseGruRB = builder.constant({type: 'float32', dimensions: [1, 144]}, noiseGruBData.subarray(144, 288));
-    const [, noiseGruY] = builder.gru(noiseGruX, noiseGruW, noiseGruR, this.frames_, 48, {bias: noiseGruB, recurrentBias: noiseGruRB, returnSequence: true, resetAfter: false, activations: ["sigmoid", "relu"]});
-    const noiseGruYTransposed = builder.transpose(noiseGruY, {permutation: [2, 0, 1, 3]});
-
-    const noiseGruTranspose1 = builder.reshape(noiseGruYTransposed, [-1, 100, 48]);
-
-    const concatenate2 = builder.concat([vadGruTranspose1, noiseGruTranspose1, input], 2);
-
-    const denoiseGruX = builder.transpose(concatenate2, {permutation: [1, 0, 2]});
-    const denoiseGruW = builder.constant({type: 'float32', dimensions: [1, 288, 114]}, denoiseGruWData);
-    const denoiseGruR = builder.constant({type: 'float32', dimensions: [1, 288, 96]}, denoiseGruRData);
-    const denoiseGruB = builder.constant({type: 'float32', dimensions: [1, 288]}, denoiseGruBData.subarray(0, 288));
-    const denoiseGruRB = builder.constant({type: 'float32', dimensions: [1, 288]}, denoiseGruBData.subarray(288, 576));
-    const [, denoiseGruY] = builder.gru(denoiseGruX, denoiseGruW, denoiseGruR, this.frames_, 96, {bias: denoiseGruB, recurrentBias: denoiseGruRB, returnSequence: true, resetAfter: false, activations: ["sigmoid", "relu"]});
-    const denoiseGruYTransposed = builder.transpose(denoiseGruY, {permutation: [2, 0, 1, 3]});
-
-    const denoiseGruTranspose1 = builder.reshape(denoiseGruYTransposed, [-1, 100, 96]);
-
-    const denoiseOutputKernel0 = builder.constant({type: 'float32', dimensions: [96, 22]}, denoiseOutputKernel0Data);
-    const denoiseOutput0 = builder.matmul(denoiseGruTranspose1, denoiseOutputKernel0);
-
-    const denoiseOutputBias0 = builder.constant({type: 'float32', dimensions: [22]}, denoiseOutputBias0Data);
-    const biasedTensorName = builder.add(denoiseOutput0, denoiseOutputBias0)
-
-    const denoiseOutput = builder.sigmoid(biasedTensorName);
-
-    this.model_ = builder.createModel({'output': denoiseOutput});
+    this.builder = nn.createModelBuilder();
+
+    const inputDenseKernel0 = await this.buildConstantByNpy('input_dense_kernel_0.npy');
+    const inputDenseBias0 = await this.buildConstantByNpy('input_dense_bias_0.npy');
+    const vadGruW = await this.buildConstantByNpy('vad_gru_W.npy');
+    const vadGruR = await this.buildConstantByNpy('vad_gru_R.npy');
+    const vadGruBData = await this.buildConstantByNpy('vad_gru_B.npy');
+    const noiseGruW = await this.buildConstantByNpy('noise_gru_W.npy');
+    const noiseGruR = await this.buildConstantByNpy('noise_gru_R.npy');
+    const noiseGruBData = await this.buildConstantByNpy('noise_gru_B.npy');
+    const denoiseGruW = await this.buildConstantByNpy('denoise_gru_W.npy');
+    const denoiseGruR = await this.buildConstantByNpy('denoise_gru_R.npy');
+    const denoiseGruBData = await this.buildConstantByNpy('denoise_gru_B.npy');
+    const denoiseOutputKernel0 = await this.buildConstantByNpy('denoise_output_kernel_0.npy');
+    const denoiseOutputBias0 = await this.buildConstantByNpy('denoise_output_bias_0.npy');
+
+    const input = this.builder.input('input', {type: 'float32', dimensions: [this.batchSize_, 100, 42]});
+    const inputDense0 = this.builder.matmul(input, inputDenseKernel0);
+    const biasedTensorName2 = this.builder.add(inputDense0, inputDenseBias0);
+    const inputDenseTanh0 = this.builder.tanh(biasedTensorName2);
+    const vadGruX = this.builder.transpose(inputDenseTanh0, {permutation: [1, 0, 2]});
+    //hiddenSize = 24
+    const vadGruB = this.builder.slice(vadGruBData, [0], [3 * 24], {axes: [1]});
+    const vadGruRB = this.builder.slice(vadGruBData, [3 * 24], [-1], {axes: [1]});
+    const vadGruInitialH = this.builder.input('vadGruInitialH', {type: 'float32', dimensions: [1, this.batchSize_, 24]});
+    const [vadGruYH, vadGruY] = this.builder.gru(
+      vadGruX, vadGruW, vadGruR, this.frames_, 24, {
+      bias: vadGruB, recurrentBias: vadGruRB, initialHiddenState: vadGruInitialH,
+      returnSequence: true, resetAfter: false, activations: ["sigmoid", "relu"]
+    });
+    const vadGruYTransposed = this.builder.transpose(vadGruY, {permutation: [2, 0, 1, 3]});
+    const vadGruTranspose1 = this.builder.reshape(vadGruYTransposed, [-1, 100, 24]);
+    const concatenate1 = this.builder.concat([inputDenseTanh0, vadGruTranspose1, input], 2);
+    const noiseGruX = this.builder.transpose(concatenate1, {permutation: [1, 0, 2]});
+    //hiddenSize = 48
+    const noiseGruB = this.builder.slice(noiseGruBData, [0], [3 * 48], {axes: [1]});
+    const noiseGruRB = this.builder.slice(noiseGruBData, [3 * 48], [-1], {axes: [1]});
+    const noiseGruInitialH = this.builder.input('noiseGruInitialH', {type: 'float32', dimensions: [1, this.batchSize_, 48]});
+    const [noiseGruYH, noiseGruY] = this.builder.gru(
+      noiseGruX, noiseGruW, noiseGruR, this.frames_, 48, {
+      bias: noiseGruB, recurrentBias: noiseGruRB, initialHiddenState: noiseGruInitialH,
+      returnSequence: true, resetAfter: false, activations: ["sigmoid", "relu"]
+    });
+    const noiseGruYTransposed = this.builder.transpose(noiseGruY, {permutation: [2, 0, 1, 3]});
+    const noiseGruTranspose1 = this.builder.reshape(noiseGruYTransposed, [-1, 100, 48]);
+    const concatenate2 = this.builder.concat([vadGruTranspose1, noiseGruTranspose1, input], 2);
+    const denoiseGruX = this.builder.transpose(concatenate2, {permutation: [1, 0, 2]});
+    //hiddenSize = 96
+    const denoiseGruB = this.builder.slice(denoiseGruBData, [0], [3 * 96], {axes: [1]});
+    const denoiseGruRB = this.builder.slice(denoiseGruBData, [3 * 96], [-1], {axes: [1]});
+    const denoiseGruInitialH = this.builder.input('denoiseGruInitialH', { type: 'float32', dimensions: [1, this.batchSize_, 96] });
+    const [denoiseGruYH, denoiseGruY] = this.builder.gru(
+      denoiseGruX, denoiseGruW, denoiseGruR, this.frames_, 96, {
+      bias: denoiseGruB, recurrentBias: denoiseGruRB, initialHiddenState: denoiseGruInitialH,
+      returnSequence: true, resetAfter: false, activations: ["sigmoid", "relu"]
+    });
+    const denoiseGruYTransposed = this.builder.transpose(denoiseGruY, {permutation: [2, 0, 1, 3]});
+    const denoiseGruTranspose1 = this.builder.reshape(denoiseGruYTransposed, [-1, 100, 96]);
+    const denoiseOutput0 = this.builder.matmul(denoiseGruTranspose1, denoiseOutputKernel0);
+    const biasedTensorName = this.builder.add(denoiseOutput0, denoiseOutputBias0)
+    const denoiseOutput = this.builder.sigmoid(biasedTensorName);
+
+    this.model_ = this.builder.createModel({denoiseOutput, vadGruYH, noiseGruYH, denoiseGruYH});
   }
 
   async compile(options) {
     this.compilation_ = await this.model_.compile(options);
   }
 
-  async compute(inputBuffer) {
-    const inputs = {input: {buffer: inputBuffer}};
+  async compute(inputBuffer, vadGruInitialHBuffer, noiseGruInitialHBuffer, denoiseGruInitialHBuffer) {
+    const inputs = {
+      input: {buffer: inputBuffer},
+      vadGruInitialH: {buffer: vadGruInitialHBuffer},
+      noiseGruInitialH: {buffer: noiseGruInitialHBuffer},
+      denoiseGruInitialH: {buffer: denoiseGruInitialHBuffer},
+    };
     const outputs = await this.compilation_.compute(inputs);
-    return outputs.output;
+    return outputs;
   }
 
   dispose() {