feat: load whisper (#234)

Co-authored-by: Neko <neko@ayaka.moe>

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<whisper::WhisperProcessor, anyhow::Error> {
+  // 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>)?;

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<Tensor> {
+    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<Tensor> {
+    match self {
+      Self::Normal(model) => model.decoder.forward(x, encoder_out, flush),
+    }
+  }
+
+  pub fn decoder_final_linear(
+    &self,
+    x: &Tensor,
+  ) -> candle_core::Result<Tensor> {
+    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<f32>,
+  pub device:      Device,
+}
+
+impl WhisperProcessor {
+  pub fn new(
+    model: WhichWhisperModel,
+    device: Device,
+  ) -> Result<Self> {
+    // 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];
+    <LittleEndian as ByteOrder>::read_f32_into(mel_bytes, &mut mel_filters);
+
+    Ok(Self { model, tokenizer, config, mel_filters, device })
+  }
+
+  pub fn transcribe(
+    &mut self,
+    audio: &[f32],
+  ) -> Result<String> {
+    // 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<Vec<u32>> {
+    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<f32> = 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<u32> {
+    self
+      .tokenizer
+      .token_to_id(token)
+      .ok_or_else(|| anyhow::anyhow!("Token not found: {}", token))
+  }
+}