use anyhow::{Context, Result, bail};
use inferencelayer::chatterbox::{WATERMARK_ALPHA, WATERMARK_KEY, watermark_detect};
fn main() -> Result<()> {
let Some(path) = std::env::args().nth(1) else {
eprintln!("usage: chatterbox-wmcheck FILE.wav");
std::process::exit(2);
};
let (samples, sr) = match read_wav(&path) {
Ok(v) => v,
Err(e) => {
eprintln!("cannot read {path}: {e}");
std::process::exit(2);
}
};
let score = watermark_detect(&samples, WATERMARK_KEY);
let thresh = WATERMARK_ALPHA / 2.0;
let present = score > thresh;
println!(
"{path}: {:.2} s @ {sr} Hz | score {score:.4} (threshold {thresh:.4}) → {}",
samples.len() as f32 / sr as f32,
if present { "MARKED by this engine" } else { "no mark found" }
);
if sr != 24_000 {
println!(
" note: {sr} Hz input — the mark sits above 8 kHz, so anything resampled below \
16 kHz will read as unmarked regardless of its origin"
);
}
std::process::exit(if present { 0 } else { 1 });
}
fn read_wav(path: &str) -> Result<(Vec<f32>, u32)> {
let b = std::fs::read(path).with_context(|| format!("read {path}"))?;
if b.len() < 44 || &b[0..4] != b"RIFF" || &b[8..12] != b"WAVE" {
bail!("not a RIFF/WAVE file");
}
let (mut pos, mut fmt) = (12usize, None);
while pos + 8 <= b.len() {
let id = &b[pos..pos + 4];
let sz = u32::from_le_bytes([b[pos + 4], b[pos + 5], b[pos + 6], b[pos + 7]]) as usize;
let body = pos + 8;
if id == b"fmt " && body + 16 <= b.len() {
let format = u16::from_le_bytes([b[body], b[body + 1]]);
let ch = u16::from_le_bytes([b[body + 2], b[body + 3]]) as usize;
let sr = u32::from_le_bytes([b[body + 4], b[body + 5], b[body + 6], b[body + 7]]);
let bits = u16::from_le_bytes([b[body + 14], b[body + 15]]);
fmt = Some((format, ch.max(1), sr, bits));
} else if id == b"data" {
let (format, ch, sr, bits) = fmt.context("data chunk before fmt chunk")?;
let end = (body + sz).min(b.len());
let raw = &b[body..end];
let samples: Vec<f32> = match (format, bits) {
(1, 16) => raw
.chunks_exact(2)
.step_by(ch) .map(|c| i16::from_le_bytes([c[0], c[1]]) as f32 / 32768.0)
.collect(),
(3, 32) => raw
.chunks_exact(4)
.step_by(ch)
.map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
.collect(),
_ => bail!("unsupported WAV encoding (format {format}, {bits} bit)"),
};
if samples.is_empty() {
bail!("no samples");
}
return Ok((samples, sr));
}
pos = body + sz + (sz & 1);
}
bail!("no data chunk")
}