Spaces:
Running
on
CPU Upgrade
Running
on
CPU Upgrade
File size: 6,774 Bytes
178417b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 |
from dataclasses import dataclass
from typing import Callable, Generator, override
import fastrtc
import librosa
import numpy as np
from ten_vad import TenVad
@dataclass
class VADEvent:
interrupt_signal: bool | None = None
full_audio: tuple[int, np.ndarray] | None = None
class RealtimeVAD:
def __init__(
self,
src_sr: int = 24000,
hop_size: int = 256,
start_threshold: float = 0.8,
end_threshold: float = 0.7,
pad_start_s: float = 0.6,
min_positive_s: float = 0.4,
min_silence_s: float = 1.2,
):
self.src_sr = src_sr
self.vad_sr = 16000
self.hop_size = hop_size
self.start_threshold = start_threshold
self.end_threshold = end_threshold
self.pad_start_s = pad_start_s
self.min_positive_s = min_positive_s
self.min_silence_s = min_silence_s
self.vad_model = TenVad(hop_size=hop_size)
self.vad_buffer = np.array([], dtype=np.int16)
"""
VAD Buffer to store audio data for VAD processing
Stores 16kHz int16 PCM. Process and cut for each `hop_size` samples.
"""
self.src_buffer = np.array([], dtype=np.int16)
"""
Source Buffer to store original audio data
Stores original sampling rate (24kHz) int16 PCM.
Cut when pause detected (after `min_silence_s`).
Sliding window `pad_start_s` when inactive.
"""
self.vad_buffer_offset = 0
self.src_buffer_offset = 0
self.active = False
self.interrupt_signal = False
self.sum_positive_s = 0.0
self.silence_start_s: float | None = None
def process(self, audio_data: np.ndarray):
if audio_data.ndim == 2:
# FastRTC style [channels, samples]
audio_data = audio_data[0]
# Append to buffers
self.src_buffer = np.concatenate((self.src_buffer, audio_data))
vad_audio_data = librosa.resample(
audio_data.astype(np.float32) / 32768.0,
orig_sr=self.src_sr,
target_sr=self.vad_sr,
)
vad_audio_data = (vad_audio_data * 32767.0).round().astype(np.int16)
self.vad_buffer = np.concatenate((self.vad_buffer, vad_audio_data))
vad_buffer_size = self.vad_buffer.shape[0]
def process_chunk(chunk_offset_s: float, vad_chunk: np.ndarray):
speech_prob, _ = self.vad_model.process(vad_chunk)
hop_s = self.hop_size / self.vad_sr
if not self.active:
if speech_prob >= self.start_threshold:
self.active = True
self.sum_positive_s = hop_s
print(f"[VAD] Active at {chunk_offset_s:.2f}s, {speech_prob=:.3f}")
else:
new_src_offset = int(
(chunk_offset_s - self.pad_start_s) * self.src_sr
)
cut_pos = new_src_offset - self.src_buffer_offset
if cut_pos > 0:
self.src_buffer = self.src_buffer[cut_pos:]
self.src_buffer_offset = new_src_offset
return
chunk_src_pos = int(chunk_offset_s * self.src_sr)
if speech_prob >= self.end_threshold:
self.silence_start_s = None
self.sum_positive_s += hop_s
if (
not self.interrupt_signal
and self.sum_positive_s >= self.min_positive_s
):
self.interrupt_signal = True
yield VADEvent(interrupt_signal=True)
print(
f"[VAD] Interrupt signal at {chunk_offset_s:.2f}s, {speech_prob=:.3f}"
)
elif self.silence_start_s is None:
self.silence_start_s = chunk_offset_s
if (
self.silence_start_s is not None
and chunk_offset_s - self.silence_start_s >= self.min_silence_s
):
# Inactive now
cut_pos = chunk_src_pos - self.src_buffer_offset
if self.interrupt_signal:
webrtc_audio = self.src_buffer[np.newaxis, :cut_pos]
yield VADEvent(full_audio=(self.src_sr, webrtc_audio))
print(
f"[VAD] Full audio at {chunk_offset_s:.2f}s, {webrtc_audio.shape=}"
)
self.src_buffer = self.src_buffer[cut_pos:]
self.src_buffer_offset = chunk_src_pos
self.active = False
self.interrupt_signal = False
self.sum_positive_s = 0.0
self.silence_start_s = None
for chunk_pos in range(0, vad_buffer_size - self.hop_size, self.hop_size):
processed_samples = chunk_pos + self.hop_size
chunk_offset_s = (self.vad_buffer_offset + chunk_pos) / self.vad_sr
vad_chunk = self.vad_buffer[chunk_pos : chunk_pos + self.hop_size]
yield from process_chunk(chunk_offset_s, vad_chunk)
self.vad_buffer = self.vad_buffer[processed_samples:]
self.vad_buffer_offset += processed_samples
type StreamerGenerator = Generator[fastrtc.tracks.EmitType, None, None]
type StreamerFn = Callable[[tuple[int, np.ndarray], str], StreamerGenerator]
class VADStreamHandler(fastrtc.StreamHandler):
def __init__(
self,
streamer_fn: StreamerFn,
input_sample_rate: int = 24000,
):
super().__init__(
"mono",
24000,
None,
input_sample_rate,
30,
)
self.streamer_fn = streamer_fn
self.realtime_vad = RealtimeVAD(src_sr=input_sample_rate)
self.generator: StreamerGenerator | None = None
@override
def emit(self) -> fastrtc.tracks.EmitType:
if self.generator is None:
return None
try:
return next(self.generator)
except StopIteration:
self.generator = None
return None
@override
def receive(self, frame: tuple[int, np.ndarray]):
_, audio_data = frame
for event in self.realtime_vad.process(audio_data):
if event.interrupt_signal:
self.generator = None
self.clear_queue()
if event.full_audio is not None:
self.wait_for_args_sync()
self.latest_args[0] = event.full_audio
self.generator = self.streamer_fn(*self.latest_args)
@override
def copy(self):
return VADStreamHandler(
self.streamer_fn,
input_sample_rate=self.input_sample_rate,
)
|