Optimize for Apple Silicon using MLX

This commit is contained in:
Adolfo Reyna
2026-02-26 20:57:31 -05:00
parent 8d39fda6bc
commit d1fa1ffb15

View File

@@ -1,4 +1,4 @@
import whisper import mlx_whisper
import numpy as np import numpy as np
import sounddevice as sd import sounddevice as sd
import queue import queue
@@ -7,13 +7,13 @@ import torch
from silero_vad import load_silero_vad, get_speech_timestamps from silero_vad import load_silero_vad, get_speech_timestamps
# Parameters # Parameters
MODEL_TYPE = "tiny.en" MODEL_PATH = "mlx-community/whisper-tiny.en-mlx" # MLX optimized model
CHANNELS = 1 CHANNELS = 1
SAMPLERATE = 16000 SAMPLERATE = 16000
BLOCK_SIZE = 512 # Silero VAD prefers specific block sizes (512, 1024, 1536) BLOCK_SIZE = 512 # Silero VAD prefers 512, 1024, or 1536
VAD_THRESHOLD = 0.5 # Confidence threshold for speech VAD_THRESHOLD = 0.5
BUFFER_LIMIT = SAMPLERATE * 30 # Max 30 seconds of audio buffer BUFFER_LIMIT = SAMPLERATE * 30
MIN_SILENCE_DURATION_MS = 500 # Silence duration to trigger transcription MIN_SILENCE_DURATION_MS = 500
audio_queue = queue.Queue() audio_queue = queue.Queue()
@@ -23,22 +23,15 @@ def callback(indata, frames, time, status):
audio_queue.put(indata.copy()) audio_queue.put(indata.copy())
def main(): def main():
print(f"Loading Whisper model '{MODEL_TYPE}'...") print(f"Loading MLX-optimized Whisper model '{MODEL_PATH}'...")
whisper_model = whisper.load_model(MODEL_TYPE) # mlx-whisper uses the same model names or Hugging Face paths
print("Loading Silero VAD model...") print("Loading Silero VAD model...")
vad_model = load_silero_vad() vad_model = load_silero_vad()
print("Models loaded.") print("Models loaded.")
print("\nAvailable Audio Devices:") print("\nStarting live transcription (MLX + VAD)... (Press Ctrl+C to stop)")
devices = sd.query_devices()
print(devices)
default_device = sd.default.device[0]
print(f"\nUsing default input device index: {default_device}")
print("\nStarting live transcription with VAD... (Press Ctrl+C to stop)")
audio_buffer = [] audio_buffer = []
speech_started = False speech_started = False
@@ -51,13 +44,9 @@ def main():
audio_buffer.append(data.flatten()) audio_buffer.append(data.flatten())
if len(audio_buffer) > 0: if len(audio_buffer) > 0:
# Concatenate buffer to check for speech
current_audio = np.concatenate(audio_buffer) current_audio = np.concatenate(audio_buffer)
# Convert to torch tensor for Silero
audio_tensor = torch.from_numpy(current_audio) audio_tensor = torch.from_numpy(current_audio)
# Get speech timestamps
speech_timestamps = get_speech_timestamps( speech_timestamps = get_speech_timestamps(
audio_tensor, audio_tensor,
vad_model, vad_model,
@@ -66,31 +55,24 @@ def main():
min_silence_duration_ms=MIN_SILENCE_DURATION_MS min_silence_duration_ms=MIN_SILENCE_DURATION_MS
) )
# If we have speech and then silence, or buffer is getting too long
if len(speech_timestamps) > 0: if len(speech_timestamps) > 0:
speech_started = True speech_started = True
# Check if the last speech segment has "ended" (i.e., we have enough silence after it)
# or if we've reached a significant buffer size
last_end = speech_timestamps[-1]['end'] last_end = speech_timestamps[-1]['end']
buffer_len_samples = len(current_audio) buffer_len_samples = len(current_audio)
# If the speech ended more than MIN_SILENCE_DURATION_MS ago
if (buffer_len_samples - last_end) > (SAMPLERATE * MIN_SILENCE_DURATION_MS / 1000) or buffer_len_samples > BUFFER_LIMIT: if (buffer_len_samples - last_end) > (SAMPLERATE * MIN_SILENCE_DURATION_MS / 1000) or buffer_len_samples > BUFFER_LIMIT:
# Transcribe the valid speech segment # Transcribe with MLX
result = whisper_model.transcribe(current_audio, fp16=False, language="en") result = mlx_whisper.transcribe(current_audio, path_or_hf_repo=MODEL_PATH)
text = result['text'].strip() text = result['text'].strip()
if text: if text:
print(f"Transcription: {text}") print(f"Transcription: {text}")
# Reset buffer
audio_buffer = [] audio_buffer = []
speech_started = False speech_started = False
elif not speech_started and len(current_audio) > SAMPLERATE * 2: elif not speech_started and len(current_audio) > SAMPLERATE * 2:
# Clear buffer if it's just silence for more than 2 seconds
audio_buffer = [] audio_buffer = []
except KeyboardInterrupt: except KeyboardInterrupt: