Optimize for Apple Silicon using MLX
This commit is contained in:
@@ -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:
|
||||||
|
|||||||
Reference in New Issue
Block a user