inferencelayer 0.2.4

//! `moshi-converse` — a persistent duplex engine driven turn-by-turn over stdin, so an
//! external orchestrator (scripts/personaplex/adaptive_conversation.py) can run a FULLY
//! ADAPTIVE conversation: an LLM plays the user and writes each next line from Moshi's ACTUAL
//! spoken reply, not a fixed script. The engine loads once and keeps session state across
//! turns (context accumulates like a real call).
//!
//! Protocol (one JSON object per line):
//!   in : {"turn_wav": "/path/user.wav"}            play the turn, keep state
//!        {"quit": true}
//!   out: {"moshi_text": "...", "moshi_wav": "/path/reply.wav", "voiced_secs": 2.3,
//!         "heard": "<whisper of the reply>"}
//!
//!   MOSHI_LM_DIR=… moshi-converse [--voice V.safetensors] [--persona personas.json:key]
#![cfg(all(feature = "encoder-cpu", feature = "cli"))]

use std::io::{BufRead, Write};
use std::path::PathBuf;

use anyhow::{Context, Result};
use inferencelayer::mimi::{FRAME_SIZE, Mimi};
use inferencelayer::mimi_gpu::MimiGpu;
use inferencelayer::moshi_gpu::{MoshiDepGpu, moshi_full_step_topk};
use inferencelayer::moshi_lm::MoshiLm;
use inferencelayer::{GpuCtx, Lfm2Gpu, Weights};

fn read_wav_24k(path: &str) -> Vec<f32> {
    let Ok(b) = std::fs::read(path) else { return Vec::new() };
    if b.len() < 44 || &b[..4] != b"RIFF" {
        return Vec::new();
    }
    let (mut fmt, mut data) = (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[2..4].try_into().unwrap()) as usize,
                u32::from_le_bytes(body[4..8].try_into().unwrap()),
            ));
        } else if id == b"data" {
            data = Some(body);
        }
        o += 8 + sz + (sz & 1);
    }
    let (Some((ch, rate)), Some(data)) = (fmt, data) else { return Vec::new() };
    let mono: Vec<f32> = 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();
    if rate == 24000 {
        return mono;
    }
    let ratio = rate as f64 / 24000.0;
    let n = (mono.len() as f64 / ratio) as usize;
    (0..n)
        .map(|i| {
            let x = i as f64 * ratio;
            let (a, t) = (x as usize, x.fract() as f32);
            mono[a] * (1.0 - t) + mono[(a + 1).min(mono.len() - 1)] * t
        })
        .collect()
}

fn write_wav(path: &str, pcm: &[f32]) {
    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());
    }
    let _ = std::fs::write(path, w);
}

fn voiced_secs(pcm: &[f32]) -> f64 {
    pcm.chunks(2400)
        .filter(|c| (c.iter().map(|v| v * v).sum::<f32>() / c.len() as f32).sqrt() > 0.01)
        .count() as f64
        * 0.1
}

