Move VAD to On-Device

app.py

@@ -12,10 +12,6 @@ import xxhash
 from datasets import Audio
 from transformers import AutoModel
 import io
-from pydub import AudioSegment
-import tempfile
-
-from utils.vad import VadOptions, collect_chunks, get_speech_timestamps
 
 if gr.NO_RELOAD:
     diva_model = AutoModel.from_pretrained(
@@ -25,7 +21,7 @@ if gr.NO_RELOAD:
     resampler = Audio(sampling_rate=16_000)
 
 
-@spaces.GPU(duration=20)
+@spaces.GPU(duration=20, progress=gr.Progress(track_tqdm=True))
 @torch.no_grad
 def diva_audio(audio_input, do_sample=False, temperature=0.001, prev_outs=None):
     sr, y = audio_input
@@ -37,7 +33,11 @@ def diva_audio(audio_input, do_sample=False, temperature=0.001, prev_outs=None):
     )
     yield from diva_model.generate_stream(
         a["array"],
-        None,
+        (
+            "Your name is DiVA, which stands for Distilled Voice Assistant. You were trained with early-fusion training to merge OpenAI's Whisper and Meta AI's Llama 3 8B to provide end-to-end voice processing. You should give brief and helpful answers, in a conversational style. The user is talking to you with their voice and you are responding with text."
+            if prev_outs == None
+            else None
+        ),
         do_sample=do_sample,
         max_new_tokens=256,
         init_outputs=prev_outs,
@@ -45,96 +45,24 @@ def diva_audio(audio_input, do_sample=False, temperature=0.001, prev_outs=None):
     )
 
 
-def run_vad(ori_audio, sr, duration):
-    _st = time.time()
-    try:
-        audio = ori_audio
-        if duration < 1:
-            return -1, ori_audio, round(time.time() - _st, 4)
-        audio = audio.astype(np.float32) / 32768.0
-        sampling_rate = 16000
-        if sr != sampling_rate:
-            audio = librosa.resample(audio, orig_sr=sr, target_sr=sampling_rate)
-
-        vad_parameters = {}
-        vad_parameters = VadOptions(**vad_parameters)
-        speech_chunks = get_speech_timestamps(audio, vad_parameters)
-        audio = collect_chunks(audio, speech_chunks)
-        duration_after_vad = audio.shape[0] / sampling_rate
-
-        if sr != sampling_rate:
-            # resample to original sampling rate
-            vad_audio = librosa.resample(audio, orig_sr=sampling_rate, target_sr=sr)
-        else:
-            vad_audio = audio
-        vad_audio = np.round(vad_audio * 32768.0).astype(np.int16)
-        vad_audio_bytes = vad_audio.tobytes()
-
-        return duration_after_vad, vad_audio_bytes, round(time.time() - _st, 4)
-    except Exception as e:
-        msg = f"[asr vad error] audio_len: {len(ori_audio)/(sr*2):.3f} s, trace: {traceback.format_exc()}"
-        print(msg)
-        return -1, ori_audio, round(time.time() - _st, 4)
-
-
-def warm_up():
-    frames = np.ones(2048)  # 1024 frames of 2 bytes each
-    dur, frames, tcost = run_vad(frames, 16000, 10)
-    print(f"warm up done, time_cost: {tcost:.3f} s")
-
-
-warm_up()
-
-
 @dataclass
 class AppState:
     stream: np.ndarray | None = None
     sampling_rate: int = 0
-    pause_detected: bool = False
-    started_talking: bool = False
     stopped: bool = False
     conversation: list = field(default_factory=list)
     model_outs: any = None
 
 
-def determine_pause(audio: np.ndarray, sampling_rate: int, state: AppState) -> bool:
-    """Take in the stream, determine if a pause happened"""
-
-    temp_audio = audio[-2 * sampling_rate :]
-
-    duration = len(audio) / sampling_rate
-    dur_vad, _, time_vad = run_vad(temp_audio, sampling_rate, duration)
-
-    if dur_vad > 0.25 and not state.started_talking:
-        print("started talking")
-        state.started_talking = True
-        return False
-
-    print(f"duration_after_vad: {dur_vad:.3f} s, time_vad: {time_vad:.3f} s")
-
-    return dur_vad < 0.5
-
-
 def process_audio(audio: tuple, state: AppState):
