rustytorch/crates/models/rtx-csm/examples/audioseal_demo.rs

//! Load the converted AudioSeal weights and run a generator + detector pass
//! on a synthetic 1-second 16 kHz signal. Verifies that the converted
//! safetensors keys match what `Generator::new` and `Detector::new` expect,
//! and that the forward pipeline produces sensible-shaped output and a
//! decodable message.
//!
//! Usage:
//! ```
//! cargo run -p rtx-csm --release --example audioseal_demo -- \
//!     --generator /tmp/audioseal_generator.safetensors \
//!     --detector  /tmp/audioseal_detector.safetensors \
//!     --message 0xBEEF
//! ```

use anyhow::{Context, Result};
use candle_core::{DType, Device, Tensor};
use clap::Parser;
use rtx_csm::audio_io;
use rtx_csm::audioseal::{AudioSealWatermarker, Detector, Generator, MESSAGE_BITS, SAMPLE_RATE};
use rtx_csm::watermark::Watermarker;
use std::path::PathBuf;

#[derive(Debug, Parser)]
#[command(name = "audioseal_demo")]
struct Cli {
    /// Path to converted generator safetensors (run `audioseal_convert` first).
    #[arg(long)]
    generator: PathBuf,
    /// Path to converted detector safetensors.
    #[arg(long)]
    detector: PathBuf,
    /// 16-bit message to embed.
    #[arg(long, default_value = "0xBEEF")]
    message: String,
    /// Force CPU device.
    #[arg(long)]
    cpu: bool,
    /// Optional input WAV (any sample rate, any channels). If provided,
    /// loaded and resampled to 16 kHz mono. Otherwise a synthetic
    /// pink-noise + tone burst signal is used (more speech-like than a
    /// pure sine but still out-of-distribution).
    #[arg(long)]
    wav: Option<PathBuf>,
    /// Optional output WAV path for the watermarked signal.
    #[arg(long)]
    out: Option<PathBuf>,
}

fn parse_message(s: &str) -> Result<u16> {
    let s = s.trim();
    let val = if let Some(rest) = s.strip_prefix("0x").or_else(|| s.strip_prefix("0X")) {
        u16::from_str_radix(rest, 16)?
    } else {
        s.parse::<u16>()?
    };
    Ok(val)
}

fn main() -> Result<()> {
    tracing_subscriber::fmt().init();
    let cli = Cli::parse();
    let device = if cli.cpu {
        Device::Cpu
    } else if candle_core::utils::metal_is_available() {
        Device::new_metal(0)?
    } else {
        Device::Cpu
    };
    let message = parse_message(&cli.message)?;
    println!("device: {:?}", device);
    println!("message: 0x{:04X} ({} bits)", message, MESSAGE_BITS);

    // Load weights — F32 since AudioSeal is small enough to leave un-cast.
    let gen_vb = unsafe {
        candle_nn::VarBuilder::from_mmaped_safetensors(&[&cli.generator], DType::F32, &device)
    }
    .context("opening generator safetensors")?;
    let det_vb = unsafe {
        candle_nn::VarBuilder::from_mmaped_safetensors(&[&cli.detector], DType::F32, &device)
    }
    .context("opening detector safetensors")?;

    let generator = Generator::new(gen_vb).context("Generator::new")?;
    let detector = Detector::new(det_vb).context("Detector::new")?;
    println!("loaded generator + detector successfully");

    // Source signal: real WAV (resampled to 16 kHz mono) if --wav is
    // provided, else a speech-like synthetic burst (multi-formant + noise).
    let signal: Vec<f32> = if let Some(wav) = cli.wav.as_ref() {
        let s = audio_io::load_mono_at_rate(wav, SAMPLE_RATE)?;
        println!(
            "loaded {}: {} samples ({:.2} s @ {} Hz)",
            wav.display(),
            s.len(),
            s.len() as f32 / SAMPLE_RATE as f32,
            SAMPLE_RATE
        );
        s
    } else {
        let n = SAMPLE_RATE as usize;
        let f1 = 200.0; // F1-ish
        let f2 = 800.0; // F2-ish
        (0..n)
            .map(|i| {
                let t = i as f32 / SAMPLE_RATE as f32;
                // Two formants + lightly-shaped pseudo-noise.
                let s1 = (2.0 * std::f32::consts::PI * f1 * t).sin();
                let s2 = (2.0 * std::f32::consts::PI * f2 * t).sin() * 0.5;
                let n_seed = (i as u32).wrapping_mul(2654435761);
                let n_val = (n_seed as f32 / u32::MAX as f32 - 0.5) * 0.3;
                (s1 + s2 + n_val) * 0.1
            })
            .collect()
    };
    let n = signal.len();
    let xs = Tensor::from_slice(&signal, (1, 1, n), &device)?;
    println!("input signal: {n} samples");

    // Generator forward — produces watermark residual.
    let residual = generator
        .forward(&xs, message as u32)
        .context("generator forward")?;
    let residual_shape = residual.dims().to_vec();
    let watermarked = (xs.clone() + &residual)?;
    println!(
        "generator output residual shape: {:?}, watermarked shape: {:?}",
        residual_shape,
        watermarked.dims()
    );

    // Detector forward on the watermarked signal.
    let logits = detector.forward(&watermarked).context("detector forward")?;
    println!("detector logits shape: {:?}", logits.dims());

    let (presence, decoded, mean_presence) = detector.decode(&logits)?;
    println!(
        "detector decode: mean_presence={mean_presence:.4} (>0.5 = watermarked)",
    );
    println!("decoded message: 0x{:04X}", decoded);
    let bits_correct = MESSAGE_BITS - (decoded ^ message).count_ones() as usize;
    println!("message bits matching: {bits_correct}/{MESSAGE_BITS}");

    let presence_len = presence.dim(candle_core::D::Minus1)?;
    println!("per-sample presence length: {presence_len}");

    // Optionally write the watermarked signal so we can A/B listen.
    if let Some(out_path) = cli.out.as_ref() {
        let wm_samples: Vec<f32> = watermarked
            .reshape((n,))?
            .to_dtype(DType::F32)?
            .to_vec1()?;
        audio_io::write_wav_mono(out_path, &wm_samples, SAMPLE_RATE)?;
        println!("wrote watermarked WAV to {}", out_path.display());
    }

    // Also exercise the public AudioSealWatermarker surface.
    let gen_vb2 = unsafe {
        candle_nn::VarBuilder::from_mmaped_safetensors(&[&cli.generator], DType::F32, &device)
    }?;
    let det_vb2 = unsafe {
        candle_nn::VarBuilder::from_mmaped_safetensors(&[&cli.detector], DType::F32, &device)
    }?;
    let mut wm = AudioSealWatermarker::from_var_builders(gen_vb2, det_vb2, device.clone(), message)?;
    wm.message = message;
    let embedded = wm.embed(&signal)?;
    let result = wm.detect(&embedded)?;
    println!(
        "Watermarker round-trip: mean_presence={:.4}, decoded=0x{:04X}",
        result.mean_presence,
        result.message.unwrap_or(0)
    );

    Ok(())
}