fn main() -> Result<()> {
    let mut args = std::env::args().skip(1);
    let mut voice: Option<PathBuf> = None;
    let mut persona: Option<String> = None;
    while let Some(a) = args.next() {
        match a.as_str() {
            "--voice" => voice = Some(PathBuf::from(args.next().context("--voice")?)),
            "--persona" => persona = Some(args.next().context("--persona")?),
            other => anyhow::bail!("unknown arg {other}"),
        }
    }
    let cache = PathBuf::from(std::env::var("HOME")?).join(".cache/inferencelayer");
    let lm_dir = std::env::var("MOSHI_LM_DIR").map(PathBuf::from).unwrap_or_else(|_| cache.join("moshi-lm"));

    eprintln!("moshi-converse: loading {} …", lm_dir.display());
    let ctx = GpuCtx::new()?;
    let mimi = Mimi::load(&cache.join("mimi"))?;
    let mut codec = MimiGpu::new(&ctx, &mimi)?;
    let lm = MoshiLm::load_i8(&lm_dir)?;
    let w = Weights::load(&ctx, lm_dir.clone())?;
    let gpu = Lfm2Gpu::new(&ctx, w);
    let dep = MoshiDepGpu::load(&ctx, lm_dir, &gpu)?;
    let pieces = inferencelayer::moshi_lm::spm_pieces(&inferencelayer::moshi_lm::find_spm_model()?)?;
    let whisper = inferencelayer::whisper::Whisper::load(&cache.join("whisper-base")).ok();

    let mut st = lm.state();
    // one-time persona priming (voice replay + role prompt), like the server
    if voice.is_some() || persona.is_some() {
        use inferencelayer::moshi_lm::{PRIME_SILENCE_FRAMES, SILENCE_TOKENS, SINE_TOKENS};
        let rows = voice.as_ref().map(|p| inferencelayer::moshi_lm::load_voice_embeddings(p)).transpose()?.unwrap_or_default();
        let ptoks = persona.as_ref().map(|s| inferencelayer::moshi_lm::load_persona_tokens(s)).transpose()?.unwrap_or_default();
        for row in &rows {
            gpu.forward_from_embeds_only(&ctx, row, st.clock())?;
            st.advance_clock();
        }
        let mut fwd = |e: &[f32], p: usize| { gpu.forward_from_embeds_only(&ctx, e, p).expect("prime"); };
        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); }
    }

    let (mut rng, mut trng) = (299_792_458u64, 424_242u64);
    let stdin = std::io::stdin();
    let mut stdout = std::io::stdout();
    eprintln!("READY");
    for line in stdin.lock().lines() {
        let line = line?;
        let cmd: serde_json::Value = serde_json::from_str(&line).unwrap_or_default();
        if cmd.get("quit").and_then(|v| v.as_bool()) == Some(true) {
            break;
        }
        let Some(wav) = cmd.get("turn_wav").and_then(|v| v.as_str()) else { continue };
        let user = read_wav_24k(wav);
        let n_user = user.len().div_ceil(FRAME_SIZE);
        let total = n_user + 40; // 3.2 s response window
        let mut agent = Vec::with_capacity(total * FRAME_SIZE);
        let mut text = String::new();
        let silence = [0f32; FRAME_SIZE];
        for f in 0..total {
            let frame = user.get(f * FRAME_SIZE..(f + 1) * FRAME_SIZE).unwrap_or(&silence);
            let uc = codec.encode_frame(&ctx, frame)?;
            dep.upload_gumbel_noise_audio(&ctx, &mut rng, 0.8);
            dep.upload_topk_noise(&ctx, &mut trng, 0.7);
            let (tt, at, _) = lm.step_full_ext(&mut st, &uc, &mut |emb: &[f32], pos: usize| {
                moshi_full_step_topk(&gpu, &dep, &ctx, emb, pos, None).expect("step")
            });
            agent.extend(codec.decode_frame(&ctx, &at)?);
            let t = tt as usize;
            if t != 0 && t != 3 && let Some(p) = pieces.get(t) {
                text.push_str(&p.replace('\u{2581}', " "));
            }
        }
        let reply_wav = format!("{wav}.reply.wav");
        write_wav(&reply_wav, &agent);
        let heard = match &whisper {
            Some(wsp) if voiced_secs(&agent) > 0.3 => {
                let m = agent.len() * 2 / 3;
                let pcm16: Vec<f32> = (0..m).map(|i| {
                    let x = i as f32 * 1.5; let (a, f) = (x as usize, x.fract());
                    agent[a] * (1.0 - f) + agent[(a + 1).min(agent.len() - 1)] * f
                }).collect();
                wsp.transcribe(&cache.join("whisper-base"), &pcm16).unwrap_or_default().trim().to_string()
            }
            _ => String::new(),
        };
        let out = serde_json::json!({
            "moshi_text": text.trim(), "moshi_wav": reply_wav,
            "voiced_secs": (voiced_secs(&agent) * 10.0).round() / 10.0, "heard": heard,
        });
        writeln!(stdout, "{out}")?;
        stdout.flush()?;
    }
    Ok(())
}