inferencelayer 0.2.3

#![allow(clippy::type_complexity)]
#![allow(clippy::large_enum_variant)]
//! `whisper` — speech-to-text on the engine. Log-mel + Whisper encoder (on the GPU when a device is
//! available, else pure-Rust CPU) + KV-cached greedy decode (CPU) + BPE detokenize. The encoder is
//! the compute-heavy part and runs on wgpu → Metal / Vulkan / DX12 / WebGPU, so the same code path
//! transcribes on a laptop, a server, or a browser.
//!
//!   whisper <model-dir> <audio.wav>            # 16 kHz mono PCM16
//!   WHISPER_GPU=0 whisper <model-dir> <wav>    # force the CPU encoder
//!
//! `<model-dir>` is a `export_whisper.py` export (config.json + model.safetensors + tokenizer.json).

use std::path::PathBuf;
use std::time::Instant;

use anyhow::{Context, Result};
use inferencelayer::GpuCtx;
use inferencelayer::whisper::{Whisper, WhisperEncoder};
use inferencelayer::whisper_gpu::{WhisperDecoderGpu, WhisperEncoderGpu};

/// The chosen encoder path. GPU keeps the context + wrapper resident (setup done once). The optional
/// decoder makes the WHOLE model GPU-resident (`WHISPER_GPU_DECODE=1`); it stays CPU by default since
/// autoregressive decode is not a native speedup (see `whisper_gpu`).
enum Enc {
    Cpu,
    Gpu(GpuCtx, WhisperEncoderGpu, Option<WhisperDecoderGpu>, String),
}

/// Minimal WAV reader: find the `data` chunk, read it as mono i16 LE @ 16 kHz (Whisper's input).
fn read_wav(path: &PathBuf) -> Result<Vec<f32>> {
    let bytes = std::fs::read(path).with_context(|| format!("read {}", path.display()))?;
    let dp = bytes
        .windows(4)
        .position(|w| w == b"data")
        .context("no `data` chunk — expected a 16 kHz mono PCM16 WAV")?
        + 8;
    Ok(bytes[dp..]
        .chunks_exact(2)
        .map(|c| i16::from_le_bytes(c.try_into().unwrap()) as f32 / 32768.0)
        .collect())
}

fn main() -> Result<()> {
    let mut args = std::env::args().skip(1);
    let usage = "usage: whisper <model-dir> <audio.wav>   (16 kHz mono PCM16)";
    let dir = PathBuf::from(args.next().context(usage)?);
    let wav = PathBuf::from(args.next().context(usage)?);

    let samples = read_wav(&wav)?;
    let audio_secs = samples.len() as f32 / 16_000.0;
    let model = Whisper::load(&dir).context("load model")?;
    let want_gpu = std::env::var("WHISPER_GPU")
        .map(|v| v != "0")
        .unwrap_or(true);
    let want_gpu_decode = std::env::var("WHISPER_GPU_DECODE")
        .map(|v| v == "1")
        .unwrap_or(false);

    // Setup (one-time; in a server this is startup): pick the device and build the encoder (and,
    // opt-in, the decoder). GPU init + shader compilation live here so they don't pollute timing.
    let engine = match want_gpu.then(GpuCtx::new).transpose() {
        // The GPU encoder wraps its own CPU encoder (for the conv stem + as parity oracle).
        Ok(Some(ctx)) => {
            let gpu = WhisperEncoderGpu::new(&ctx, WhisperEncoder::load(&dir)?, 1500)?;
            let dec = if want_gpu_decode {
                Some(WhisperDecoderGpu::new(&ctx, &model)?)
            } else {
                None
            };
            let device = format!("gpu ({})", ctx.backend);
            Enc::Gpu(ctx, gpu, dec, device)
        }
        _ => Enc::Cpu,
    };

    let t0 = Instant::now();
    let enc = match &engine {
        Enc::Gpu(ctx, gpu, _, _) => gpu.forward(ctx, &model.encoder.logmel(&samples))?,
        Enc::Cpu => model.encoder.encode(&samples),
    };
    let enc_ms = t0.elapsed().as_secs_f32() * 1e3;
    let enc_dev = match &engine {
        Enc::Gpu(_, _, _, d) => d.as_str(),
        Enc::Cpu => "cpu",
    };

    let t1 = Instant::now();
    let (ids, dec_dev) = match &engine {
        Enc::Gpu(ctx, _, Some(dec), _) => (dec.generate(ctx, &enc, model.max_target())?, "gpu"),
        _ => (model.generate_from_enc(&enc, model.max_target()), "cpu"),
    };
    let text = model.decode_text(&dir, &ids)?;
    let dec_ms = t1.elapsed().as_secs_f32() * 1e3;

    let total = enc_ms + dec_ms;
    eprintln!(
        "[whisper] {audio_secs:.1}s audio · encoder {enc_dev} {enc_ms:.0} ms · decode {dec_dev} \
         {dec_ms:.0} ms · total {total:.0} ms ({:.0}x realtime)",
        audio_secs * 1000.0 / total.max(1e-3),
    );
    println!("{text}");
    Ok(())
}