rust-candle-whisper

name	rust-candle-whisper
description	Implement native Rust ML inference with Candle framework. Use when building GPU-accelerated ML pipelines without Python dependencies.

Native ML with Candle

Pure Rust ML inference using the Candle framework for GPU-accelerated models.

Setup

# Cargo.toml
[dependencies]
candle-core = "0.4"
candle-nn = "0.4"
candle-transformers = "0.4"
hf-hub = "0.3"
tokenizers = "0.15"
symphonia = { version = "0.5", features = ["all"] }

[features]
cuda = ["candle-core/cuda", "candle-nn/cuda", "candle-transformers/cuda"]

Model Structure

use candle_core::{DType, Device, Tensor};
use candle_nn::VarBuilder;
use candle_transformers::models::whisper::{self as m, Config};
use hf_hub::{Repo, RepoType};
use std::path::Path;

pub struct WhisperModel {
    model: m::model::Whisper,
    tokenizer: WhisperTokenizer,
    mel_filters: Vec<f32>,
    device: Device,
}

Device Initialization

impl WhisperModel {
    fn init_device() -> Result<Device> {
        // Try CUDA first
        #[cfg(feature = "cuda")]
        {
            if let Ok(device) = Device::new_cuda(0) {
                tracing::info!("Using CUDA device");
                return Ok(device);
            }
        }

        // Fall back to CPU
        tracing::info!("Using CPU device");
        Ok(Device::Cpu)
    }
}

Loading from HuggingFace Hub

impl WhisperModel {
    pub fn load(model_id: &str, cache_dir: Option<&Path>) -> Result<Self> {
        tracing::info!("Loading model: {}", model_id);

        // Setup HF Hub with custom cache
        let cache_path = cache_dir
            .map(|p| p.to_path_buf())
            .unwrap_or_else(|| PathBuf::from("models/hf"));

        std::env::set_var("HF_HOME", &cache_path);

        let api = hf_hub::api::sync::ApiBuilder::new()
            .with_cache_dir(cache_path)
            .build()?;

        let repo = api.repo(Repo::new(model_id.to_string(), RepoType::Model));

        // Download model files
        let config_path = repo.get("config.json")?;
        let tokenizer_path = repo.get("tokenizer.json")?;
        let weights_path = repo.get("model.safetensors")?;

        // Load configuration
         : Config = {
              = std::fs::(&config_path)?;
            serde_json::(&content)?
        };

        tracing::info!(
            ,
            config.encoder_layers,
            config.decoder_layers
        );

        
          = ::()?;

        
          =  {
            VarBuilder::(&[weights_path], DType::F32, &device)?
        };

        
          = m::model::Whisper::(&vb, config)?;

        
          = WhisperTokenizer::(&tokenizer_path)?;
          = ()?;

        ( {
            model,
            tokenizer,
            mel_filters,
            device,
        })
    }
}

Audio Loading with Symphonia

use symphonia::core::audio::SampleBuffer;
use symphonia::core::codecs::DecoderOptions;
use symphonia::core::formats::FormatOptions;
use symphonia::core::io::MediaSourceStream;
use symphonia::core::probe::Hint;

pub fn load_audio(path: &Path) -> Result<Vec<f32>> {
    let file = std::fs::File::open(path)?;
    let mss = MediaSourceStream::new(Box::new(file), Default::default());

    let mut hint = Hint::new();
    if let Some(ext) = path.extension().and_then(|e| e.to_str()) {
        hint.with_extension(ext);
    }

    let probed = symphonia::default::get_probe()
        .format(&hint, mss, &FormatOptions::default(), &Default::default())?;

    let mut format = probed.format;
    let track = format.default_track()
        .ok_or_else(|| Error::(.()))?;

      = symphonia::default::()
        .(&track.codec_params, &DecoderOptions::())?;

      = track.id;
      = ::();

     {
          =  format.() {
            (p) => p,
            (symphonia::core::errors::Error::( e))
                 e.() == std::io::ErrorKind::UnexpectedEof => ,
            (e) =>  (e.()),
        };

         packet.() != track_id {
            ;
        }

          = decoder.(&packet)?;
          = *decoded.();

          = SampleBuffer::<>::(decoded.()  , spec);
        sample_buf.(decoded);

        
          = sample_buf.();
         spec.channels.() >  {
              = spec.channels.();
               channel_samples.(channels) {
                 :  = chunk.().sum::<>() / channels  ;
                samples.(avg);
            }
        }  {
            samples.(channel_samples);
        }
    }

    (samples)
}

Mel Spectrogram Computation

const N_FFT: usize = 400;
const HOP_LENGTH: usize = 160;
const N_MELS: usize = 128;

