diff --git a/apps/stage-tamagotchi/src-tauri/src/whisper/mod.rs b/apps/stage-tamagotchi/src-tauri/src/whisper/mod.rs index a60d1a39b..105252993 100644 --- a/apps/stage-tamagotchi/src-tauri/src/whisper/mod.rs +++ b/apps/stage-tamagotchi/src-tauri/src/whisper/mod.rs @@ -5,7 +5,11 @@ 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 hf_hub::{ + Repo, + RepoType, + api::sync::{Api, ApiBuilder}, +}; use tokenizers::Tokenizer; pub enum WhisperModel { @@ -97,18 +101,58 @@ pub struct WhisperProcessor { pub device: Device, } +pub struct WhisperProgressEmitter { + filename: String, +} + +impl WhisperProgressEmitter { + fn new(filename: &str) -> Self { + Self { + filename: filename.to_string(), + } + } +} + +impl hf_hub::api::Progress for WhisperProgressEmitter { + fn init( + &mut self, + size: usize, + filename: &str, + ) { + todo!(); + } + + fn update( + &mut self, + size: usize, + ) { + todo!(); + } + + fn finish(&mut self) { + todo!(); + } +} + 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 api = ApiBuilder::new().with_progress(false).build()?; 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 repo = api.repo(Repo::with_revision( + model_id.to_string(), + RepoType::Model, + revision.to_string(), + )); - let config_filename = repo.get("config.json")?; + let config_filename = + repo.download_with_progress("config.json", WhisperProgressEmitter::new("config.json"))?; + // TODO: for RainbowBird let tokenizer_filename = repo.get("tokenizer.json")?; + // TODO: for RainbowBird let model_filename = repo.get("model.safetensors")?; let config: Config = serde_json::from_str(&std::fs::read_to_string(config_filename)?)?; @@ -116,10 +160,15 @@ impl WhisperProcessor { 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)? }; + // 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 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(), @@ -130,7 +179,13 @@ impl WhisperProcessor { 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 }) + Ok(Self { + model, + tokenizer, + config, + mel_filters, + device, + }) } pub fn transcribe( @@ -140,7 +195,15 @@ impl WhisperProcessor { // 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)?; + 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)?; @@ -159,7 +222,11 @@ impl WhisperProcessor { &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 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