whisper-translation/engine_transcribe.py

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