mirror of
https://github.com/HumanAIGC-Engineering/gradio-webrtc.git
synced 2026-02-04 09:29:23 +08:00
aefb08150f
sync code of fastrtc, add text support through datachannel, fix safari connect problem support chat without camera or mic
330 lines
12 KiB
Python
330 lines
12 KiB
Python
import logging
|
|
import warnings
|
|
from dataclasses import dataclass
|
|
from functools import lru_cache
|
|
from typing import List
|
|
|
|
import click
|
|
import numpy as np
|
|
from huggingface_hub import hf_hub_download
|
|
from numpy.typing import NDArray
|
|
|
|
from ..utils import AudioChunk
|
|
from .protocol import PauseDetectionModel
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# The code below is adapted from https://github.com/snakers4/silero-vad.
|
|
# The code below is adapted from https://github.com/gpt-omni/mini-omni/blob/main/utils/vad.py
|
|
|
|
|
|
@lru_cache
|
|
def get_silero_model() -> PauseDetectionModel:
|
|
"""Returns the VAD model instance and warms it up with dummy data."""
|
|
# Warm up the model with dummy data
|
|
|
|
try:
|
|
import importlib.util
|
|
|
|
mod = importlib.util.find_spec("onnxruntime")
|
|
if mod is None:
|
|
raise RuntimeError("Install fastrtc[vad] to use ReplyOnPause")
|
|
except (ValueError, ModuleNotFoundError):
|
|
raise RuntimeError("Install fastrtc[vad] to use ReplyOnPause")
|
|
model = SileroVADModel()
|
|
print(click.style("INFO", fg="green") + ":\t Warming up VAD model.")
|
|
model.warmup()
|
|
print(click.style("INFO", fg="green") + ":\t VAD model warmed up.")
|
|
return model
|
|
|
|
|
|
@dataclass
|
|
class SileroVadOptions:
|
|
"""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
|
|
speech_duration: If the length of the speech is less than this value, a pause will be detected.
|
|
"""
|
|
|
|
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
|
|
|
|
|
|
class SileroVADModel:
|
|
@staticmethod
|
|
def download_model() -> str:
|
|
return hf_hub_download(
|
|
repo_id="freddyaboulton/silero-vad", filename="silero_vad.onnx"
|
|
)
|
|
|
|
def __init__(self):
|
|
try:
|
|
import onnxruntime
|
|
except ImportError as e:
|
|
raise RuntimeError(
|
|
"Applying the VAD filter requires the onnxruntime package"
|
|
) from e
|
|
|
|
path = self.download_model()
|
|
|
|
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
|
|
|
|
@staticmethod
|
|
def collect_chunks(audio: np.ndarray, chunks: List[AudioChunk]) -> 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]
|
|
)
|
|
|
|
def get_speech_timestamps(
|
|
self,
|
|
audio: np.ndarray,
|
|
vad_options: SileroVadOptions,
|
|
**kwargs,
|
|
) -> List[AudioChunk]:
|
|
"""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.
|
|
"""
|
|
|
|
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)
|
|
|
|
state = self.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 = self(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 warmup(self):
|
|
for _ in range(10):
|
|
dummy_audio = np.zeros(102400, dtype=np.float32)
|
|
self.vad((24000, dummy_audio), None)
|
|
|
|
def vad(
|
|
self,
|
|
audio: tuple[int, NDArray[np.float32] | NDArray[np.int16]],
|
|
options: None | SileroVadOptions,
|
|
) -> tuple[float, list[AudioChunk]]:
|
|
sampling_rate, audio_ = audio
|
|
logger.debug("VAD audio shape input: %s", audio_.shape)
|
|
try:
|
|
if audio_.dtype != np.float32:
|
|
audio_ = audio_.astype(np.float32) / 32768.0
|
|
sr = 16000
|
|
if sr != sampling_rate:
|
|
try:
|
|
import librosa # type: ignore
|
|
except ImportError as e:
|
|
raise RuntimeError(
|
|
"Applying the VAD filter requires the librosa if the input sampling rate is not 16000hz"
|
|
) from e
|
|
audio_ = librosa.resample(audio_, orig_sr=sampling_rate, target_sr=sr)
|
|
|
|
if not options:
|
|
options = SileroVadOptions()
|
|
speech_chunks = self.get_speech_timestamps(audio_, options)
|
|
logger.debug("VAD speech chunks: %s", speech_chunks)
|
|
audio_ = self.collect_chunks(audio_, speech_chunks)
|
|
logger.debug("VAD audio shape: %s", audio_.shape)
|
|
duration_after_vad = audio_.shape[0] / sr
|
|
return duration_after_vad, speech_chunks
|
|
except Exception as e:
|
|
import math
|
|
import traceback
|
|
|
|
logger.debug("VAD Exception: %s", str(e))
|
|
exec = traceback.format_exc()
|
|
logger.debug("traceback %s", exec)
|
|
return math.inf, []
|
|
|
|
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: # type: ignore
|
|
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
|