inferencelayer 0.2.9

//! `moshi-talk` — leave Moshi a voice message, get a spoken reply + transcript. The
//! file-mode demo of the full-duplex stack (our Mimi + our 7B on wgpu/CPU): no realtime
//! constraint, so it runs anywhere the checkpoints do.
//!
//! ```text
//! moshi-talk <input.wav> [--out reply.wav] [--cpu] [--tail-seconds 5] [--seed 42]
//! ```
//!
//! Input: any mono/stereo PCM16 or f32 WAV at any rate (resampled to 24 kHz mono). The input
//! is followed by `--tail-seconds` of silence so Moshi can keep talking after you stop.
//! Output: `reply.wav` (Moshi's voice, 24 kHz) + its inner-monologue transcript on stdout.
//! Checkpoints: `~/.cache/inferencelayer/{mimi,moshi-lm}` (see tests/fixtures/export_*.py);
//! the text tokenizer is read from the HF cache (tokenizer_spm_32k_3.model).

use std::path::PathBuf;

use anyhow::{Context, Result, bail};
use inferencelayer::mimi::{FRAME_SIZE, Mimi};
use inferencelayer::moshi_lm::{MoshiLm, Sampling};

fn main() -> Result<()> {
    let mut args = std::env::args().skip(1);
    let mut input = None;
    let mut out = PathBuf::from("reply.wav");
    let mut gpu = true;
    let mut tail_s = 5.0f64;
    let mut seed = 299_792_458u64; // the reference demo's seed
    let mut greedy = false;
    let mut gumbel = false;
    let mut full_gpu = false;
    let mut voice: Option<PathBuf> = None;
    let mut persona: Option<String> = None;
    let mut lm_dir_override: Option<PathBuf> = None;
    while let Some(a) = args.next() {
        match a.as_str() {
            "--out" => out = PathBuf::from(args.next().context("--out")?),
            "--cpu" => gpu = false,
            "--tail-seconds" => tail_s = args.next().context("--tail-seconds")?.parse()?,
            "--seed" => seed = args.next().context("--seed")?.parse()?,
            "--greedy" => greedy = true, // argmax decode — the live server's --greedy-gpu mode
            "--gumbel" => gumbel = true, // full-vocab Gumbel-max — the GPU sampler's semantics
            "--full-gpu" => full_gpu = true, // the SERVER's exact path: fused GPU step + noise
            "--voice" => voice = Some(PathBuf::from(args.next().context("--voice")?)),
            "--persona" => persona = Some(args.next().context("--persona")?),
            "--lm-dir" => lm_dir_override = Some(PathBuf::from(args.next().context("--lm-dir")?)),
            _ => input = Some(PathBuf::from(a)),
        }
    }
    let input = input.context("usage: moshi-talk <input.wav> [--out reply.wav] [--cpu]")?;
    let cache = PathBuf::from(std::env::var("HOME")?).join(".cache/inferencelayer");
    let lm_dir = lm_dir_override.unwrap_or_else(|| cache.join("moshi-lm"));

    // ---- input audio → 24 kHz mono f32, padded to whole frames + a silent tail -------------
    let mut pcm = read_wav_24k_mono(&input)?;
    let n_talk = pcm.len().div_ceil(FRAME_SIZE);
    pcm.resize((n_talk + (tail_s * 12.5) as usize) * FRAME_SIZE, 0.0);
    let n_frames = pcm.len() / FRAME_SIZE;
    eprintln!(
        "input: {:.1}s speech + {tail_s:.0}s tail = {n_frames} frames",
        n_talk as f64 * 0.08
    );

    // ---- models ------------------------------------------------------------------------------
    eprintln!("loading Mimi (CPU) …");
    let mimi = Mimi::load(&cache.join("mimi")).context("mimi checkpoint")?;
    eprintln!("loading Moshi 7B embeddings+depformer (CPU i8) …");
    let mut lm = MoshiLm::load_i8(&lm_dir).context("lm checkpoint")?;
    if !greedy {
        let s = Sampling { gumbel_full: gumbel, ..Sampling::default() };
        lm.set_sampling(s, seed);
    }
    let gpu_lm = if gpu {
        eprintln!("loading temporal onto wgpu (q8 FFN + q4 attention) …");
        let ctx = inferencelayer::GpuCtx::new().context("no GPU — rerun with --cpu")?;
        let w = inferencelayer::Weights::load(&ctx, lm_dir.clone())?;
        let g = inferencelayer::Lfm2Gpu::new(&ctx, w);
        Some((ctx, g))
    } else {
        eprintln!("running the temporal on CPU i8 (slow: ~0.3 s/frame)");
        None
    };
    let dep = match (&gpu_lm, full_gpu) {
        (Some((ctx, g)), true) => {
            eprintln!("loading depformer onto wgpu (fused full-GPU step — the server's path) …");
            Some(inferencelayer::moshi_gpu::MoshiDepGpu::load(ctx, lm_dir.clone(), g)?)
        }
        (None, true) => bail!("--full-gpu needs the GPU (drop --cpu)"),
        _ => None,
    };
    let mut noise_rng = seed.max(1);
    let mut text_rng = seed.rotate_left(17).max(1);
    let pieces = load_spm_pieces()?;

    // ---- the duplex loop ---------------------------------------------------------------------
    let mut enc = mimi.stream();
    let mut dec = mimi.stream();
    let mut st = lm.state();
    // PersonaPlex prompt priming (voice embeddings replay + forced silence/text frames)
    if voice.is_some() || persona.is_some() {
        use inferencelayer::moshi_lm::{PRIME_SILENCE_FRAMES, SILENCE_TOKENS, SINE_TOKENS};
        let (ctx, g) = gpu_lm.as_ref().context("persona priming needs the GPU")?;
        let rows = match &voice {
            Some(p) => inferencelayer::moshi_lm::load_voice_embeddings(p)?,
            None => Vec::new(),
        };
        let ptoks = match &persona {
            Some(spec) => inferencelayer::moshi_lm::load_persona_tokens(spec)?,
            None => Vec::new(),
        };
        let t0 = std::time::Instant::now();
        for row in &rows {
            g.forward_from_embeds_only(ctx, row, st.clock())?;
            st.advance_clock();
        }
        let mut fwd = |e: &[f32], p: usize| {
            g.forward_from_embeds_only(ctx, e, p).expect("prime fwd");
        };
        for _ in 0..PRIME_SILENCE_FRAMES {
            lm.step_forced(&mut st, 3, &SILENCE_TOKENS, &SINE_TOKENS, &mut fwd);
        }
        for &t in &ptoks {
            lm.step_forced(&mut st, t, &SILENCE_TOKENS, &SINE_TOKENS, &mut fwd);
        }
        for _ in 0..PRIME_SILENCE_FRAMES {
            lm.step_forced(&mut st, 3, &SILENCE_TOKENS, &SINE_TOKENS, &mut fwd);
        }
        eprintln!(
            "persona primed: {} voice + {} text frames in {:.1}s",
            rows.len(),
            ptoks.len(),
            t0.elapsed().as_secs_f64()
        );
    }
    let mut voice = Vec::with_capacity(pcm.len());
    let mut text = String::new();
    let t0 = std::time::Instant::now();
    for f in 0..n_frames {
        let user = enc.encode_frame(&pcm[f * FRAME_SIZE..(f + 1) * FRAME_SIZE]);
        if let Some(d) = &dep {
            let (ctx, g) = gpu_lm.as_ref().unwrap();
            let (tt, at, _) = if greedy {
                lm.step_full_ext(&mut st, &user, &mut |emb: &[f32], pos: usize| {
                    inferencelayer::moshi_gpu::moshi_full_step(g, d, ctx, emb, pos)
                        .expect("full gpu step")
                })
            } else {
                // the SERVER's default path: GPU top-k text + Gumbel audio, one submit
                d.upload_gumbel_noise_audio(ctx, &mut noise_rng, 0.8);
                d.upload_topk_noise(ctx, &mut text_rng, 0.7);
                lm.step_full_ext(&mut st, &user, &mut |emb: &[f32], pos: usize| {
                    inferencelayer::moshi_gpu::moshi_full_step_topk(g, d, ctx, emb, pos, None)
                        .expect("full gpu step")
                })
            };
            voice.extend_from_slice(&dec.decode_frame(&at));
            let tok = tt as usize;
            if tok != 0 && tok != 3
                && let Some(p) = pieces.get(tok) {
                    text.push_str(&p.replace('\u{2581}', " "));
                }
            if f % 25 == 24 {
                eprintln!("  {:>4.1}s / {:.1}s …", f as f64 * 0.08, n_frames as f64 * 0.08);
            }
            continue;
        }
        let tr = match &gpu_lm {
            Some((ctx, g)) => lm.step_ext(
                &mut st,
                &user,
                Some(&mut |emb: &[f32], pos: usize| {
                    let logits = g.forward_from_embeds(ctx, emb, pos).expect("gpu forward");
                    let hidden = g.read_hnorm_for_tests(ctx).expect("hnorm");
                    (hidden, logits)
                }),
            ),
            None => lm.step(&mut st, &user),
        };
        voice.extend_from_slice(&dec.decode_frame(&tr.audio_tokens));
        let tok = tr.text_token as usize;
        if tok != 0 && tok != 3
            && let Some(p) = pieces.get(tok) {
                text.push_str(&p.replace('\u{2581}', " "));
            }
        if f % 25 == 24 {
            eprintln!("  {:>4.1}s / {:.1}s …", f as f64 * 0.08, n_frames as f64 * 0.08);
        }
    }
    let dt = t0.elapsed().as_secs_f64();
    write_wav(&out, &voice)?;
    println!("\nMoshi said: {}", text.trim());
    println!(
        "\nreply: {} ({:.1}s of audio, generated in {dt:.1}s — {:.0} ms/frame)",
        out.display(),
        voice.len() as f64 / 24000.0,
        dt * 1e3 / n_frames as f64
    );
    Ok(())
}

