updata index tts

This commit is contained in:
jhqxxx
2026-02-14 15:52:30 +08:00
parent 7c832e0ce8
commit 2c34fc2d79
40 changed files with 2602 additions and 540 deletions
+21 -10
View File
@@ -2,7 +2,11 @@ use std::collections::HashMap;
use aha::utils::{find_type_files, get_device, read_pth_tensor_info_cycle};
use anyhow::Result;
use candle_core::{Device, pickle::read_all_with_key, safetensors};
use candle_core::{
Device,
pickle::{read_all_with_key, read_pth_tensor_info},
safetensors,
};
use candle_nn::VarBuilder;
#[test]
@@ -203,19 +207,26 @@ fn index_tts2_weight() -> Result<()> {
// RUST_BACKTRACE=1 cargo test -F cuda index_tts2_weight -r -- --nocapture
let save_dir: String =
aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?;
let model_path = format!("{}/IndexTeam/IndexTTS-2/", save_dir);
let s2mel_path = model_path+ "/s2mel.pth";
let model_path = format!("{}/IndexTeam/IndexTTS-2/", save_dir);
let bigvgan_path = format!(
"{}/nv-community/bigvgan_v2_22khz_80band_256x/bigvgan_generator.pt",
save_dir
);
// let gpt_path = model_path+ "/gpt.pth";
// let spk_matrix_path = model_path+ "/feat1.pt";
// let s2mel_path = model_path+ "/s2mel.pth";
// let wac2vec2_path = model_path+ "/wav2vec2bert_stats.pt";
// let model_path = format!("{}/iic/speech_campplus_sv_zh-cn_16k-common/", save_dir);
// let campplus_path = model_path+ "/campplus_cn_common.bin";
// let model_list = find_type_files(&model_path, "safetensors")?;
let model_list = vec![s2mel_path];
// let mut dict_to_hashmap = HashMap::new();
// let mut dtype = candle_core::DType::F32;
let model_list = vec![bigvgan_path];
// // let mut dict_to_hashmap = HashMap::new();
// // let mut dtype = candle_core::DType::F32;
for m in model_list {
// let dict = read_all_with_key(m, Some("state_dict"))?;
// let dict = read_all_with_key(m, Some("net"))?;
let dict = read_pth_tensor_info_cycle(m, Some("net.cfm"))?;
let dict = read_all_with_key(m, Some("generator"))?;
// let dict = read_pth_tensor_info_cycle(m, Some("net.cfm"))?;
// dtype = dict[0].1.dtype();
for (k, v) in dict {
// if k.contains("model") {
@@ -230,9 +241,9 @@ fn index_tts2_weight() -> Result<()> {
// let model_list = vec![semantic_codec_path];
// for m in model_list {
// let weights = safetensors::load(m, &device)?;
// for (key, tensor) in weights.iter() {
// for (key, tensor) in weights.iter() {
// println!("=== {} === {:?}", key, tensor.shape());
// }
// }
Ok(())
}
}