wehub-resource-sync
203 lines
8.8 KiB
Python
203 lines
8.8 KiB
Python
import threading, warnings
|
|
|
|
from lib.conf import DEVICE_SYSTEM, systems, tts_dir, devices
|
|
from lib.conf_models import default_voice_detection_model
|
|
|
|
_pipeline_cache = {}
|
|
_pipeline_lock = threading.Lock()
|
|
|
|
def pyannote_patch()->None:
|
|
'''Restore APIs removed in torchaudio >=2.9 that pyannote.audio's
|
|
transitive deps (speechbrain, asteroid-filterbanks, silero-vad, …)
|
|
still call at import time, and route pyannote.audio 4.x's I/O
|
|
through soundfile instead of torchcodec (which is broken on
|
|
Windows ROCm builds and unnecessary for our preloaded-dict path).
|
|
Idempotent: safe to call from both app entrypoint and lazy paths.
|
|
'''
|
|
import torchaudio
|
|
# Silence pyannote.audio.core.io's torchcodec warning. Must run
|
|
# before `from pyannote.audio import ...` triggers io.py import.
|
|
warnings.filterwarnings(
|
|
'ignore',
|
|
message=r'(?s).*torchcodec is not installed correctly.*',
|
|
category=UserWarning,
|
|
)
|
|
if not hasattr(torchaudio, 'list_audio_backends'):
|
|
def _list_audio_backends()->list:
|
|
backends = []
|
|
try:
|
|
from torchaudio.utils import ffmpeg_utils
|
|
if ffmpeg_utils.get_versions():
|
|
backends.append('ffmpeg')
|
|
except Exception:
|
|
pass
|
|
try:
|
|
import soundfile
|
|
backends.append('soundfile')
|
|
except Exception:
|
|
pass
|
|
return backends
|
|
torchaudio.list_audio_backends = _list_audio_backends
|
|
if not hasattr(torchaudio, 'AudioMetaData'):
|
|
class _AudioMetaData:
|
|
def __init__(self, sample_rate:int=0, num_frames:int=0,
|
|
num_channels:int=0, bits_per_sample:int=0,
|
|
encoding:str='UNKNOWN')->None:
|
|
self.sample_rate = sample_rate
|
|
self.num_frames = num_frames
|
|
self.num_channels = num_channels
|
|
self.bits_per_sample = bits_per_sample
|
|
self.encoding = encoding
|
|
torchaudio.AudioMetaData = _AudioMetaData
|
|
# Replace pyannote.audio.core.io.Audio's torchcodec-backed loader
|
|
# with a soundfile-backed one. Preloaded {'waveform','sample_rate'}
|
|
# dicts short-circuit to the original code path (no decoding); any
|
|
# file-path input is decoded by soundfile then handed back to the
|
|
# original as an in-memory dict so its resampling/channel logic
|
|
# still applies.
|
|
try:
|
|
import torch, soundfile as sf
|
|
from pyannote.audio.core import io as _pa_io
|
|
if not getattr(_pa_io, '_e2a_sf_patched', False):
|
|
_orig_call = _pa_io.Audio.__call__
|
|
def _sf_call(self, file, **kw):
|
|
if isinstance(file, dict) and 'waveform' in file:
|
|
return _orig_call(self, file, **kw)
|
|
path = file['audio'] if isinstance(file, dict) else file
|
|
data, sr = sf.read(str(path), dtype='float32', always_2d=True)
|
|
# soundfile -> (frames, channels); pyannote -> (channels, frames)
|
|
waveform = torch.from_numpy(data.T.copy())
|
|
return _orig_call(
|
|
self,
|
|
{'waveform': waveform, 'sample_rate': sr},
|
|
**kw,
|
|
)
|
|
_pa_io.Audio.__call__ = _sf_call
|
|
_pa_io._e2a_sf_patched = True
|
|
except Exception:
|
|
pass
|
|
|
|
class BackgroundDetector:
|
|
def __init__(self, wav_file:str, device:str)->None:
|
|
self.wav_file = wav_file
|
|
self.device = device
|
|
self.torch_device = None
|
|
self.total_duration = self._get_duration()
|
|
|
|
def _get_duration(self)->float:
|
|
try:
|
|
import librosa
|
|
return float(librosa.get_duration(path=self.wav_file))
|
|
except Exception:
|
|
return 0.0
|
|
|
|
def _get_props(self)->tuple:
|
|
import torch, librosa
|
|
pyannote_patch()
|
|
from pyannote.audio import Model
|
|
from pyannote.audio.pipelines import VoiceActivityDetection
|
|
from pyannote.audio.utils.reproducibility import ReproducibilityWarning
|
|
warnings.filterwarnings('ignore', category=ReproducibilityWarning)
|
|
# MIOpen on Windows ROCm fails to JIT pyannote's dropout kernel
|
|
# (miopenStatusUnknownError). On ROCm builds, torch.backends.cudnn
|
|
# maps to MIOpen; disabling it routes ops to native ATen kernels.
|
|
# Cost is negligible for VAD-sized models.
|
|
if getattr(torch.version, 'hip', None) is not None:
|
|
if DEVICE_SYSTEM == systems['WINDOWS']:
|
|
torch.backends.cudnn.enabled = False
|
|
self.torch_device = torch.device(
|
|
devices['CUDA']['proc'] if (self.device == devices['CUDA']['proc'] and getattr(torch.version, 'hip', None) is None) or (getattr(torch.version, 'hip', None) is not None and DEVICE_SYSTEM != systems['WINDOWS'])
|
|
else devices['XPU']['proc'] if hasattr(torch, devices['XPU']['proc']) and torch.xpu.is_available()
|
|
else devices['MPS']['proc'] if torch.backends.mps.is_available()
|
|
else devices['CPU']['proc']
|
|
)
|
|
key = self.torch_device.type
|
|
pipeline = _pipeline_cache.get(key)
|
|
if pipeline is None:
|
|
with _pipeline_lock:
|
|
pipeline = _pipeline_cache.get(key)
|
|
if pipeline is None:
|
|
model = Model.from_pretrained(
|
|
default_voice_detection_model,
|
|
cache_dir=tts_dir
|
|
)
|
|
model.eval()
|
|
batch_size = 1 if devices['JETSON']['found'] else 32
|
|
pipeline = VoiceActivityDetection(segmentation=model, batch_size=batch_size)
|
|
if key == devices['CUDA']['proc'] and not devices['JETSON']['found']:
|
|
torch.backends.cuda.matmul.allow_tf32 = True
|
|
torch.backends.cudnn.allow_tf32 = True
|
|
pipeline.instantiate({
|
|
"min_duration_on": 0.0,
|
|
"min_duration_off": 0.0
|
|
})
|
|
pipeline.to(self.torch_device)
|
|
_pipeline_cache[key] = pipeline
|
|
if pipeline:
|
|
y, sr = librosa.load(self.wav_file, sr=16000, mono=True)
|
|
waveform = torch.from_numpy(y).float().unsqueeze(0)
|
|
return pipeline, waveform, sr
|
|
return None, None, None
|
|
|
|
def detect(self, vad_ratio_thresh:float=0.05)->tuple[bool, dict[str, float|bool]]:
|
|
import gc, torch
|
|
pipeline, waveform, sr = self._get_props()
|
|
try:
|
|
if (
|
|
pipeline is not None
|
|
and waveform is not None
|
|
and waveform.numel() > 0
|
|
and sr is not None
|
|
and sr > 0
|
|
):
|
|
file = {
|
|
"waveform": waveform,
|
|
"sample_rate": sr
|
|
}
|
|
with torch.inference_mode():
|
|
annotation = pipeline(file)
|
|
speech_time = sum(
|
|
segment.end - segment.start
|
|
for segment in annotation.itersegments()
|
|
)
|
|
non_speech_ratio = 1.0 - (
|
|
speech_time / self.total_duration if self.total_duration > 0 else 0.0
|
|
)
|
|
background_detected = non_speech_ratio > vad_ratio_thresh
|
|
return background_detected, {
|
|
"non_speech_ratio": non_speech_ratio,
|
|
"background_detected": background_detected
|
|
}
|
|
return False, {}
|
|
finally:
|
|
# Park cached pipelines on CPU so they don't hold accelerator memory
|
|
# while Demucs (or other heavy models) loads later in extract_voice().
|
|
# Keep them in _pipeline_cache so subsequent detect() calls can reuse
|
|
# them without reloading from disk; _get_props() is expected to call
|
|
# .to(self.torch_device) on retrieval.
|
|
with _pipeline_lock:
|
|
for p in _pipeline_cache.values():
|
|
try:
|
|
p.to('cpu')
|
|
except Exception:
|
|
pass
|
|
# Drop local refs to free any short-lived tensors / annotations
|
|
pipeline = waveform = sr = None
|
|
gc.collect()
|
|
if torch.cuda.is_available():
|
|
try:
|
|
torch.cuda.empty_cache()
|
|
torch.cuda.ipc_collect()
|
|
except Exception:
|
|
pass
|
|
if hasattr(torch, 'xpu') and torch.xpu.is_available():
|
|
try:
|
|
torch.xpu.empty_cache()
|
|
except Exception:
|
|
pass
|
|
if torch.backends.mps.is_available():
|
|
try:
|
|
torch.mps.empty_cache()
|
|
except Exception:
|
|
pass
|