// ---- helpers --------------------------------------------------------------------------------

fn read_wav_24k_mono(path: &PathBuf) -> Result<Vec<f32>> {
    let b = std::fs::read(path).with_context(|| format!("read {}", path.display()))?;
    anyhow::ensure!(b.len() > 44 && &b[..4] == b"RIFF" && &b[8..12] == b"WAVE", "not a WAV");
    let (mut fmt, mut data): (Option<(u16, u16, u32, u16)>, Option<&[u8]>) = (None, None);
    let mut o = 12usize;
    while o + 8 <= b.len() {
        let id = &b[o..o + 4];
        let sz = u32::from_le_bytes(b[o + 4..o + 8].try_into().unwrap()) as usize;
        let body = &b[o + 8..(o + 8 + sz).min(b.len())];
        if id == b"fmt " {
            fmt = Some((
                u16::from_le_bytes(body[0..2].try_into().unwrap()),
                u16::from_le_bytes(body[2..4].try_into().unwrap()),
                u32::from_le_bytes(body[4..8].try_into().unwrap()),
                u16::from_le_bytes(body[14..16].try_into().unwrap()),
            ));
        } else if id == b"data" {
            data = Some(body);
        }
        o += 8 + sz + (sz & 1);
    }
    let ((tag, ch, rate, bits), data) = (fmt.context("no fmt chunk")?, data.context("no data")?);
    let ch = ch as usize;
    let mono: Vec<f32> = match (tag, bits) {
        (1, 16) => data
            .chunks_exact(2 * ch)
            .map(|fr| {
                fr.chunks_exact(2)
                    .map(|s| i16::from_le_bytes(s.try_into().unwrap()) as f32 / 32768.0)
                    .sum::<f32>()
                    / ch as f32
            })
            .collect(),
        (3, 32) => data
            .chunks_exact(4 * ch)
            .map(|fr| {
                fr.chunks_exact(4)
                    .map(|s| f32::from_le_bytes(s.try_into().unwrap()))
                    .sum::<f32>()
                    / ch as f32
            })
            .collect(),
        _ => bail!("unsupported WAV: format {tag}, {bits}-bit (use PCM16 or f32)"),
    };
    if rate == 24000 {
        return Ok(mono);
    }
    // linear resample — demo-grade, fine for speech input
    let ratio = rate as f64 / 24000.0;
    let n = (mono.len() as f64 / ratio) as usize;
    Ok((0..n)
        .map(|i| {
            let x = i as f64 * ratio;
            let (a, t) = (x as usize, x.fract() as f32);
            let b = (a + 1).min(mono.len() - 1);
            mono[a] * (1.0 - t) + mono[b] * t
        })
        .collect())
}