-    if state.stream is None:
-        state.stream = audio[1]
-        state.sampling_rate = audio[0]
-    elif audio is not None and audio[1] is not None:
-        state.stream = np.concatenate((state.stream, audio[1]))
-    else:
-        return None, state
+    return audio, state
 
-    pause_detected = determine_pause(state.stream, state.sampling_rate, state)
-    state.pause_detected = pause_detected
 
-    if state.pause_detected and state.started_talking:
-        return gr.Audio(recording=False), state
-    return None, state
-
-
-def response(state: AppState):
-    if not state.pause_detected and not state.started_talking:
+def response(state: AppState, audio: tuple):
+    if not audio:
         return AppState()
+    state.stream = audio[1]
+    state.sampling_rate = audio[0]
 
     file_name = f"/tmp/{xxhash.xxh32(bytes(state.stream)).hexdigest()}.wav"
 
@@ -159,8 +87,7 @@ def response(state: AppState):
 
 
 def start_recording_user(state: AppState):
-    if not state.stopped:
-        return gr.Audio(recording=True)
+    return None
 
 
 theme = gr.themes.Soft(
@@ -181,29 +108,80 @@ theme = gr.themes.Soft(
     neutral_hue="stone",
 )
 
-with gr.Blocks(theme=theme) as demo:
+js = """
+async function main() {
+  const script1 = document.createElement("script");
+  script1.src = "https://cdn.jsdelivr.net/npm/onnxruntime-web@1.14.0/dist/ort.js";
+  document.head.appendChild(script1)
+  const script2 = document.createElement("script");
+  script2.onload = async () =>  {
+    console.log("vad loaded") ;
+    var record = document.querySelector('.record-button');
+    record.textContent = "Just Start Talking!"
+    record.style = "width: 11vw"
+    const myvad = await vad.MicVAD.new({
+      onSpeechStart: () => {
+        var record = document.querySelector('.record-button');
+        if (record != null) {
+          console.log(record);
+          record.click();
+        }
+      },
+      onSpeechEnd: (audio) => {
+        var stop = document.querySelector('.stop-button');
+        if (stop != null) {
+          console.log(stop);
+          stop.click();
+        }
+      }
+    })
+    myvad.start()
+  }
+  script2.src = "https://cdn.jsdelivr.net/npm/@ricky0123/vad-web@0.0.7/dist/bundle.min.js";
+  script1.onload = () =>  {
+    console.log("onnx loaded") 
+    document.head.appendChild(script2)
+  };
+}
+"""
+
+js_reset = """
+() => {
+  var record = document.querySelector('.record-button');
+  record.textContent = "Just Start Talking!"
+  record.style = "width: 11vw"
+}
+"""
+
+with gr.Blocks(theme=theme, js=js) as demo:
     with gr.Row():
-        input_audio = gr.Audio(label="Input Audio", sources="microphone", type="numpy")
+        input_audio = gr.Audio(
+            label="Input Audio",
+            sources=["microphone"],
+            type="numpy",
+            streaming=False,
+        )
     with gr.Row():
         chatbot = gr.Chatbot(label="Conversation", type="messages")
     state = gr.State(value=AppState())
-
-    stream = input_audio.stream(
+    stream = input_audio.start_recording(
         process_audio,
         [input_audio, state],
         [input_audio, state],
-        stream_every=0.25,
-        time_limit=10,
     )
-    respond = input_audio.stop_recording(response, [state], [state, chatbot])
-    respond.then(start_recording_user, [state], [input_audio])
+    respond = input_audio.stop_recording(
+        response, [state, input_audio], [state, chatbot]
+    )
+    restart = respond.then(start_recording_user, [state], [input_audio]).then(
+        lambda: None, None, None, js=js_reset
+    )
 
-    cancel = gr.Button("Stop Conversation", variant="stop")
+    cancel = gr.Button("Restart Conversation", variant="stop")
     cancel.click(
         lambda: (AppState(stopped=True), gr.Audio(recording=False)),
         None,
         [state, input_audio],
-        cancels=[respond, stream],
+        cancels=[respond, restart],
     )
 
 if __name__ == "__main__":

utils/assets/silero_vad.onnx

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:591f853590d11ddde2f2a54f9e7ccecb2533a8af7716330e8adfa6f3849787a9
-size 1807524

utils/snac_utils.py

@@ -1,146 +0,0 @@
-import torch
-import time
-import numpy as np
-
-
-class SnacConfig:
-    audio_vocab_size = 4096
-    padded_vocab_size = 4160
-    end_of_audio = 4097
-
-
-snac_config = SnacConfig()    
-
-
-def get_time_str():
-    time_str = time.strftime("%Y%m%d_%H%M%S", time.localtime())
-    return time_str
-
-
-def layershift(input_id, layer, stride=4160, shift=152000):
-    return input_id + shift + layer * stride
-
-    
-def generate_audio_data(snac_tokens, snacmodel, device=None):
-    audio = reconstruct_tensors(snac_tokens, device)
-    with torch.inference_mode():
-        audio_hat = snacmodel.decode(audio)
-    audio_data = audio_hat.cpu().numpy().astype(np.float64) * 32768.0
-    audio_data = audio_data.astype(np.int16)
-    audio_data = audio_data.tobytes()
-    return audio_data
-
-    
-def get_snac(list_output, index, nums_generate):
-
-    snac = []
-    start = index
-    for i in range(nums_generate):
-        snac.append("#")
-        for j in range(7):
-            snac.append(list_output[j][start - nums_generate - 5 + j + i])
-    return snac
-
-
-def reconscruct_snac(output_list):
-    if len(output_list) == 8:
-        output_list = output_list[:-1]
-    output = []
-    for i in range(7):
-        output_list[i] = output_list[i][i + 1 :]
-    for i in range(len(output_list[-1])):
-        output.append("#")
-        for j in range(7):
-            output.append(output_list[j][i])
-    return output
-
-
-def reconstruct_tensors(flattened_output, device=None):
-    """Reconstructs the list of tensors from the flattened output."""
-
-    if device is None:
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-    def count_elements_between_hashes(lst):
-        try:
-            # Find the index of the first '#'
-            first_index = lst.index("#")
-            # Find the index of the second '#' after the first
-            second_index = lst.index("#", first_index + 1)
-            # Count the elements between the two indices
-            return second_index - first_index - 1
-        except ValueError:
-            # Handle the case where there aren't enough '#' symbols
-            return "List does not contain two '#' symbols"
-
-    def remove_elements_before_hash(flattened_list):
-        try:
-            # Find the index of the first '#'
-            first_hash_index = flattened_list.index("#")
-            # Return the list starting from the first '#'
-            return flattened_list[first_hash_index:]
-        except ValueError:
-            # Handle the case where there is no '#'
-            return "List does not contain the symbol '#'"
-
-    def list_to_torch_tensor(tensor1):
-        # Convert the list to a torch tensor
-        tensor = torch.tensor(tensor1)
-        # Reshape the tensor to have size (1, n)
-        tensor = tensor.unsqueeze(0)
-        return tensor
-
-    flattened_output = remove_elements_before_hash(flattened_output)
-    codes = []
-    tensor1 = []
-    tensor2 = []
-    tensor3 = []
-    tensor4 = []
-
-    n_tensors = count_elements_between_hashes(flattened_output)
-    if n_tensors == 7:
-        for i in range(0, len(flattened_output), 8):
-
-            tensor1.append(flattened_output[i + 1])
-            tensor2.append(flattened_output[i + 2])
-            tensor3.append(flattened_output[i + 3])
-            tensor3.append(flattened_output[i + 4])
-
-            tensor2.append(flattened_output[i + 5])
-            tensor3.append(flattened_output[i + 6])
-            tensor3.append(flattened_output[i + 7])
-            codes = [
-                list_to_torch_tensor(tensor1).to(device),
-                list_to_torch_tensor(tensor2).to(device),
-                list_to_torch_tensor(tensor3).to(device),
-            ]
-
-    if n_tensors == 15:
-        for i in range(0, len(flattened_output), 16):
-
-            tensor1.append(flattened_output[i + 1])
-            tensor2.append(flattened_output[i + 2])
-            tensor3.append(flattened_output[i + 3])
-            tensor4.append(flattened_output[i + 4])
-            tensor4.append(flattened_output[i + 5])
-            tensor3.append(flattened_output[i + 6])
-            tensor4.append(flattened_output[i + 7])
-            tensor4.append(flattened_output[i + 8])
-
-            tensor2.append(flattened_output[i + 9])
-            tensor3.append(flattened_output[i + 10])
-            tensor4.append(flattened_output[i + 11])
-            tensor4.append(flattened_output[i + 12])
-            tensor3.append(flattened_output[i + 13])
-            tensor4.append(flattened_output[i + 14])
-            tensor4.append(flattened_output[i + 15])
-
-            codes = [
-                list_to_torch_tensor(tensor1).to(device),
-                list_to_torch_tensor(tensor2).to(device),
-                list_to_torch_tensor(tensor3).to(device),
-                list_to_torch_tensor(tensor4).to(device),
-            ]
-
-    return codes
-

utils/vad.py

@@ -1,290 +0,0 @@
-import bisect
-import functools
-import os
-import warnings
-
-from typing import List, NamedTuple, Optional
-
-import numpy as np
-
-
-# The code below is adapted from https://github.com/snakers4/silero-vad.
-class VadOptions(NamedTuple):
-    """VAD options.
-
-    Attributes:
-      threshold: Speech threshold. Silero VAD outputs speech probabilities for each audio chunk,
-        probabilities ABOVE this value are considered as SPEECH. It is better to tune this
-        parameter for each dataset separately, but "lazy" 0.5 is pretty good for most datasets.
-      min_speech_duration_ms: Final speech chunks shorter min_speech_duration_ms are thrown out.
-      max_speech_duration_s: Maximum duration of speech chunks in seconds. Chunks longer
-        than max_speech_duration_s will be split at the timestamp of the last silence that
-        lasts more than 100ms (if any), to prevent aggressive cutting. Otherwise, they will be
-        split aggressively just before max_speech_duration_s.
-      min_silence_duration_ms: In the end of each speech chunk wait for min_silence_duration_ms
-        before separating it
-      window_size_samples: Audio chunks of window_size_samples size are fed to the silero VAD model.
-        WARNING! Silero VAD models were trained using 512, 1024, 1536 samples for 16000 sample rate.
-        Values other than these may affect model performance!!
-      speech_pad_ms: Final speech chunks are padded by speech_pad_ms each side
-    """
-
-    threshold: float = 0.5
-    min_speech_duration_ms: int = 250
-    max_speech_duration_s: float = float("inf")
-    min_silence_duration_ms: int = 2000
-    window_size_samples: int = 1024
-    speech_pad_ms: int = 400
-
-
-def get_speech_timestamps(
-    audio: np.ndarray,
-    vad_options: Optional[VadOptions] = None,
-    **kwargs,
-) -> List[dict]:
-    """This method is used for splitting long audios into speech chunks using silero VAD.
-
-    Args:
-      audio: One dimensional float array.
-      vad_options: Options for VAD processing.
-      kwargs: VAD options passed as keyword arguments for backward compatibility.
-
-    Returns:
-      List of dicts containing begin and end samples of each speech chunk.
-    """
-    if vad_options is None:
-        vad_options = VadOptions(**kwargs)
-
-    threshold = vad_options.threshold
-    min_speech_duration_ms = vad_options.min_speech_duration_ms
-    max_speech_duration_s = vad_options.max_speech_duration_s
-    min_silence_duration_ms = vad_options.min_silence_duration_ms
-    window_size_samples = vad_options.window_size_samples
-    speech_pad_ms = vad_options.speech_pad_ms
-
-    if window_size_samples not in [512, 1024, 1536]:
-        warnings.warn(
-            "Unusual window_size_samples! Supported window_size_samples:\n"
-            " - [512, 1024, 1536] for 16000 sampling_rate"
-        )
-
-    sampling_rate = 16000
-    min_speech_samples = sampling_rate * min_speech_duration_ms / 1000
-    speech_pad_samples = sampling_rate * speech_pad_ms / 1000
-    max_speech_samples = (
-        sampling_rate * max_speech_duration_s
-        - window_size_samples
-        - 2 * speech_pad_samples
-    )
-    min_silence_samples = sampling_rate * min_silence_duration_ms / 1000
-    min_silence_samples_at_max_speech = sampling_rate * 98 / 1000
-
-    audio_length_samples = len(audio)
-
-    model = get_vad_model()
-    state = model.get_initial_state(batch_size=1)
-
-    speech_probs = []
-    for current_start_sample in range(0, audio_length_samples, window_size_samples):
-        chunk = audio[current_start_sample : current_start_sample + window_size_samples]
-        if len(chunk) < window_size_samples:
-            chunk = np.pad(chunk, (0, int(window_size_samples - len(chunk))))
-        speech_prob, state = model(chunk, state, sampling_rate)
-        speech_probs.append(speech_prob)
-
-    triggered = False
-    speeches = []
-    current_speech = {}
-    neg_threshold = threshold - 0.15
-
-    # to save potential segment end (and tolerate some silence)
-    temp_end = 0
-    # to save potential segment limits in case of maximum segment size reached
-    prev_end = next_start = 0
-
-    for i, speech_prob in enumerate(speech_probs):
-        if (speech_prob >= threshold) and temp_end:
-            temp_end = 0
-            if next_start < prev_end:
-                next_start = window_size_samples * i
-
-        if (speech_prob >= threshold) and not triggered:
-            triggered = True
-            current_speech["start"] = window_size_samples * i
-            continue
-
-        if (
-            triggered
-            and (window_size_samples * i) - current_speech["start"] > max_speech_samples
-        ):
-            if prev_end:
-                current_speech["end"] = prev_end
-                speeches.append(current_speech)
-                current_speech = {}
-                # previously reached silence (< neg_thres) and is still not speech (< thres)
-                if next_start < prev_end:
-                    triggered = False
-                else:
-                    current_speech["start"] = next_start
-                prev_end = next_start = temp_end = 0
-            else:
-                current_speech["end"] = window_size_samples * i
-                speeches.append(current_speech)
-                current_speech = {}
-                prev_end = next_start = temp_end = 0
-                triggered = False
-                continue
-
-        if (speech_prob < neg_threshold) and triggered:
-            if not temp_end:
-                temp_end = window_size_samples * i
-            # condition to avoid cutting in very short silence
-            if (window_size_samples * i) - temp_end > min_silence_samples_at_max_speech:
-                prev_end = temp_end
-            if (window_size_samples * i) - temp_end < min_silence_samples:
-                continue
-            else:
-                current_speech["end"] = temp_end
-                if (
-                    current_speech["end"] - current_speech["start"]
-                ) > min_speech_samples:
-                    speeches.append(current_speech)
-                current_speech = {}
-                prev_end = next_start = temp_end = 0
-                triggered = False
-                continue
-
-    if (
-        current_speech
-        and (audio_length_samples - current_speech["start"]) > min_speech_samples
-    ):
-        current_speech["end"] = audio_length_samples
-        speeches.append(current_speech)
-
-    for i, speech in enumerate(speeches):
-        if i == 0:
-            speech["start"] = int(max(0, speech["start"] - speech_pad_samples))
-        if i != len(speeches) - 1:
-            silence_duration = speeches[i + 1]["start"] - speech["end"]
-            if silence_duration < 2 * speech_pad_samples:
-                speech["end"] += int(silence_duration // 2)
-                speeches[i + 1]["start"] = int(
-                    max(0, speeches[i + 1]["start"] - silence_duration // 2)
-                )
-            else:
-                speech["end"] = int(
-                    min(audio_length_samples, speech["end"] + speech_pad_samples)
-                )
-                speeches[i + 1]["start"] = int(
-                    max(0, speeches[i + 1]["start"] - speech_pad_samples)
-                )
-        else:
-            speech["end"] = int(
-                min(audio_length_samples, speech["end"] + speech_pad_samples)
-            )
-
-    return speeches
-
-
-def collect_chunks(audio: np.ndarray, chunks: List[dict]) -> np.ndarray:
-    """Collects and concatenates audio chunks."""
-    if not chunks:
-        return np.array([], dtype=np.float32)
-
-    return np.concatenate([audio[chunk["start"] : chunk["end"]] for chunk in chunks])
-
-
-class SpeechTimestampsMap:
-    """Helper class to restore original speech timestamps."""
-
-    def __init__(self, chunks: List[dict], sampling_rate: int, time_precision: int = 2):
-        self.sampling_rate = sampling_rate
-        self.time_precision = time_precision
-        self.chunk_end_sample = []
-        self.total_silence_before = []
-
-        previous_end = 0
-        silent_samples = 0
-
-        for chunk in chunks:
-            silent_samples += chunk["start"] - previous_end
-            previous_end = chunk["end"]
-
-            self.chunk_end_sample.append(chunk["end"] - silent_samples)
-            self.total_silence_before.append(silent_samples / sampling_rate)
-
-    def get_original_time(
-        self,
-        time: float,
-        chunk_index: Optional[int] = None,
-    ) -> float:
-        if chunk_index is None:
-            chunk_index = self.get_chunk_index(time)
-
-        total_silence_before = self.total_silence_before[chunk_index]
-        return round(total_silence_before + time, self.time_precision)
-
-    def get_chunk_index(self, time: float) -> int:
-        sample = int(time * self.sampling_rate)
-        return min(
-            bisect.bisect(self.chunk_end_sample, sample),
-            len(self.chunk_end_sample) - 1,
-        )
-
-
-@functools.lru_cache
-def get_vad_model():
-    """Returns the VAD model instance."""
-    asset_dir = os.path.join(os.path.dirname(__file__), "assets")
-    path = os.path.join(asset_dir, "silero_vad.onnx")
-    return SileroVADModel(path)
-
-
-class SileroVADModel:
-    def __init__(self, path):
-        try:
-            import onnxruntime
-        except ImportError as e:
-            raise RuntimeError(
-                "Applying the VAD filter requires the onnxruntime package"
-            ) from e
-
-        opts = onnxruntime.SessionOptions()
-        opts.inter_op_num_threads = 1
-        opts.intra_op_num_threads = 1
-        opts.log_severity_level = 4
-
-        self.session = onnxruntime.InferenceSession(
-            path,
-            providers=["CPUExecutionProvider"],
-            sess_options=opts,
-        )
-
-    def get_initial_state(self, batch_size: int):
-        h = np.zeros((2, batch_size, 64), dtype=np.float32)
-        c = np.zeros((2, batch_size, 64), dtype=np.float32)
-        return h, c
-
-    def __call__(self, x, state, sr: int):
-        if len(x.shape) == 1:
-            x = np.expand_dims(x, 0)
-        if len(x.shape) > 2:
-            raise ValueError(
-                f"Too many dimensions for input audio chunk {len(x.shape)}"
-            )
-        if sr / x.shape[1] > 31.25:
-            raise ValueError("Input audio chunk is too short")
-
-        h, c = state
-
-        ort_inputs = {
-            "input": x,
-            "h": h,
-            "c": c,
-            "sr": np.array(sr, dtype="int64"),
-        }
-
-        out, h, c = self.session.run(None, ort_inputs)
-        state = (h, c)
-
-        return out, state