use std::sync::Arc; use serde::Serialize; use tauri::Emitter; use crate::AppState; use kon_core::model_registry::{ self, Engine, LanguageSupport, ModelEntry, }; use kon_core::types::ModelId; use kon_transcription::model_manager; use kon_transcription::{load_parakeet, load_whisper, LocalEngine}; /// Map legacy size strings to ModelId. fn whisper_model_id(size: &str) -> ModelId { match size.to_lowercase().as_str() { "tiny" => ModelId::new("whisper-tiny-en"), "base" => ModelId::new("whisper-base-en"), "small" => ModelId::new("whisper-small-en"), "medium" => ModelId::new("whisper-medium-en"), other => ModelId::new(other), } } fn parakeet_model_id(name: &str) -> ModelId { match name { "ctc-int8" => ModelId::new("parakeet-ctc-0.6b-int8"), other => ModelId::new(other), } } fn engine_for_name( state: &AppState, engine_name: &str, ) -> Result, String> { match engine_name { "whisper" => Ok(state.whisper_engine.clone()), "parakeet" => Ok(state.parakeet_engine.clone()), other => Err(format!("Unknown engine: {other}")), } } fn language_support_info( language_support: LanguageSupport, ) -> LanguageSupportInfo { match language_support { LanguageSupport::EnglishOnly => LanguageSupportInfo { kind: "english-only".into(), language_count: 1, }, LanguageSupport::Multilingual(count) => LanguageSupportInfo { kind: "multilingual".into(), language_count: count, }, } } fn model_capability( entry: &'static ModelEntry, engine: &Arc, ) -> ModelRuntimeCapabilities { let loaded_model_id = engine.loaded_model_id(); ModelRuntimeCapabilities { id: entry.id.to_string(), display_name: entry.display_name.to_string(), downloaded: model_manager::is_downloaded(&entry.id), loaded: loaded_model_id .as_ref() .map(|id| id == &entry.id) .unwrap_or(false), language_support: language_support_info(entry.languages), } } pub fn load_model_from_disk(model_id: &ModelId) -> Result, String> { let entry = model_registry::find_model(model_id) .ok_or_else(|| format!("Unknown model: {model_id}"))?; match entry.engine { Engine::Whisper => { let dir = model_manager::model_dir(model_id); let model_file = entry .files .first() .map(|file| dir.join(file.filename)) .ok_or_else(|| format!("No files registered for model: {model_id}"))?; if !model_file.exists() { return Err(format!("Model not downloaded: {model_id}")); } load_whisper(&model_file).map_err(|e| e.to_string()) } Engine::Parakeet => { let dir = model_manager::model_dir(model_id); if !dir.exists() { return Err(format!("Model not downloaded: {model_id}")); } load_parakeet(&dir).map_err(|e| e.to_string()) } Engine::Moonshine => Err("Moonshine models are not yet supported in this build".into()), } } pub fn default_model_id_for_engine(engine: &str) -> ModelId { match engine { "parakeet" => ModelId::new("parakeet-ctc-0.6b-int8"), _ => ModelId::new("whisper-base-en"), } } pub async fn ensure_model_loaded( state: &AppState, engine_name: &str, model_id: &str, ) -> Result<(), String> { let model_id = ModelId::new(model_id); let entry = model_registry::find_model(&model_id) .ok_or_else(|| format!("Unknown model: {model_id}"))?; let expected_engine = match entry.engine { Engine::Whisper => "whisper", Engine::Parakeet => "parakeet", Engine::Moonshine => "moonshine", }; if expected_engine != engine_name { return Err(format!( "Model {} belongs to {}, not {}", model_id, expected_engine, engine_name )); } if !model_manager::is_downloaded(&model_id) { return Err(format!("Model not downloaded: {model_id}")); } let engine = engine_for_name(state, engine_name)?; if engine .loaded_model_id() .as_ref() .map(|loaded| loaded == &model_id) .unwrap_or(false) { return Ok(()); } let engine_clone = engine.clone(); let model_id_clone = model_id.clone(); tokio::task::spawn_blocking(move || { let model = load_model_from_disk(&model_id_clone)?; engine_clone.load(model, model_id_clone); Ok::<_, String>(()) }) .await .map_err(|e| e.to_string())??; Ok(()) } /// Spawns a background task that loads the default Whisper model into memory /// if it is already downloaded. Returns immediately — errors are logged, never panicked. /// /// Call this once from app setup() after AppState is managed. pub fn prewarm_default_model(whisper_engine: Arc) { let model_id = default_model_id_for_engine("whisper"); if !model_manager::is_downloaded(&model_id) { return; } if whisper_engine .loaded_model_id() .as_ref() .map(|id| id == &model_id) .unwrap_or(false) { return; } tauri::async_runtime::spawn(async move { let result = tauri::async_runtime::spawn_blocking(move || { load_model_from_disk(&model_id).map(|model| { whisper_engine.load(model, model_id); }) }) .await; match result { Ok(Ok(())) => eprintln!("[startup] Whisper model pre-warmed successfully"), Ok(Err(e)) => eprintln!("[startup] Whisper pre-warm failed: {e}"), Err(e) => eprintln!("[startup] Whisper pre-warm task panicked: {e}"), } }); } #[derive(Serialize)] #[serde(rename_all = "camelCase")] pub struct RuntimeCapabilities { pub accelerators: Vec, pub engines: Vec, } #[derive(Serialize)] #[serde(rename_all = "camelCase")] pub struct EngineRuntimeCapabilities { pub id: String, pub default_model_id: String, pub loaded_model_id: Option, pub supports_gpu: bool, pub models: Vec, } #[derive(Serialize)] #[serde(rename_all = "camelCase")] pub struct ModelRuntimeCapabilities { pub id: String, pub display_name: String, pub downloaded: bool, pub loaded: bool, pub language_support: LanguageSupportInfo, } #[derive(Serialize)] #[serde(rename_all = "camelCase")] pub struct LanguageSupportInfo { pub kind: String, pub language_count: u16, } #[tauri::command] pub fn get_runtime_capabilities( state: tauri::State<'_, AppState>, ) -> Result { let whisper = state.whisper_engine.clone(); let parakeet = state.parakeet_engine.clone(); let whisper_models = model_registry::all_models() .iter() .filter(|entry| entry.engine == Engine::Whisper) .map(|entry| model_capability(entry, &whisper)) .collect(); let parakeet_models = model_registry::all_models() .iter() .filter(|entry| entry.engine == Engine::Parakeet) .map(|entry| model_capability(entry, ¶keet)) .collect(); Ok(RuntimeCapabilities { // Current desktop build ships CPU-only inference backends. accelerators: vec!["cpu".into()], engines: vec![ EngineRuntimeCapabilities { id: "whisper".into(), default_model_id: default_model_id_for_engine("whisper") .to_string(), loaded_model_id: whisper .loaded_model_id() .map(|model_id| model_id.to_string()), supports_gpu: false, models: whisper_models, }, EngineRuntimeCapabilities { id: "parakeet".into(), default_model_id: default_model_id_for_engine("parakeet") .to_string(), loaded_model_id: parakeet .loaded_model_id() .map(|model_id| model_id.to_string()), supports_gpu: false, models: parakeet_models, }, ], }) } // --- Whisper model commands --- #[tauri::command] pub async fn download_model( app: tauri::AppHandle, size: String, ) -> Result { let id = whisper_model_id(&size); let app_clone = app.clone(); model_manager::download(&id, move |progress| { let _ = app_clone.emit("model-download-progress", &progress); }) .await .map_err(|e| e.to_string())?; Ok(format!("Model {} downloaded", size)) } #[tauri::command] pub fn check_model(size: String) -> Result { let id = whisper_model_id(&size); Ok(model_manager::is_downloaded(&id)) } #[tauri::command] pub fn list_models() -> Result, String> { let ids = model_manager::list_downloaded(); Ok(ids .into_iter() .filter(|id| id.as_str().starts_with("whisper-")) .map(|id| match id.as_str() { "whisper-tiny-en" => "Tiny".to_string(), "whisper-base-en" => "Base".to_string(), "whisper-small-en" => "Small".to_string(), "whisper-medium-en" => "Medium".to_string(), other => other.to_string(), }) .collect()) } #[tauri::command] pub async fn load_model( state: tauri::State<'_, AppState>, size: String, ) -> Result { let id = whisper_model_id(&size); ensure_model_loaded(&state, "whisper", id.as_str()).await?; Ok(format!("Model {} loaded", size)) } #[tauri::command] pub fn check_engine(state: tauri::State) -> Result { Ok(state.whisper_engine.is_loaded()) } // --- Parakeet model commands --- #[tauri::command] pub async fn download_parakeet_model( app: tauri::AppHandle, name: String, ) -> Result { let id = parakeet_model_id(&name); let app_clone = app.clone(); model_manager::download(&id, move |progress| { let _ = app_clone.emit("parakeet-download-progress", &progress); }) .await .map_err(|e| e.to_string())?; Ok(format!("Parakeet model {} downloaded", name)) } #[tauri::command] pub fn check_parakeet_model(name: String) -> Result { let id = parakeet_model_id(&name); Ok(model_manager::is_downloaded(&id)) } #[tauri::command] pub fn list_parakeet_models() -> Result, String> { let ids = model_manager::list_downloaded(); Ok(ids .into_iter() .filter(|id| id.as_str().starts_with("parakeet-")) .map(|id| match id.as_str() { "parakeet-ctc-0.6b-int8" => "ctc-int8".to_string(), other => other.to_string(), }) .collect()) } #[tauri::command] pub async fn load_parakeet_model( state: tauri::State<'_, AppState>, name: String, ) -> Result { let id = parakeet_model_id(&name); ensure_model_loaded(&state, "parakeet", id.as_str()).await?; Ok(format!("Parakeet model {} loaded", name)) } #[tauri::command] pub fn check_parakeet_engine( state: tauri::State<'_, AppState>, ) -> Result { Ok(state.parakeet_engine.is_loaded()) }