diff --git a/crates/transcription/src/lib.rs b/crates/transcription/src/lib.rs index 84afc65..1d5ff6b 100644 --- a/crates/transcription/src/lib.rs +++ b/crates/transcription/src/lib.rs @@ -5,7 +5,7 @@ pub mod whisper_rs_backend; pub use concurrency::run_inference; pub use local_engine::{ - load_parakeet, load_whisper, LocalEngine, TimedTranscript, + load_parakeet, load_whisper, LocalEngine, SpeechBackend, TimedTranscript, }; pub use transcribe_rs::SpeechModel; pub use model_manager::{ diff --git a/crates/transcription/src/local_engine.rs b/crates/transcription/src/local_engine.rs index 66612f1..8329a4f 100644 --- a/crates/transcription/src/local_engine.rs +++ b/crates/transcription/src/local_engine.rs @@ -10,17 +10,30 @@ use kon_core::types::{ TranscriptionOptions, }; +use crate::whisper_rs_backend::WhisperRsBackend; + /// Result of a timed transcription: transcript + inference duration. pub struct TimedTranscript { pub transcript: Transcript, pub inference_ms: u64, } +/// Public discriminator selected by the loaders (`load_parakeet`, `load_whisper`) +/// and passed to `LocalEngine::load`. `src-tauri::commands::models` names this +/// type as the return of `load_model_from_disk`, so it must be `pub`. +pub enum SpeechBackend { + /// transcribe-rs-owned model. Used for Parakeet ONNX (wrapped in + /// ParakeetWordGranularity for word-level timestamps). + Adapter(Box), + /// Direct whisper-rs. The only path that actually forwards `initial_prompt`. + WhisperRs(WhisperRsBackend), +} + /// Wraps any transcribe-rs engine in Kon's SpeechToText trait. /// Encapsulates threading: inference always runs on a blocking thread. /// The rest of the app never imports transcribe-rs directly. pub struct LocalEngine { - engine: Mutex>>, + engine: Mutex>, engine_name: EngineName, loaded_model_id: Mutex>, } @@ -34,10 +47,10 @@ impl LocalEngine { } } - pub fn load(&self, model: Box, model_id: ModelId) { + pub fn load(&self, backend: SpeechBackend, model_id: ModelId) { let mut guard = self.engine.lock().unwrap_or_else(|e| e.into_inner()); - *guard = Some(model); + *guard = Some(backend); let mut id_guard = self .loaded_model_id .lock() @@ -72,33 +85,37 @@ impl LocalEngine { ) -> Result { let mut guard = self.engine.lock().unwrap_or_else(|e| e.into_inner()); - let engine = - guard.as_mut().ok_or(KonError::EngineNotLoaded)?; - - let opts = TranscribeOptions { - language: options.language.clone(), - translate: false, - leading_silence_ms: None, - trailing_silence_ms: None, - }; + let backend = guard.as_mut().ok_or(KonError::EngineNotLoaded)?; let start = Instant::now(); - let result: TranscriptionResult = engine - .transcribe(audio.samples(), &opts) - .map_err(|e| KonError::TranscriptionFailed(e.to_string()))?; + let segments: Vec = match backend { + SpeechBackend::Adapter(model) => { + let opts = TranscribeOptions { + language: options.language.clone(), + translate: false, + leading_silence_ms: None, + trailing_silence_ms: None, + }; + let result: TranscriptionResult = model + .transcribe(audio.samples(), &opts) + .map_err(|e| KonError::TranscriptionFailed(e.to_string()))?; + result + .segments + .unwrap_or_default() + .into_iter() + .map(|s| Segment { + start: s.start as f64, + end: s.end as f64, + text: s.text, + }) + .collect() + } + SpeechBackend::WhisperRs(w) => w + .transcribe_sync(audio.samples(), options) + .map_err(|e| KonError::TranscriptionFailed(e.to_string()))?, + }; let inference_ms = start.elapsed().as_millis() as u64; - let segments = result - .segments - .unwrap_or_default() - .into_iter() - .map(|s| Segment { - start: s.start as f64, - end: s.end as f64, - text: s.text, - }) - .collect(); - Ok(TimedTranscript { transcript: Transcript::new( segments, @@ -150,9 +167,7 @@ impl transcribe_rs::SpeechModel for ParakeetWordGranularity { } /// Load a Parakeet model from a directory path. -pub fn load_parakeet( - model_dir: &Path, -) -> Result> { +pub fn load_parakeet(model_dir: &Path) -> Result { use transcribe_rs::onnx::Quantization; let model = transcribe_rs::onnx::parakeet::ParakeetModel::load( model_dir, @@ -163,21 +178,15 @@ pub fn load_parakeet( "Failed to load Parakeet: {e}" )) })?; - Ok(Box::new(ParakeetWordGranularity(model))) + Ok(SpeechBackend::Adapter(Box::new(ParakeetWordGranularity(model)))) } -/// Load a Whisper model from a GGML file path. -pub fn load_whisper( - model_path: &Path, -) -> Result> { - let engine = - transcribe_rs::whisper_cpp::WhisperEngine::load(model_path) - .map_err(|e| { - KonError::TranscriptionFailed(format!( - "Failed to load Whisper: {e}" - )) - })?; - Ok(Box::new(engine)) +/// Load a Whisper model from a GGML file path via whisper-rs. +pub fn load_whisper(model_path: &Path) -> Result { + let backend = WhisperRsBackend::load(model_path).map_err(|e| { + KonError::TranscriptionFailed(format!("Failed to load Whisper: {e}")) + })?; + Ok(SpeechBackend::WhisperRs(backend)) } #[cfg(test)] diff --git a/src-tauri/src/commands/models.rs b/src-tauri/src/commands/models.rs index d9000e0..a7dfeb0 100644 --- a/src-tauri/src/commands/models.rs +++ b/src-tauri/src/commands/models.rs @@ -72,7 +72,7 @@ fn model_capability( } } -pub fn load_model_from_disk(model_id: &ModelId) -> Result, String> { +pub fn load_model_from_disk(model_id: &ModelId) -> Result { let entry = model_registry::find_model(model_id) .ok_or_else(|| format!("Unknown model: {model_id}"))?;