diff --git a/apps/stage-tamagotchi/src-tauri/src/lib.rs b/apps/stage-tamagotchi/src-tauri/src/lib.rs index 966792de3..9fb81a774 100644 --- a/apps/stage-tamagotchi/src-tauri/src/lib.rs +++ b/apps/stage-tamagotchi/src-tauri/src/lib.rs @@ -1,13 +1,13 @@ use std::{sync::atomic::Ordering, time::Duration}; use tauri::{ - menu::{Menu, MenuItem}, - tray::TrayIconBuilder, Emitter, Manager, RunEvent, WebviewUrl, WebviewWindowBuilder, + menu::{Menu, MenuItem}, + tray::TrayIconBuilder, }; use tauri_plugin_prevent_default::Flags; use tokio::time::sleep; @@ -15,12 +15,13 @@ use tokio::time::sleep; mod app_click_through; mod app_windows; mod commands; +mod whisper; #[cfg(target_os = "macos")] use app_click_through::native_macos::{get_mouse_location, get_window_frame}; #[cfg(target_os = "windows")] use app_click_through::native_windows::{get_mouse_location, get_window_frame}; -use app_click_through::state::{set_click_through_enabled, WindowClickThroughState}; +use app_click_through::state::{WindowClickThroughState, set_click_through_enabled}; #[tauri::command] async fn start_monitor(window: tauri::Window) -> Result<(), String> { @@ -95,6 +96,26 @@ async fn stop_click_through(window: tauri::Window) -> Result<(), String> { Ok(()) } +fn load_whisper_model() -> Result { + // Determine device to use + let device = if candle_core::utils::cuda_is_available() { + candle_core::Device::new_cuda(0)? + } else if candle_core::utils::metal_is_available() { + candle_core::Device::new_metal(0)? + } else { + candle_core::Device::Cpu + }; + + println!("Using device: {device:?}"); + + let whisper_model = whisper::WhichWhisperModel::Tiny; + + println!("Loading whisper model: {:?}", whisper_model); + + let model = whisper::WhisperProcessor::new(whisper_model, device.clone())?; + Ok(model) +} + #[cfg_attr(mobile, tauri::mobile_entry_point)] #[allow(clippy::missing_panics_doc)] pub fn run() { @@ -140,6 +161,9 @@ pub fn run() { )?; } + // Load whisper model + load_whisper_model()?; + // TODO: i18n let quit_item = MenuItem::with_id(app, "quit", "Quit", true, None::<&str>)?; let settings_item = MenuItem::with_id(app, "settings", "Settings", true, None::<&str>)?; diff --git a/apps/stage-tamagotchi/src-tauri/src/whisper/melfilters.bytes b/apps/stage-tamagotchi/src-tauri/src/whisper/melfilters.bytes new file mode 100644 index 000000000..0874829e2 Binary files /dev/null and b/apps/stage-tamagotchi/src-tauri/src/whisper/melfilters.bytes differ diff --git a/apps/stage-tamagotchi/src-tauri/src/whisper/melfilters128.bytes b/apps/stage-tamagotchi/src-tauri/src/whisper/melfilters128.bytes new file mode 100644 index 000000000..f287c5b1d Binary files /dev/null and b/apps/stage-tamagotchi/src-tauri/src/whisper/melfilters128.bytes differ diff --git a/apps/stage-tamagotchi/src-tauri/src/whisper/mod.rs b/apps/stage-tamagotchi/src-tauri/src/whisper/mod.rs new file mode 100644 index 000000000..a60d1a39b --- /dev/null +++ b/apps/stage-tamagotchi/src-tauri/src/whisper/mod.rs @@ -0,0 +1,207 @@ +// https://github.com/proj-airi/candle-examples/blob/3a6783a4788333e7a117f4f8548f02b2fbfec2ed/apps/silero-vad-whisper-realtime/src/whisper.rs +use anyhow::Result; +use byteorder::{ByteOrder, LittleEndian}; +use candle_core::{Device, IndexOp, Tensor}; +use candle_nn::VarBuilder; +use candle_transformers::models::whisper::{self as whisper_model, Config, audio}; +use clap::ValueEnum; +use hf_hub::{Repo, RepoType, api::sync::Api}; +use tokenizers::Tokenizer; + +pub enum WhisperModel { + Normal(whisper_model::model::Whisper), +} + +impl WhisperModel { + pub fn encoder_forward( + &mut self, + x: &Tensor, + flush: bool, + ) -> candle_core::Result { + match self { + Self::Normal(model) => model.encoder.forward(x, flush), + } + } + + pub fn decoder_forward( + &mut self, + x: &Tensor, + encoder_out: &Tensor, + flush: bool, + ) -> candle_core::Result { + match self { + Self::Normal(model) => model.decoder.forward(x, encoder_out, flush), + } + } + + pub fn decoder_final_linear( + &self, + x: &Tensor, + ) -> candle_core::Result { + match self { + Self::Normal(model) => model.decoder.final_linear(x), + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, ValueEnum)] +pub enum WhichWhisperModel { + Tiny, + #[value(name = "tiny.en")] + TinyEn, + Base, + #[value(name = "base.en")] + BaseEn, + Small, + #[value(name = "small.en")] + SmallEn, + Medium, + #[value(name = "medium.en")] + MediumEn, + Large, + LargeV2, + LargeV3, + LargeV3Turbo, + #[value(name = "distil-medium.en")] + DistilMediumEn, + #[value(name = "distil-large-v2")] + DistilLargeV2, +} + +impl WhichWhisperModel { + pub const fn model_and_revision(self) -> (&'static str, &'static str) { + match self { + Self::Tiny => ("openai/whisper-tiny", "main"), + Self::TinyEn => ("openai/whisper-tiny.en", "refs/pr/15"), + Self::Base => ("openai/whisper-base", "refs/pr/22"), + Self::BaseEn => ("openai/whisper-base.en", "refs/pr/13"), + Self::Small => ("openai/whisper-small", "main"), + Self::SmallEn => ("openai/whisper-small.en", "refs/pr/10"), + Self::Medium => ("openai/whisper-medium", "main"), + Self::MediumEn => ("openai/whisper-medium.en", "main"), + Self::Large => ("openai/whisper-large", "refs/pr/36"), + Self::LargeV2 => ("openai/whisper-large-v2", "refs/pr/57"), + Self::LargeV3 => ("openai/whisper-large-v3", "main"), + Self::LargeV3Turbo => ("openai/whisper-large-v3-turbo", "main"), + Self::DistilMediumEn => ("distil-whisper/distil-medium.en", "main"), + Self::DistilLargeV2 => ("distil-whisper/distil-large-v2", "main"), + } + } +} + +pub struct WhisperProcessor { + pub model: WhisperModel, + pub tokenizer: Tokenizer, + pub config: Config, + pub mel_filters: Vec, + pub device: Device, +} + +impl WhisperProcessor { + pub fn new( + model: WhichWhisperModel, + device: Device, + ) -> Result { + // Load the Whisper model based on the provided model type + let api = Api::new()?; + let (model_id, revision) = model.model_and_revision(); + let repo = api.repo(Repo::with_revision(model_id.to_string(), RepoType::Model, revision.to_string())); + + let config_filename = repo.get("config.json")?; + let tokenizer_filename = repo.get("tokenizer.json")?; + let model_filename = repo.get("model.safetensors")?; + + let config: Config = serde_json::from_str(&std::fs::read_to_string(config_filename)?)?; + let tokenizer = Tokenizer::from_file(tokenizer_filename).map_err(anyhow::Error::msg)?; + + println!("Loading Whisper model from: {:?}", model_filename.display()); + +// SAFETY: This is safe because we are using a mmaped file and the safetensors library guarantees that the data is valid. + let var_builder = unsafe { VarBuilder::from_mmaped_safetensors(&[model_filename], whisper_model::DTYPE, &device)? }; + + let model = WhisperModel::Normal(whisper_model::model::Whisper::load(&var_builder, config.clone())?); + + let mel_bytes = match config.num_mel_bins { + 80 => include_bytes!("./melfilters.bytes").as_slice(), + 128 => include_bytes!("./melfilters128.bytes").as_slice(), + num_mel_bins => anyhow::bail!("Unsupported number of mel bins: {}", num_mel_bins), + }; + + let mut mel_filters = vec![0f32; mel_bytes.len() / 4]; + ::read_f32_into(mel_bytes, &mut mel_filters); + + Ok(Self { model, tokenizer, config, mel_filters, device }) + } + + pub fn transcribe( + &mut self, + audio: &[f32], + ) -> Result { + // Convert PCM to mel spectrogram + let mel = audio::pcm_to_mel(&self.config, audio, &self.mel_filters); + let mel_len = mel.len(); + let mel = Tensor::from_vec(mel, (1, self.config.num_mel_bins, mel_len / self.config.num_mel_bins), &self.device)?; + + // Run inference + let audio_features = self.model.encoder_forward(&mel, true)?; + // Simple greedy decoding + let tokens = self.decode_greedy(&audio_features)?; + + let text = self + .tokenizer + .decode(&tokens, true) + .map_err(anyhow::Error::msg)?; + + Ok(text) + } + + fn decode_greedy( + &mut self, + audio_features: &Tensor, + ) -> Result> { + let mut tokens = vec![self.token_id(whisper_model::SOT_TOKEN)?, self.token_id(whisper_model::TRANSCRIBE_TOKEN)?, self.token_id(whisper_model::NO_TIMESTAMPS_TOKEN)?]; + + let max_len = 50; // Short sequence for real-time processing + + for i in 0..max_len { + let tokens_t = Tensor::new(tokens.as_slice(), &self.device)?.unsqueeze(0)?; + let ys = self + .model + .decoder_forward(&tokens_t, audio_features, i == 0)?; + + let (_, seq_len, _) = ys.dims3()?; + let logits = self + .model + .decoder_final_linear(&ys.i((..1, seq_len - 1..))?)? + .i(0)? + .i(0)?; + + // Get most likely token + let logits_v: Vec = logits.to_vec1()?; + let next_token = logits_v + .iter() + .enumerate() + .max_by(|(_, a), (_, b)| a.total_cmp(b)) + .map(|(i, _)| u32::try_from(i).unwrap()) + .unwrap(); + + if next_token == self.token_id(whisper_model::EOT_TOKEN)? { + break; + } + + tokens.push(next_token); + } + + Ok(tokens) + } + + fn token_id( + &self, + token: &str, + ) -> Result { + self + .tokenizer + .token_to_id(token) + .ok_or_else(|| anyhow::anyhow!("Token not found: {}", token)) + } +}