inferencelayer 0.2.3

//! `qwen3tts-dbg` — Phase-0 manifest dumper for the Qwen3-TTS checkpoints.
//!
//! Rust-only (no torch): parses each `.safetensors` header and prints a COLLAPSED tensor manifest
//! (per-layer repetition folded to `...layers.{i}...` with a count), plus the parsed/validated
//! config identity. This is how the P1+ loaders learn ground-truth tensor names/shapes instead of
//! guessing PyTorch module paths. Manifest tooling is shared with the chatterbox port.
//!
//! Usage:
//!   cargo run --release --features "cli encoder-cpu" --bin qwen3tts-dbg -- [DIR]
//!   (DIR defaults to ~/.cache/inferencelayer/qwen3tts/1.7b-base)

use anyhow::Result;
use inferencelayer::chatterbox::{TensorInfo, read_manifest};
use inferencelayer::qwen3tts::{Qwen3TtsConfig, cfg, default_root};
use std::collections::BTreeMap;
use std::path::PathBuf;

fn main() -> Result<()> {
    let dir: PathBuf = std::env::args()
        .nth(1)
        .map(PathBuf::from)
        .unwrap_or_else(|| default_root().join("1.7b-base"));
    println!("qwen3tts checkpoint dir: {}\n", dir.display());

    match Qwen3TtsConfig::load(&dir) {
        Ok(c) => println!(
            "config OK: size={} type={} talker h={} spk_enc={:?} langs={:?} spk_ids={}\n",
            c.top.tts_model_size,
            c.top.tts_model_type,
            c.top.talker_config.hidden_size,
            c.top.speaker_encoder_config,
            c.top
                .talker_config
                .codec_language_id
                .keys()
                .collect::<Vec<_>>(),
            c.top.talker_config.spk_id.len(),
        ),
        Err(e) => println!("config: [ERROR: {e:#}]\n"),
    }

    for (label, path) in [
        (
            "talker (+ code predictor + speaker encoder)",
            dir.join(cfg::FILE_MODEL),
        ),
        (
            "codec (speech_tokenizer)",
            dir.join(cfg::DIR_CODEC).join(cfg::FILE_MODEL),
        ),
    ] {
        if !path.exists() {
            println!("── {label}: {}  [not present]\n", path.display());
            continue;
        }
        match read_manifest(&path) {
            Ok(m) => dump(label, &m),
            Err(e) => println!("── {label}: [ERROR: {e:#}]\n"),
        }
    }
    Ok(())
}

/// Fold digit-runs between dots to `{i}` so per-layer tensors collapse into one templated row.
fn template(name: &str) -> String {
    name.split('.')
        .map(|seg| {
            if !seg.is_empty() && seg.bytes().all(|b| b.is_ascii_digit()) {
                "{i}"
            } else {
                seg
            }
        })
        .collect::<Vec<_>>()
        .join(".")
}

struct Group {
    count: usize,
    dtype: String,
    example_shape: Vec<usize>,
}

fn dump(label: &str, manifest: &[TensorInfo]) {
    let total_numel: usize = manifest.iter().map(|t| t.numel()).sum();
    let mut groups: BTreeMap<String, Group> = BTreeMap::new();
    for t in manifest {
        let g = groups.entry(template(&t.name)).or_insert_with(|| Group {
            count: 0,
            dtype: t.dtype.clone(),
            example_shape: t.shape.clone(),
        });
        g.count += 1;
    }
    println!("── {label}");
    println!(
        "   {} tensors, {} templated groups, {:.1}M params",
        manifest.len(),
        groups.len(),
        total_numel as f64 / 1e6
    );
    for (tmpl, g) in &groups {
        let mult = if g.count > 1 {
            format!("×{}", g.count)
        } else {
            String::new()
        };
        println!(
            "   {:<70} {:<5} {:<4} {:?}",
            tmpl, g.dtype, mult, g.example_shape
        );
    }
    println!();
}