import sys from unittest.mock import MagicMock # Comprehensive workaround for missing _lzma in some Python builds try: import lzma except ImportError: mock_lzma = MagicMock() mock_lzma.FORMAT_XZ = 1 mock_lzma.FORMAT_ALONE = 2 mock_lzma.FORMAT_RAW = 3 mock_lzma.CHECK_NONE = 0 mock_lzma.CHECK_CRC32 = 1 mock_lzma.CHECK_CRC64 = 4 mock_lzma.CHECK_SHA256 = 10 sys.modules["_lzma"] = MagicMock() sys.modules["lzma"] = mock_lzma import time import os import numpy as np import sounddevice as sd import torch import mlx_whisper from silero_vad import load_silero_vad, get_speech_timestamps from pyannote.audio import Pipeline import queue import argparse # Audio Constants SAMPLERATE = 16000 BLOCK_SIZE = 512 VAD_THRESHOLD = 0.5 def transcribe_with_controls(audio, transcribe_kwargs): """Call mlx_whisper.transcribe with graceful fallback for unsupported args.""" try: return mlx_whisper.transcribe(audio, **transcribe_kwargs) except Exception as exc: message = str(exc) if "beam_size" in transcribe_kwargs: transcribe_kwargs.pop("beam_size") return mlx_whisper.transcribe(audio, **transcribe_kwargs) raise exc def audio_callback(indata, frames, time, status, audio_queue): if status: print(status, file=sys.stderr) audio_queue.put(indata.copy()) def assign_speakers_to_segments(segments, diarization): if not diarization or not segments: return segments for segment in segments: segment_start, segment_end = segment['start'], segment['end'] speaker_intersections = {} for turn, _, speaker in diarization.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 if speaker_intersections: segment['speaker'] = max(speaker_intersections, key=speaker_intersections.get) else: segment['speaker'] = 'UNKNOWN' return segments def run_transcription(out_queue, args): 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: print("[Transcribe] Loading speaker diarization pipeline...") hf_token = os.environ.get("HF_TOKEN") try: diarization_pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization-3.1", use_auth_token=hf_token) diarization_pipeline.to(torch.device(device)) except Exception as e: print(f"[Transcribe] Diarization load failed: {e}") audio_queue = queue.Queue() device_index = args.device def callback_wrapper(indata, frames, time, status): audio_callback(indata, frames, time, status, audio_queue) audio_buffer = [] speech_started = False buffer_limit = SAMPLERATE * args.max_buffer rolling_context = "" last_stream_time = time.time() print(f"[Transcribe] Audio stream ready on device {device_index}.") try: with sd.InputStream(samplerate=SAMPLERATE, channels=args.channels, callback=callback_wrapper, blocksize=BLOCK_SIZE, device=device_index): while True: while not audio_queue.empty(): data = audio_queue.get() 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) buffer_duration = len(current_audio) / SAMPLERATE 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 last_end = speech_timestamps[-1]['end'] buffer_len_samples = len(current_audio) # Check for flush if (buffer_len_samples - last_end) > (SAMPLERATE * args.silence / 1000) or buffer_len_samples > buffer_limit: 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 result = transcribe_with_controls(current_audio, transcribe_kwargs) text = result['text'].strip() detected_lang = result.get('language', args.lang or 'en') if args.filter_lang and args.lang and detected_lang != args.lang: print(f"[Transcribe] Filtered {detected_lang}") text = "" if text: segments = result.get('segments', []) if diarization_pipeline: audio_for_diarization = torch.from_numpy(current_audio).float().unsqueeze(0) try: diarization_result = diarization_pipeline({"waveform": audio_for_diarization, "sample_rate": SAMPLERATE}) segments = assign_speakers_to_speakers(segments, diarization_result) except: pass out_queue.put({"original": text, "detected_lang": detected_lang, "segments": segments, "ts": time.time()}) print(f"[Transcribe] {detected_lang.upper()}: {text}") if args.context: rolling_context = (rolling_context + " " + text)[-200:].strip() audio_buffer = [] speech_started = False elif args.stream and (time.time() - last_stream_time) > 1.5: # Draft Mode draft_kwargs = {"path_or_hf_repo": model_name, "temperature": 0.0} if args.lang: draft_kwargs["language"] = args.lang if args.context and rolling_context: draft_kwargs["initial_prompt"] = rolling_context draft_result = transcribe_with_controls(current_audio, draft_kwargs) draft_text = draft_result['text'].strip() if draft_text: out_queue.put({"draft": draft_text, "ts": time.time()}) last_stream_time = time.time() elif not speech_started and len(current_audio) > SAMPLERATE * 2: audio_buffer = [] except KeyboardInterrupt: pass except Exception as e: print(f"[Transcribe] Error: {e}") if __name__ == "__main__": import multiprocessing run_transcription(multiprocessing.Queue(), argparse.Namespace())