Adolfo Reyna
207 lines
10 KiB
Python
207 lines
10 KiB
Python
import sys
|
|
from unittest.mock import MagicMock
|
|
|
|
# MANDATORY: Mock lzma BEFORE any other imports (especially pyannote/torchvision)
|
|
try:
|
|
import lzma
|
|
except ImportError:
|
|
mock_lzma = MagicMock()
|
|
mock_lzma.FORMAT_XZ, mock_lzma.FORMAT_ALONE, mock_lzma.FORMAT_RAW = 1, 2, 3
|
|
mock_lzma.CHECK_NONE, mock_lzma.CHECK_CRC32, mock_lzma.CHECK_CRC64, mock_lzma.CHECK_SHA256 = 0, 1, 4, 10
|
|
sys.modules["_lzma"] = MagicMock()
|
|
sys.modules["lzma"] = mock_lzma
|
|
|
|
import time
|
|
import os
|
|
import argparse
|
|
|
|
def list_audio_devices():
|
|
try:
|
|
import sounddevice as sd
|
|
print("\nAvailable Audio Devices:")
|
|
print(sd.query_devices())
|
|
except ImportError:
|
|
print("[Error] sounddevice not installed. Cannot list devices.")
|
|
|
|
def assign_speakers_to_segments(segments, diarization_result):
|
|
if not diarization_result or not segments:
|
|
return segments
|
|
|
|
# pyannote 4.0.4 returns a DiarizeOutput object with a speaker_diarization attribute
|
|
annotation = None
|
|
if hasattr(diarization_result, 'speaker_diarization'):
|
|
annotation = diarization_result.speaker_diarization
|
|
elif hasattr(diarization_result, 'itertracks'):
|
|
annotation = diarization_result
|
|
|
|
if annotation is None:
|
|
return segments
|
|
|
|
for segment in segments:
|
|
segment_start, segment_end = segment.get('start', 0), segment.get('end', 0)
|
|
speaker_intersections = {}
|
|
|
|
try:
|
|
for turn, _, speaker in annotation.itertracks(yield_label=True):
|
|
intersection_start = max(segment_start, turn.start)
|
|
intersection_end = min(segment_end, turn.end)
|
|
if intersection_end > intersection_start:
|
|
duration = intersection_end - intersection_start
|
|
speaker_intersections[speaker] = speaker_intersections.get(speaker, 0) + duration
|
|
except Exception:
|
|
continue
|
|
|
|
if speaker_intersections:
|
|
segment['speaker'] = max(speaker_intersections, key=speaker_intersections.get)
|
|
else:
|
|
segment['speaker'] = 'UNKNOWN'
|
|
return segments
|
|
|
|
def run_transcription(out_queue, args):
|
|
import numpy as np
|
|
import sounddevice as sd
|
|
import torch
|
|
import mlx_whisper
|
|
from silero_vad import load_silero_vad, get_speech_timestamps
|
|
import queue
|
|
|
|
SAMPLERATE, BLOCK_SIZE, VAD_THRESHOLD = 16000, 512, 0.5
|
|
|
|
def transcribe_with_controls(audio, transcribe_kwargs):
|
|
try:
|
|
return mlx_whisper.transcribe(audio, **transcribe_kwargs)
|
|
except Exception as exc:
|
|
if "beam_size" in transcribe_kwargs:
|
|
transcribe_kwargs.pop("beam_size")
|
|
return mlx_whisper.transcribe(audio, **transcribe_kwargs)
|
|
raise exc
|
|
|
|
device = "mps" if torch.backends.mps.is_available() else "cpu"
|
|
model_name = args.model
|
|
if args.quantize and "4bit" not in model_name:
|
|
model_name = "mlx-community/whisper-small-mlx-4bit"
|
|
|
|
print(f"[Transcribe] Loading Whisper model '{model_name}' on {device}...")
|
|
vad_model = load_silero_vad()
|
|
|
|
diarization_pipeline = None
|
|
if args.speaker_diarization:
|
|
from pyannote.audio import Pipeline
|
|
print("[Transcribe] Loading speaker diarization pipeline...")
|
|
hf_token = os.environ.get("HF_TOKEN")
|
|
try:
|
|
diarization_pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization-3.1", token=hf_token)
|
|
diarization_pipeline.to(torch.device(device))
|
|
except Exception as e:
|
|
print(f"[Transcribe] Diarization load failed: {e}")
|
|
|
|
audio_queue = queue.Queue()
|
|
def audio_callback(indata, frames, time, status):
|
|
if status: print(status, file=sys.stderr)
|
|
audio_queue.put(indata.copy())
|
|
|
|
audio_buffer, speech_started, rolling_context, rolling_context_en = [], False, "", ""
|
|
last_stream_time = time.time()
|
|
|
|
try:
|
|
with sd.InputStream(samplerate=SAMPLERATE, channels=args.channels, callback=audio_callback, blocksize=BLOCK_SIZE, device=args.device):
|
|
while True:
|
|
while not audio_queue.empty():
|
|
data = audio_queue.get()
|
|
if data is not None and data.size > 0:
|
|
if args.channels > 1:
|
|
data = np.mean(data, axis=1)
|
|
audio_buffer.append(data.flatten())
|
|
|
|
if audio_buffer:
|
|
current_audio = np.concatenate(audio_buffer)
|
|
audio_tensor = torch.from_numpy(current_audio)
|
|
|
|
speech_timestamps = get_speech_timestamps(audio_tensor, vad_model, sampling_rate=SAMPLERATE, threshold=VAD_THRESHOLD, min_silence_duration_ms=args.silence)
|
|
|
|
if speech_timestamps:
|
|
speech_started = True
|
|
if (len(current_audio) - speech_timestamps[-1]['end']) > (SAMPLERATE * args.silence / 1000) or len(current_audio) > (SAMPLERATE * args.max_buffer):
|
|
|
|
# STEP 1: ALWAYS TRANSCRIBE FIRST
|
|
transcribe_kwargs = {
|
|
"path_or_hf_repo": model_name,
|
|
"temperature": args.temperature_fallback,
|
|
"word_timestamps": args.speaker_diarization,
|
|
"logprob_threshold": args.logprob_threshold,
|
|
"compression_ratio_threshold": args.compression_threshold,
|
|
}
|
|
if args.lang: transcribe_kwargs["language"] = args.lang
|
|
if args.context and rolling_context: transcribe_kwargs["initial_prompt"] = rolling_context
|
|
|
|
res_asr = transcribe_with_controls(current_audio, transcribe_kwargs)
|
|
text = res_asr['text'].strip()
|
|
detected_lang = res_asr.get('language', args.lang or 'en')
|
|
|
|
if text:
|
|
if args.filter_lang and args.lang and detected_lang != args.lang:
|
|
print(f"[Transcribe] Discarding segment (Detected: {detected_lang}, Expected: {args.lang})")
|
|
audio_buffer, speech_started = [], False
|
|
continue
|
|
|
|
english_text = text
|
|
if detected_lang != "en":
|
|
bridge_kwargs = {
|
|
"path_or_hf_repo": model_name,
|
|
"task": "translate",
|
|
"temperature": args.temperature_fallback,
|
|
"logprob_threshold": args.logprob_threshold,
|
|
"compression_ratio_threshold": args.compression_threshold,
|
|
}
|
|
if args.lang:
|
|
bridge_kwargs["language"] = args.lang
|
|
if args.context and rolling_context_en:
|
|
bridge_kwargs["initial_prompt"] = rolling_context_en
|
|
res_bridge = transcribe_with_controls(current_audio, bridge_kwargs)
|
|
english_text = res_bridge.get("text", "").strip() or text
|
|
|
|
speaker_label = None
|
|
if diarization_pipeline:
|
|
try:
|
|
diar_out = diarization_pipeline({"waveform": torch.from_numpy(current_audio).float().unsqueeze(0), "sample_rate": SAMPLERATE})
|
|
asr_segments = res_asr.get('segments', [])
|
|
diarized_segments = assign_speakers_to_segments(asr_segments, diar_out)
|
|
|
|
speakers = [s.get('speaker') for s in diarized_segments if s.get('speaker') and s.get('speaker') != 'UNKNOWN']
|
|
if speakers:
|
|
from collections import Counter
|
|
speaker_label = Counter(speakers).most_common(1)[0][0]
|
|
except Exception as diar_err:
|
|
print(f"[Transcribe] Diarization error: {diar_err}")
|
|
|
|
out_queue.put({
|
|
"original": text,
|
|
"en_bridge": english_text,
|
|
"detected_lang": detected_lang,
|
|
"speaker": speaker_label,
|
|
"ts": time.time()
|
|
})
|
|
|
|
speaker_tag = f" ({speaker_label})" if speaker_label else ""
|
|
print(f"[Transcribe] {detected_lang.upper()}{speaker_tag}: {text}")
|
|
if args.context:
|
|
rolling_context = (rolling_context + " " + text)[-200:].strip()
|
|
rolling_context_en = (rolling_context_en + " " + english_text)[-200:].strip()
|
|
|
|
audio_buffer, speech_started = [], False
|
|
|
|
elif args.stream and (time.time() - last_stream_time) > 1.5:
|
|
draft_result = transcribe_with_controls(current_audio, {"path_or_hf_repo": model_name, "temperature": 0.0})
|
|
if draft_result['text'].strip():
|
|
out_queue.put({"draft": draft_result['text'].strip(), "ts": time.time()})
|
|
last_stream_time = time.time()
|
|
|
|
elif not speech_started and len(current_audio) > SAMPLERATE * 2:
|
|
audio_buffer = []
|
|
|
|
except Exception as e: print(f"[Transcribe] Error: {e}")
|
|
|
|
if __name__ == "__main__":
|
|
import multiprocessing
|
|
run_transcription(multiprocessing.Queue(), argparse.Namespace())
|