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

//! Phase 10.1 — download SilentCipher 16 kHz checkpoint from HuggingFace
//! (`sony/silentcipher`) and dump every tensor's shape + dtype. Used to
//! design the candle module shape and verify our port maps the right
//! weights.
//!
//! Architecture (from `SesameAILabs/silentcipher/src/silentcipher/model.py`):
//!   Three small networks of gated 2D convs on STFT:
//!     enc_c  : 3 layers, 1 -> 32 channels
//!     dec_c  : 4 layers, 96 -> 1 channels
//!     dec_m  : 10 layers, 1 -> 128 -> message_dim, plus Linear
//!
//! Each `Layer` is `Conv2d(...) * sigmoid(Conv2d(...))` followed by a
//! `BatchNorm2d`. The PyTorch state_dict per Layer carries:
//!     conv.weight, conv.bias
//!     gate.weight, gate.bias
//!     bn.weight, bn.bias, bn.running_mean, bn.running_var,
//!     bn.num_batches_tracked
//!
//! Weight files are PyTorch `.ckpt` (pickle), same format as AudioSeal's
//! `.pth`. We can read directly with `candle_core::pickle::read_all`.
//!
//! Usage:
//! ```bash
//! cargo run -p rtx-csm --release --example silentcipher_inspect
//! ```

use anyhow::{Context, Result};
use hf_hub::api::sync::Api;

const REPO: &str = "sony/silentcipher";
const CKPT_DIR: &str = "16_khz/97561_iteration";
const CKPT_FILES: &[&str] = &["enc_c.ckpt", "dec_c.ckpt", "dec_m_0.ckpt"];

fn main() -> Result<()> {
    let api = Api::new().context("hf_hub init")?;
    let repo = api.model(REPO.to_string());

    eprintln!("=== SilentCipher 16 kHz checkpoint inspector ===");
    eprintln!("repo:     {REPO}");
    eprintln!("ckpt dir: {CKPT_DIR}");
    eprintln!();

    // Download hparams.yaml first — it contains all the architecture
    // hyperparameters (n_fft, hop_length, message_dim, etc.).
    let hparams_path = repo
        .get(&format!("{CKPT_DIR}/hparams.yaml"))
        .context("download hparams.yaml")?;
    let hparams = std::fs::read_to_string(&hparams_path)?;
    eprintln!("--- hparams.yaml ---");
    eprintln!("{hparams}");

    // Each ckpt file is a PyTorch state_dict pickle. Use the same
    // candle_core::pickle::read_all path as audioseal_convert.rs.
    let dev = candle_core::Device::Cpu;
    for ckpt_file in CKPT_FILES {
        let path = repo
            .get(&format!("{CKPT_DIR}/{ckpt_file}"))
            .with_context(|| format!("download {ckpt_file}"))?;
        let size = std::fs::metadata(&path)?.len();
        println!();
        println!("=== {ckpt_file} ({:.2} MB) ===", size as f64 / 1e6);

        let tensors = candle_core::pickle::read_all(&path)
            .with_context(|| format!("pickle read {ckpt_file}"))?;
        let mut total_params: usize = 0;
        let mut grouped: std::collections::BTreeMap<String, Vec<(String, Vec<usize>)>> =
            Default::default();
        for (name, tensor) in tensors.iter() {
            let _ = dev; // suppress unused if device path not needed
            let shape = tensor.dims().to_vec();
            total_params += shape.iter().product::<usize>();
            let prefix = name
                .split('.')
                .take(2)
                .collect::<Vec<_>>()
                .join(".");
            grouped
                .entry(prefix)
                .or_default()
                .push((name.clone(), shape));
        }
        for (prefix, items) in grouped.iter() {
            println!("  [{prefix}]  ({} tensors)", items.len());
            for (name, shape) in items.iter().take(8) {
                println!("    {name:<55} {shape:?}");
            }
            if items.len() > 8 {
                println!("    ... {} more", items.len() - 8);
            }
        }
        println!(
            "  total params: {total_params} ({:.2} M)",
            total_params as f64 / 1e6
        );
    }

    println!();
    println!("=== summary ===");
    println!("Three small networks; combined params likely 5-10 M.");
    println!("Each layer is gated conv + BatchNorm2d. STFT params");
    println!("come from hparams.yaml above.");
    println!();
    println!("Next: implement Layer/Encoder/CarrierDecoder/MsgDecoder");
    println!("in src/silentcipher.rs (Phase 10.2).");

    Ok(())
}