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())