pub fn pcm_to_mel(samples: &[f32], filters: &[f32], device: &Device) -> Result<Tensor> {
    let n_frames = (samples.len() - N_FFT) / HOP_LENGTH + 1;

    // Pre-compute Hann window
    let hann_window: Vec<f32> = (0..N_FFT)
        .map(|i| 0.5 * (1.0 - (2.0 * std::f32::consts::PI * i as f32 / N_FFT as f32).cos()))
        .collect();

    // Compute STFT magnitudes
    let fft_size = N_FFT / 2 + 1;
    let mut magnitudes = vec![0.0f32; n_frames * fft_size];

    for frame_idx in 0..n_frames {
        let start = frame_idx * HOP_LENGTH;

        
         : <> = samples[start..start + N_FFT]
            .()
            .(&hann_window)
            .(|(s, w)| s * w)
            .();

        
           ..fft_size {
              = ;
              = ;

             (n, &sample)  windowed.().() {
                  = - * std::::consts::PI * k   * n   / N_FFT  ;
                real += sample * angle.();
                imag += sample * angle.();
            }

            magnitudes[frame_idx * fft_size + k] = real * real + imag * imag;
        }
    }

    
      = [; n_frames * N_MELS];

       ..n_frames {
           ..N_MELS {
              = ;
               ..fft_size {
                sum += filters[mel * fft_size + k] * magnitudes[frame * fft_size + k];
            }
            mel_spec[frame * N_MELS + mel] = sum.();
        }
    }

    
     : <> = mel_spec.().(|&x| x.().(-)).();
      = log_spec.().().(::NEG_INFINITY, ::max);
     : <> = log_spec
        .()
        .(|&x| ((x - max_val) / ).(-, ))
        .();

    
    Tensor::(normalized, (, N_MELS, n_frames), device)
        .(::into)
}

Autoregressive Decoding

use candle_nn::ops::softmax;

pub struct Decoder<'a> {
    model: &'a mut m::model::Whisper,
    tokenizer: &'a WhisperTokenizer,
    device: &'a Device,
    suppress_tokens: Vec<u32>,
}

impl<'a> Decoder<'a> {
    pub fn decode(&mut self, audio_features: &Tensor) -> Result<String> {
        // Initial tokens: <|startoftranscript|><|en|><|transcribe|><|notimestamps|>
        let mut tokens: Vec<u32> = vec![50258, 50259, 50359, 50363];
        let mut all_tokens = tokens.clone();

        // Autoregressive loop
        for step in 0..448 {
            let token_tensor = Tensor::new(tokens.as_slice(), self.device)?
                .unsqueeze(0)?;

            // Run decoder
              = .model.decoder
                .forward(&token_tensor, audio_features, step == )?;

            
              = logits.()?;
              = logits.((.., seq_len - , ..))?;

            
              = .(&last_logits)?;
              = (&last_logits, candle_core::D::Minus1)?;

              = probs
                .(candle_core::D::Minus1)?
                .(DType::U32)?
                .to_vec1::<>()?[];

            
             next_token ==  {  
                ;
            }

            all_tokens.(next_token);
            tokens = [next_token];
        }

        
         : <> = all_tokens
            .()
            .(|&&t| t < )  
            .()
            .();

        .tokenizer.(&text_tokens)
    }

     (&, logits: &Tensor)  <Tensor> {
          = logits.to_vec2::<>()?;

        
         &token  &.suppress_tokens {
            logits_vec[][token  ] = ::NEG_INFINITY;
        }

        Tensor::(logits_vec, .device).(::into)
    }
}

Global Model Caching

use std::sync::OnceLock;
use parking_lot::Mutex;

static WHISPER_MODEL: OnceLock<Mutex<WhisperModel>> = OnceLock::new();

pub fn transcribe(audio_path: &Path) -> Result<String> {
    let model = WHISPER_MODEL.get_or_init(|| {
        tracing::info!("Loading Whisper model (first use)...");
        Mutex::new(WhisperModel::load_default().expect("Failed to load model"))
    });

    let mut model_guard = model.lock();
    model_guard.transcribe(audio_path)
}

pub fn preload_model() -> Result<()> {
    if WHISPER_MODEL.get().is_some() {
        return Ok(());
    }

    let model = WhisperModel::load_default()?;
    let _ = WHISPER_MODEL.get_or_init(|| Mutex::new(model));
    Ok(())
}

VRAM Estimation

fn estimate_vram_gb(config: &Config) -> f32 {
    let encoder_params = config.encoder_layers * config.d_model * config.d_model * 4;
    let decoder_params = config.decoder_layers * config.d_model * config.d_model * 4;
    let vocab_params = config.vocab_size * config.d_model;
    let total_params = encoder_params + decoder_params + vocab_params;

    // float32 = 4 bytes, plus 20% overhead
    (total_params as f32 * 4.0 * 1.2) / (1024.0 * 1024.0 * 1024.0)
}

Guidelines

• Use cuda feature for GPU acceleration
• Memory-map weights with from_mmaped_safetensors
• Cache models globally with OnceLock
• Use Symphonia for pure-Rust audio decoding
• Pre-compute mel filterbank coefficients
• Implement token suppression for stable decoding
• Estimate VRAM before loading models

Examples

See hercules-local-algo/src/whisper/ for complete Whisper implementation.