Support VoxLock 1.2 models in the speech denoisers (#4120)

.github/scripts/test-offline-speech-denoiser.sh

@@ -35,7 +35,7 @@ rm ./gtcrn_simple.onnx
 log "------------------------------------------------------------"
 log "Run voxlock"
 log "------------------------------------------------------------"
-voxlock_url=https://github.com/arbazsiddiqui/voxlock/releases/download/sherpa-onnx-models
+voxlock_url=https://github.com/arbazsiddiqui/voxlock/releases/download/sherpa-onnx-models-v2
 download $voxlock_url/voxlock_enroll.onnx
 download $voxlock_url/voxlock_chunk_int8.onnx
 download $voxlock_url/voxlock_test_call.wav

scripts/voxlock/README.md

@@ -2,7 +2,7 @@
 
 This folder contains scripts for adding meta data to the
 [VoxLock](https://github.com/arbazsiddiqui/voxlock) ONNX models from
-<https://huggingface.co/arbazsiddiqui/voxlock/tree/main/onnx>.
+<https://huggingface.co/arbazsiddiqui/voxlock/tree/main/onnx/v2>.
 
 VoxLock is a target-speaker voice isolation model. It keeps the voice of one
 speaker, e.g., the caller in a phone call, and removes the other voices.
@@ -21,7 +21,14 @@ It has two models:
 | fp32 | `voxlock_enroll.onnx` | `voxlock_chunk.onnx` |
 
 The models with meta data can be downloaded from
-<https://github.com/arbazsiddiqui/voxlock/releases/tag/sherpa-onnx-models>
+<https://github.com/arbazsiddiqui/voxlock/releases/tag/sherpa-onnx-models-v2>
+
+These are the models of VoxLock 1.2.0. They give the same output as the
+models of VoxLock 1.1.1
+(<https://github.com/arbazsiddiqui/voxlock/releases/tag/sherpa-onnx-models>),
+which still work. The 1.2.0 models use less CPU: their enroll model also
+computes everything that depends only on the voice sample, instead of the chunk
+model computing it again for every chunk.
 
 - [./run.sh](./run.sh) downloads the models and runs
   [./add_meta_data.py](./add_meta_data.py), which also checks that the outputs

scripts/voxlock/add_meta_data.py

@@ -3,20 +3,25 @@
 
 """
 Add meta data to the VoxLock ONNX models from
-https://huggingface.co/arbazsiddiqui/voxlock/tree/main/onnx
+https://huggingface.co/arbazsiddiqui/voxlock/tree/main/onnx/v2
 
-enroll.onnx runs once on a voice sample of the caller:
+enroll.onnx runs once on a voice sample of the caller. Besides the normalized
+magnitude of the sample, it returns the keys and values of the speaker
+attention for every band, so the chunk model does not recompute them for
+every chunk:
 
-NodeArg(name='enroll', type='tensor(float)', shape=[1, 'Te'])
+NodeArg(name='enroll', type='tensor(float)', shape=[1, 48000])
 -----
-NodeArg(name='enroll_mag', type='tensor(float)', shape=[1, 257, 'Fe'])
-NodeArg(name='context', type='tensor(float)', shape=[1, 512, 'Fc'])
+NodeArg(name='enroll_n', type='tensor(float)', shape=[257, 376])
+NodeArg(name='k', type='tensor(float)', shape=[32, 2, 64, 298])
+NodeArg(name='v', type='tensor(float)', shape=[32, 2, 298, 64])
 
 chunk.onnx runs on each chunk of the call and carries the LSTM state:
 
 NodeArg(name='wav', type='tensor(float)', shape=[1, 'S'])
-NodeArg(name='enroll_mag', type='tensor(float)', shape=[1, 257, 'Fe'])
-NodeArg(name='context', type='tensor(float)', shape=[1, 512, 'Fc'])
+NodeArg(name='enroll_n', type='tensor(float)', shape=[257, 376])
+NodeArg(name='k', type='tensor(float)', shape=[32, 2, 64, 298])
+NodeArg(name='v', type='tensor(float)', shape=[32, 2, 298, 64])
 NodeArg(name='h', type='tensor(float)', shape=[6, 32, 256])
 NodeArg(name='c', type='tensor(float)', shape=[6, 32, 256])
 -----
@@ -36,8 +41,8 @@ import numpy as np
 import onnx
 import onnxruntime as ort
 
-# Weights of VoxLock 1.1.1
-HF_REVISION = "a50f4e1d109301d4b6f61dec57ebb1a61e7d8f03"
+# Models of VoxLock 1.2.0
+HF_REVISION = "2413b292bc2c337bfc19a9802aaf3e011e68e9c4"
 
 # (input file, output file, role, precision)
 MODELS = [
@@ -71,8 +76,8 @@ def get_meta_data(role: str, precision: str):
     meta_data = {
         "model_type": "voxlock",
         "model_role": role,
-        "version": 1,
-        "model_version": "1.1.1",
+        "version": 2,
+        "model_version": "1.2.0",
         "precision": precision,
         "sample_rate": 16000,
         "n_fft": 512,
@@ -80,7 +85,6 @@ def get_meta_data(role: str, precision: str):
         "window_length": 512,
         "window_type": "hann",
         "num_freq_bins": 257,
-        "context_dim": 512,
         # The chunk model accepts only the outputs of the enroll model for a
         # voice sample of exactly this length
         "voice_sample_seconds": 3,
@@ -118,25 +122,35 @@ def add_meta_data(in_filename: str, out_filename: str, meta_data):
     onnx.save(model, out_filename)
 
 
-def make_inputs(sess, rng):
-    """Random inputs for a session; enroll_mag and context come from enroll."""
+def make_inputs(sess, rng, enroll_filename: str):
+    """Random inputs for a session. For a chunk model, enroll_n, k and v are
+    the outputs of the enroll model on a random voice sample."""
+    opts = ort.SessionOptions()
+    opts.log_severity_level = 3
+    enroll = ort.InferenceSession(
+        enroll_filename, opts, providers=["CPUExecutionProvider"]
+    )
+    voice_sample = {
+        "enroll": (rng.standard_normal((1, 48000)) * 0.1).astype(np.float32)
+    }
+
     names = [i.name for i in sess.get_inputs()]
     if names == ["enroll"]:
-        return {"enroll": (rng.standard_normal((1, 48000)) * 0.1).astype(np.float32)}
+        return voice_sample
 
     num_frames = 100
+    voice = enroll.run(None, voice_sample)
     return {
         "wav": (rng.standard_normal((1, (num_frames - 1) * 128 + 512)) * 0.1).astype(
             np.float32
         ),
-        "enroll_mag": np.abs(rng.standard_normal((1, 257, 376))).astype(np.float32),
-        "context": rng.standard_normal((1, 512, 298)).astype(np.float32),
+        **{o.name: x for o, x in zip(enroll.get_outputs(), voice)},
         "h": (rng.standard_normal((6, 32, 256)) * 0.1).astype(np.float32),
         "c": (rng.standard_normal((6, 32, 256)) * 0.1).astype(np.float32),
     }
 
 
-def check(in_filename: str, out_filename: str, meta_data):
+def check(in_filename: str, out_filename: str, meta_data, enroll_filename: str):
     opts = ort.SessionOptions()
     opts.log_severity_level = 3
     opts.intra_op_num_threads = 1
@@ -147,7 +161,7 @@ def check(in_filename: str, out_filename: str, meta_data):
     for key, value in meta_data.items():
         assert got[key] == str(value), (key, got[key], value)
 
-    inputs = make_inputs(a, np.random.default_rng(0))
+    inputs = make_inputs(a, np.random.default_rng(0), enroll_filename)
     for x, y in zip(a.run(None, inputs), b.run(None, inputs)):
         assert np.array_equal(x, y), out_filename
 
@@ -163,7 +177,9 @@ def main():
         meta_data = get_meta_data(role, precision)
 
         add_meta_data(in_filename, out_filename, meta_data)
-        check(in_filename, out_filename, meta_data)
+        enroll_name = "enroll_fp16.onnx" if precision == "fp16" else "enroll.onnx"
+        enroll_filename = str(Path(args.in_dir) / enroll_name)
+        check(in_filename, out_filename, meta_data, enroll_filename)
         print(f"{in_filename} -> {out_filename}: meta data added, outputs unchanged")
 
 

scripts/voxlock/run.sh

@@ -3,8 +3,8 @@
 
 set -ex
 
-# Weights of VoxLock 1.1.1
-revision=a50f4e1d109301d4b6f61dec57ebb1a61e7d8f03
+# Models of VoxLock 1.2.0
+revision=2413b292bc2c337bfc19a9802aaf3e011e68e9c4
 
 models=(
 enroll.onnx
@@ -15,14 +15,15 @@ enroll_fp16.onnx
 chunk_fp16.onnx
 )
 
-mkdir -p onnx
+onnx_dir=onnx/$revision
+mkdir -p $onnx_dir
 for m in ${models[@]}; do
-  if [ ! -f onnx/$m ]; then
-    curl -fSL --output onnx/$m.tmp https://huggingface.co/arbazsiddiqui/voxlock/resolve/$revision/onnx/$m
-    mv onnx/$m.tmp onnx/$m
+  if [ ! -f $onnx_dir/$m ]; then
+    curl -fSL --output $onnx_dir/$m.tmp https://huggingface.co/arbazsiddiqui/voxlock/resolve/$revision/onnx/v2/$m
+    mv $onnx_dir/$m.tmp $onnx_dir/$m
   fi
 done
 
-python3 ./add_meta_data.py --in-dir ./onnx --out-dir ./models
+python3 ./add_meta_data.py --in-dir ./$onnx_dir --out-dir ./models
 
 ls -lh models

scripts/voxlock/test.py

@@ -3,7 +3,7 @@
 
 """
 Compare the output of sherpa-onnx-offline-denoiser with the reference
-implementation in the voxlock Python package (pip install voxlock==1.1.1).
+implementation in the voxlock Python package (pip install voxlock==1.2.1).
 
 ./build/bin/sherpa-onnx-offline-denoiser \
   --speech-denoiser-voxlock-enroll-model=./models/voxlock_enroll.onnx \

sherpa-onnx/csrc/offline-speech-denoiser-voxlock-model.cc

@@ -112,8 +112,7 @@ class OfflineSpeechDenoiserVoxLockModel::Impl {
     }
   }
 
-  std::pair<Ort::Value, Ort::Value> Enroll(const float *samples,
-                                           int32_t n) const {
+  Voice Enroll(const float *samples, int32_t n) const {
     if (n != meta_.voice_sample_seconds * meta_.sample_rate) {
       SHERPA_ONNX_LOGE(
           "The voxlock voice sample must have exactly %d samples (%d seconds "
@@ -136,7 +135,7 @@ class OfflineSpeechDenoiserVoxLockModel::Impl {
                                  1, enroll_output_names_ptr_.data(),
                                  enroll_output_names_ptr_.size());
 
-    return {std::move(out[0]), std::move(out[1])};
+    return out;
   }
 
   States GetInitStates() {
@@ -157,13 +156,14 @@ class OfflineSpeechDenoiserVoxLockModel::Impl {
     return states;
   }
 
-  std::pair<Ort::Value, States> Run(Ort::Value wav, Ort::Value enroll_mag,
-                                    Ort::Value context, States states) const {
+  std::pair<Ort::Value, States> Run(Ort::Value wav, Voice voice,
+                                    States states) const {
     std::vector<Ort::Value> inputs;
-    inputs.reserve(3 + states.size());
+    inputs.reserve(1 + voice.size() + states.size());
     inputs.push_back(std::move(wav));
-    inputs.push_back(std::move(enroll_mag));
-    inputs.push_back(std::move(context));
+    for (auto &v : voice) {
+      inputs.push_back(std::move(v));
+    }
     for (auto &s : states) {
       inputs.push_back(std::move(s));
     }
@@ -204,10 +204,12 @@ class OfflineSpeechDenoiserVoxLockModel::Impl {
     auto shared =
         ReadSharedMetaData(meta_data, "enroll", config_.voxlock.enroll_model);
 
-    if (enroll_input_names_.size() != 1 || enroll_output_names_.size() != 2) {
+    // version 1 outputs: enroll_mag, context
+    // version 2 outputs: enroll_n, k, v
+    if (enroll_input_names_.size() != 1 || enroll_output_names_.size() < 2) {
       SHERPA_ONNX_LOGE(
-          "Expect the voxlock enroll model to have 1 input and 2 outputs. "
-          "Given %d inputs and %d outputs",
+          "Expect the voxlock enroll model to have 1 input and at least 2 "
+          "outputs. Given %d inputs and %d outputs",
           static_cast<int32_t>(enroll_input_names_.size()),
           static_cast<int32_t>(enroll_output_names_.size()));
       SHERPA_ONNX_EXIT(-1);
@@ -256,17 +258,32 @@ class OfflineSpeechDenoiserVoxLockModel::Impl {
     auto shared =
         ReadSharedMetaData(meta_data, "chunk", config_.voxlock.chunk_model);
 
-    // inputs: wav, enroll_mag, context, h, c
+    // inputs: wav, the outputs of the enroll model, h, c
     // outputs: frames, h_out, c_out
-    if (chunk_input_names_.size() != 5 || chunk_output_names_.size() != 3) {
+    int32_t num_voice = static_cast<int32_t>(enroll_output_names_.size());
+    if (static_cast<int32_t>(chunk_input_names_.size()) != num_voice + 3 ||
+        chunk_output_names_.size() != 3) {
       SHERPA_ONNX_LOGE(
-          "Expect the voxlock chunk model to have 5 inputs and 3 outputs. "
-          "Given %d inputs and %d outputs",
-          static_cast<int32_t>(chunk_input_names_.size()),
+          "Expect the voxlock chunk model to have %d inputs and 3 outputs. "
+          "Given %d inputs and %d outputs. Are the enroll and chunk models "
+          "from the same release?",
+          num_voice + 3, static_cast<int32_t>(chunk_input_names_.size()),
           static_cast<int32_t>(chunk_output_names_.size()));
       SHERPA_ONNX_EXIT(-1);
     }
 
+    for (int32_t i = 0; i != num_voice; ++i) {
+      if (chunk_input_names_[i + 1] != enroll_output_names_[i]) {
+        SHERPA_ONNX_LOGE(
+            "Input %d of the voxlock chunk model is '%s'. Expect '%s', output "
+            "%d of the enroll model. Are the enroll and chunk models from the "
+            "same release?",
+            i + 1, chunk_input_names_[i + 1].c_str(),
+            enroll_output_names_[i].c_str(), i);
+        SHERPA_ONNX_EXIT(-1);
+      }
+    }
+
     if (shared.version != meta_.version ||
         shared.sample_rate != meta_.sample_rate ||
         shared.n_fft != meta_.n_fft || shared.hop_length != meta_.hop_length ||
@@ -310,10 +327,12 @@ class OfflineSpeechDenoiserVoxLockModel::Impl {
       SHERPA_ONNX_EXIT(-1);
     }
 
-    auto h_shape =
-        chunk_sess_->GetInputTypeInfo(3).GetTensorTypeAndShapeInfo().GetShape();
-    auto c_shape =
-        chunk_sess_->GetInputTypeInfo(4).GetTensorTypeAndShapeInfo().GetShape();
+    auto h_shape = chunk_sess_->GetInputTypeInfo(num_voice + 1)
+                       .GetTensorTypeAndShapeInfo()
+                       .GetShape();
+    auto c_shape = chunk_sess_->GetInputTypeInfo(num_voice + 2)
+                       .GetTensorTypeAndShapeInfo()
+                       .GetShape();
     if (h_shape != meta_.state_shape || c_shape != meta_.state_shape) {
       SHERPA_ONNX_LOGE(
           "The state shape in the voxlock chunk model does not match its meta "
@@ -387,8 +406,9 @@ OfflineSpeechDenoiserVoxLockModel::OfflineSpeechDenoiserVoxLockModel(
     Manager *mgr, const OfflineSpeechDenoiserModelConfig &config)
     : impl_(std::make_unique<Impl>(mgr, config)) {}
 
-std::pair<Ort::Value, Ort::Value> OfflineSpeechDenoiserVoxLockModel::Enroll(
-    const float *samples, int32_t n) const {
+OfflineSpeechDenoiserVoxLockModel::Voice
+OfflineSpeechDenoiserVoxLockModel::Enroll(const float *samples,
+                                          int32_t n) const {
   return impl_->Enroll(samples, n);
 }
 
@@ -398,11 +418,9 @@ OfflineSpeechDenoiserVoxLockModel::GetInitStates() const {
 }
 
 std::pair<Ort::Value, OfflineSpeechDenoiserVoxLockModel::States>
-OfflineSpeechDenoiserVoxLockModel::Run(Ort::Value wav, Ort::Value enroll_mag,
-                                       Ort::Value context,
+OfflineSpeechDenoiserVoxLockModel::Run(Ort::Value wav, Voice voice,
                                        States states) const {
-  return impl_->Run(std::move(wav), std::move(enroll_mag), std::move(context),
-                    std::move(states));
+  return impl_->Run(std::move(wav), std::move(voice), std::move(states));
 }
 
 const OfflineSpeechDenoiserVoxLockModelMetaData &

sherpa-onnx/csrc/offline-speech-denoiser-voxlock-model.h

@@ -26,18 +26,21 @@ class OfflineSpeechDenoiserVoxLockModel {
 
   using States = std::vector<Ort::Value>;
 
+  // The outputs of the enroll model, in the order the chunk model takes them
+  using Voice = std::vector<Ort::Value>;
+
   /* Run the enroll model on a voice sample of the caller.
    *
    * @param samples Voice sample at GetMetaData().sample_rate. It must have
    *                exactly GetMetaData().voice_sample_seconds seconds.
    * @param n Number of samples.
    *
-   * @return Return a pair containing:
-   *  - enroll_mag, a tensor of shape (1, num_freq_bins, num_frames)
-   *  - context, a tensor of shape (1, context_dim, num_context_frames)
+   * @return Return the outputs of the enroll model. Version 1 models return
+   *         enroll_mag and context; version 2 models return enroll_n, k and v,
+   *         which also hold the speaker attention's keys and values, so the
+   *         chunk model does not recompute them for every chunk.
    */
-  std::pair<Ort::Value, Ort::Value> Enroll(const float *samples,
-                                           int32_t n) const;
+  Voice Enroll(const float *samples, int32_t n) const;
 
   // Zero LSTM states h and c
   States GetInitStates() const;
@@ -45,8 +48,7 @@ class OfflineSpeechDenoiserVoxLockModel {
   /* Run the chunk model.
    *
    * @param wav A tensor of shape (1, (K - 1) * hop_length + n_fft)
-   * @param enroll_mag Returned by Enroll()
-   * @param context Returned by Enroll()
+   * @param voice Returned by Enroll()
    * @param states LSTM states from GetInitStates() or a previous call
    *
    * @return Return a pair containing:
@@ -54,8 +56,8 @@ class OfflineSpeechDenoiserVoxLockModel {
    *    the synthesis window and is ready for overlap-add.
    *  - next states
    */
-  std::pair<Ort::Value, States> Run(Ort::Value wav, Ort::Value enroll_mag,
-                                    Ort::Value context, States states) const;
+  std::pair<Ort::Value, States> Run(Ort::Value wav, Voice voice,
+                                    States states) const;
 
   const OfflineSpeechDenoiserVoxLockModelMetaData &GetMetaData() const;
 

sherpa-onnx/csrc/online-speech-denoiser-voxlock-impl.h

@@ -30,9 +30,11 @@ namespace sherpa_onnx {
  *
  * Speech is 10 ms frames at least 10 dB above the noise floor, the quietest
  * frame of the last 2 seconds. Digital silence (below -90 dBFS) is left out
- * of the floor. At the switch, the stream first runs over the last 250 ms of
- * audio, whose output is dropped, so it does not start from an empty LSTM
- * state, and the isolated audio fades in over 10 ms.
+ * of the floor. At the switch, the stream first runs over the audio so far
+ * (up to 30 seconds), whose output is dropped, and the isolated audio fades
+ * in over 10 ms. A stream started mid-speech from an empty LSTM state can
+ * take the caller for background and keep them muted for the rest of the
+ * call; one that has heard the call from its start does not.
  */
 class OnlineSpeechDenoiserVoxLockImpl : public OnlineSpeechDenoiserImpl {
  public:
@@ -188,7 +190,7 @@ class OnlineSpeechDenoiserVoxLockImpl : public OnlineSpeechDenoiserImpl {
     const auto &meta = model_.GetMetaData();
 
     frame_length_ = meta.sample_rate / 100;
-    warm_up_length_ = meta.sample_rate / 4;
+    warm_up_length_ = 30 * meta.sample_rate;
     fade_length_ = frame_length_;
     num_frames_needed_ = meta.voice_sample_seconds * 100;
 
@@ -213,7 +215,8 @@ class OnlineSpeechDenoiserVoxLockImpl : public OnlineSpeechDenoiserImpl {
 
   void FindVoiceSample(const float *samples, int32_t n) {
     recent_.insert(recent_.end(), samples, samples + n);
-    if (static_cast<int32_t>(recent_.size()) > warm_up_length_) {
+    // Trim only when twice as long, so the buffer is not moved on every call
+    if (static_cast<int32_t>(recent_.size()) > 2 * warm_up_length_) {
       recent_.erase(recent_.begin(), recent_.end() - warm_up_length_);
     }
 
@@ -258,8 +261,13 @@ class OnlineSpeechDenoiserVoxLockImpl : public OnlineSpeechDenoiserImpl {
     voiced_.clear();
     energy_.clear();
 
+    if (static_cast<int32_t>(recent_.size()) > warm_up_length_) {
+      recent_.erase(recent_.begin(), recent_.end() - warm_up_length_);
+    }
+
     num_to_skip_ = recent_.size();
     auto warm_up = stream_->Push(recent_.data(), recent_.size());
+    recent_.clear();
     num_to_skip_ -= std::min<int32_t>(num_to_skip_, warm_up.size());
   }
 

sherpa-onnx/csrc/speech-denoiser-voxlock-stream.h

@@ -201,7 +201,7 @@ class SpeechDenoiserVoxLockStream {
         window_square_(n_fft_),
         overlap_add_(n_fft_),
         window_sum_(n_fft_) {
-    std::tie(enroll_mag_, context_) = model_->Enroll(voice_sample, n);
+    voice_ = model_->Enroll(voice_sample, n);
     states_ = model_->GetInitStates();
 
     constexpr double kPi = 3.14159265358979323846;
@@ -288,10 +288,15 @@ class SpeechDenoiserVoxLockStream {
           memory_info_, buffer_.data() + offset_, num_samples, wav_shape.data(),
           wav_shape.size());
 
+      OfflineSpeechDenoiserVoxLockModel::Voice voice;
+      voice.reserve(voice_.size());
+      for (auto &v : voice_) {
+        voice.push_back(View(&v));
+      }
+
       Ort::Value frames{nullptr};
       std::tie(frames, states_) =
-          model_->Run(std::move(wav), View(&enroll_mag_), View(&context_),
-                      std::move(states_));
+          model_->Run(std::move(wav), std::move(voice), std::move(states_));
 
       if (frames.GetTensorTypeAndShapeInfo().GetElementCount() !=
           static_cast<size_t>(num_frames) * n_fft_) {
@@ -347,8 +352,7 @@ class SpeechDenoiserVoxLockStream {
   float output_gain_;
   Ort::MemoryInfo memory_info_;
 
-  Ort::Value enroll_mag_{nullptr};
-  Ort::Value context_{nullptr};
+  OfflineSpeechDenoiserVoxLockModel::Voice voice_;
   OfflineSpeechDenoiserVoxLockModel::States states_;
 
   std::vector<float> window_square_;