fn write_wav(path: &PathBuf, pcm: &[f32]) -> Result<()> {
    let n = pcm.len() as u32;
    let mut w = Vec::with_capacity(44 + 2 * pcm.len());
    w.extend_from_slice(b"RIFF");
    w.extend_from_slice(&(36 + 2 * n).to_le_bytes());
    w.extend_from_slice(b"WAVEfmt ");
    w.extend_from_slice(&16u32.to_le_bytes());
    w.extend_from_slice(&1u16.to_le_bytes());
    w.extend_from_slice(&1u16.to_le_bytes());
    w.extend_from_slice(&24000u32.to_le_bytes());
    w.extend_from_slice(&48000u32.to_le_bytes());
    w.extend_from_slice(&2u16.to_le_bytes());
    w.extend_from_slice(&16u16.to_le_bytes());
    w.extend_from_slice(b"data");
    w.extend_from_slice(&(2 * n).to_le_bytes());
    for &s in pcm {
        w.extend_from_slice(&((s.clamp(-1.0, 1.0) * 32767.0) as i16).to_le_bytes());
    }
    Ok(std::fs::write(path, w)?)
}

fn load_spm_pieces() -> Result<Vec<String>> {
    inferencelayer::moshi_lm::spm_pieces(&inferencelayer::moshi_lm::find_spm_model()?)
}