diff --git a/transcribe.py b/transcribe.py index 42cd622..f921f8b 100644 --- a/transcribe.py +++ b/transcribe.py @@ -51,20 +51,23 @@ def main(): device = "mps" if torch.backends.mps.is_available() else "cpu" print(f"Using device: {device}") + # 1. Load models first print(f"Loading Whisper model '{WHISPER_MODEL}'...") - # Dictionary to hold models and tokenizers translation_engines = {} - for lang_name, model_id in TARGET_LANGS.items(): print(f"Loading {lang_name} translation model ({model_id})...") tokenizer = MarianTokenizer.from_pretrained(model_id) model = MarianMTModel.from_pretrained(model_id).to(device) translation_engines[lang_name] = (model, tokenizer) + print("Loading Silero VAD model...") + vad_model = load_silero_vad() + print("Models loaded.") + + # 2. Select Audio Device print("\nAvailable Audio Devices:") - devices = sd.query_devices() - print(devices) + print(sd.query_devices()) try: device_input = input("\nSelect input device index (or press Enter for default): ") @@ -79,6 +82,7 @@ def main(): speech_started = False try: + # 3. Start stream with sd.InputStream(samplerate=SAMPLERATE, channels=CHANNELS, callback=callback, blocksize=BLOCK_SIZE, device=device_index): while True: while not audio_queue.empty(): @@ -104,22 +108,20 @@ def main(): if (buffer_len_samples - last_end) > (SAMPLERATE * MIN_SILENCE_DURATION_MS / 1000) or buffer_len_samples > BUFFER_LIMIT: - # 1. Transcribe once + # Transcribe once result = mlx_whisper.transcribe(current_audio, path_or_hf_repo=WHISPER_MODEL) original_text = result['text'].strip() if original_text: - # Limit the input length to avoid memory spikes or model glitches if len(original_text) > 250: original_text = original_text[:247] + "..." print(f"\n[EN]: {original_text}") - # 2. Translate to all targets + # Translate to all targets for lang_name, (model, tokenizer) in translation_engines.items(): inputs = tokenizer(original_text, return_tensors="pt", padding=True).to(device) with torch.no_grad(): - # Added max_new_tokens to prevent runaway generation translated_tokens = model.generate(**inputs, max_new_tokens=150) translated_text = tokenizer.decode(translated_tokens[0], skip_special_tokens=True) print(f"[{lang_name[:2].upper()}]: {translated_text}")