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

//! Phase 10.4 — apply SilentCipher to an arbitrary WAV. Mirror of
//! `audioseal_apply`. Loads the released 16 kHz checkpoint from
//! `sony/silentcipher`, embeds a payload, optionally re-detects to
//! verify round-trip.
//!
//! Usage (embed + verify):
//! ```bash
//! cargo run -p rtx-csm --release --features metal --example silentcipher_apply -- \
//!     --in /tmp/asr_test.flac --out /tmp/asr_test_silent.wav --payload 12345678
//! ```
//!
//! Usage (detect only):
//! ```bash
//! cargo run -p rtx-csm --release --features metal --example silentcipher_apply -- \
//!     --in /tmp/asr_test_silent.wav --detect-only
//! ```

use anyhow::{Context, Result};
use candle_core::Device;
use clap::Parser;
use hf_hub::api::sync::Api;
use rtx_csm::{
    audio_io,
    silentcipher::{SilentCipherConfig, SilentCipherWatermarker},
};
use std::path::PathBuf;
use std::time::Instant;

const REPO: &str = "sony/silentcipher";
const CKPT_DIR: &str = "16_khz/97561_iteration";
const SR: u32 = 16_000;

#[derive(Debug, Parser)]
struct Cli {
    /// Input WAV (any rate, mono OK; resampled to 16 kHz).
    #[arg(long = "in")]
    input: PathBuf,
    /// Output WAV. Optional in `--detect-only` mode.
    #[arg(long)]
    out: Option<PathBuf>,
    /// 30-bit payload to embed. Default 0 (no payload). 16 kHz model
    /// can carry up to ~23.78 bits (15 base-3 codes).
    #[arg(long, default_value_t = 0)]
    payload: u32,
    /// Skip embedding; just run detect on the input. Useful for
    /// verifying that a previously-watermarked file still carries the
    /// signature.
    #[arg(long)]
    detect_only: bool,
    /// Source sample rate of the input WAV. Default = read from file.
    #[arg(long)]
    source_rate: Option<u32>,
    #[arg(long)]
    cpu: bool,
}

fn main() -> Result<()> {
    tracing_subscriber::fmt().with_max_level(tracing::Level::WARN).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
    };
    eprintln!("device: {device:?}");

    // Load model.
    let api = Api::new()?;
    let repo = api.model(REPO.to_string());
    let enc_c = repo.get(&format!("{CKPT_DIR}/enc_c.ckpt"))?;
    let dec_c = repo.get(&format!("{CKPT_DIR}/dec_c.ckpt"))?;
    let dec_m_0 = repo.get(&format!("{CKPT_DIR}/dec_m_0.ckpt"))?;
    let cfg = SilentCipherConfig::sixteen_khz();
    let load_t = Instant::now();
    let wm = SilentCipherWatermarker::from_ckpts(cfg, &enc_c, &dec_c, &dec_m_0, &device)?;
    eprintln!("model loaded in {} ms", load_t.elapsed().as_millis());

    // Load + resample input to 16 kHz.
    let src_rate = cli.source_rate.unwrap_or_else(|| {
        // Read native rate via hound where possible — but we just use
        // load_mono_at_rate which resamples to whatever target we
        // pass. For round-trip output at the original rate, we read
        // the source rate manually.
        match hound::WavReader::open(&cli.input) {
            Ok(r) => r.spec().sample_rate,
            Err(_) => SR, // assume 16 kHz if we can't read header
        }
    });
    let src = audio_io::load_mono_at_rate(&cli.input, src_rate)?;
    let src_16k = audio_io::resample(&src, src_rate, SR)?;
    eprintln!(
        "loaded {} ({} samples @ {src_rate} Hz, {} samples @ {SR} Hz)",
        cli.input.display(),
        src.len(),
        src_16k.len()
    );

    if cli.detect_only {
        // Detect on the input directly.
        let t = Instant::now();
        let r = wm
            .detect(&src_16k)
            .context("silentcipher detect")?;
        let ms = t.elapsed().as_millis();
        let recovered = wm.decode_bits(&r.codes);
        println!();
        println!("=== detect-only ===");
        println!("input:        {}", cli.input.display());
        println!("audio:        {} samples ({:.2}s) at 16 kHz", src_16k.len(), src_16k.len() as f32 / SR as f32);
        println!("detect:       {ms} ms");
        println!("confidence:   {:.4}", r.confidence);
        println!("payload:      0x{recovered:08X} ({recovered})");
        println!("codes:        {:?}", r.codes);
        return Ok(());
    }

    // Embed.
    let codes = wm.encode_bits(cli.payload);
    eprintln!(
        "payload 0x{:08X} -> {} codes (each in 0..{}): {:?}",
        cli.payload,
        codes.len(),
        wm.cfg.message_dim,
        codes
    );
    let embed_t = Instant::now();
    let watermarked_16k = wm
        .embed(&src_16k, &codes)
        .context("silentcipher embed")?;
    let embed_ms = embed_t.elapsed().as_millis();
    eprintln!(
        "embed: {embed_ms} ms ({:.3}× realtime)",
        embed_ms as f32 / (src_16k.len() as f32 / SR as f32 * 1000.0)
    );

    // Resample back to source rate and write output.
    let out_path = cli
        .out
        .as_ref()
        .context("--out is required unless --detect-only is set")?;
    let watermarked = audio_io::resample(&watermarked_16k, SR, src_rate)?;
    let mut watermarked = watermarked;
    watermarked.truncate(src.len());
    audio_io::write_wav_mono(out_path.as_path(), &watermarked, src_rate)?;
    eprintln!("wrote {} ({} samples @ {} Hz)", out_path.display(), watermarked.len(), src_rate);

    // Round-trip detect on the resampled output (the realistic test).
    let probe_16k = audio_io::resample(&watermarked, src_rate, SR)?;
    let detect_t = Instant::now();
    let r = wm
        .detect(&probe_16k)
        .context("silentcipher detect")?;
    let detect_ms = detect_t.elapsed().as_millis();
    let recovered = wm.decode_bits(&r.codes);
    let matching = r
        .codes
        .iter()
        .zip(codes.iter())
        .filter(|(a, b)| a == b)
        .count();
    println!();
    println!("=== round-trip ===");
    println!("payload:      0x{:08X}  (expected)", cli.payload);
    println!("recovered:    0x{recovered:08X}");
    println!("codes match:  {matching} / {}", codes.len());
    println!("confidence:   {:.4}", r.confidence);
    println!("detect:       {detect_ms} ms");
    if matching == codes.len() && recovered == cli.payload {
        println!("PASS");
    } else {
        println!("WEAK — bit accuracy {:.0}%", 100.0 * matching as f32 / codes.len() as f32);
    }
    Ok(())
}