diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index da92d3d..a3950df 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -31,7 +31,7 @@ jobs: - name: Install make dependencies run: | sudo apt-get update - sudo apt-get install -y build-essential + sudo apt-get install -y build-essential cmake - name: Install FFmpeg dependencies run: | @@ -61,7 +61,7 @@ jobs: - name: Install make dependencies run: | sudo apt-get update - sudo apt-get install -y build-essential + sudo apt-get install -y build-essential cmake - name: Install FFmpeg development packages run: | diff --git a/Cargo.lock b/Cargo.lock index 1791314..216a541 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -52,8 +52,10 @@ dependencies = [ "num", "rayon", "realfft", + "regex", "reqwest 0.12.24", "rocket", + "sentencepiece", "serde", "serde_json", "serde_yaml", @@ -3305,6 +3307,29 @@ dependencies = [ "syn", ] +[[package]] +name = "prost" +version = "0.14.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d2ea70524a2f82d518bce41317d0fae74151505651af45faf1ffbd6fd33f0568" +dependencies = [ + "bytes", + "prost-derive", +] + +[[package]] +name = "prost-derive" +version = "0.14.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "27c6023962132f4b30eb4c172c91ce92d933da334c59c23cddee82358ddafb0b" +dependencies = [ + "anyhow", + "itertools 0.14.0", + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "psl-types" version = "2.0.11" @@ -3667,9 +3692,9 @@ dependencies = [ [[package]] name = "regex" -version = "1.12.2" +version = "1.12.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "843bc0191f75f3e22651ae5f1e72939ab2f72a4bc30fa80a066bd66edefc24d4" +checksum = "e10754a14b9137dd7b1e3e5b0493cc9171fdd105e0ab477f51b72e7f3ac0e276" dependencies = [ "aho-corasick", "memchr", @@ -4104,6 +4129,32 @@ dependencies = [ "libc", ] +[[package]] +name = "sentencepiece" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3baa1506c7718f6b70bcac5475e563176dd51bb669dae38181abbf3070517453" +dependencies = [ + "libc", + "num-derive", + "num-traits", + "prost", + "prost-derive", + "sentencepiece-sys", + "thiserror 2.0.17", +] + +[[package]] +name = "sentencepiece-sys" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fa4f9b54dc005df8ec1c3f9e2347cea6b6657ec5460977a982e7a4e09ef49411" +dependencies = [ + "cc", + "cmake", + "pkg-config", +] + [[package]] name = "seq-macro" version = "0.3.6" diff --git a/Cargo.toml b/Cargo.toml index 3f5c6e3..6a3e13c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -40,6 +40,8 @@ serde_yaml = "0.9.34" zip = "7.2.0" half = "2.7.1" byteorder = "1.5.0" +sentencepiece = "0.13.1" +regex = "1.12.3" [features] flash-attn = ["candle-flash-attn"] diff --git a/README.md b/README.md index 6a08900..fdce98d 100644 --- a/README.md +++ b/README.md @@ -62,6 +62,13 @@ ⭐ 如果这个项目对你有帮助，请给我们一个 Star！ ## 环境依赖 +* cmake: + * ubuntu/WSL + ```bash + sudo apt-get update + sudo apt-get install cmake + ``` + * 启用ffmpeg的feature时： * ubuntu/WSL ```bash diff --git a/src/models/common/mod.rs b/src/models/common/mod.rs index 2b4f46f..c3222c3 100644 --- a/src/models/common/mod.rs +++ b/src/models/common/mod.rs @@ -2,9 +2,9 @@ use anyhow::{Result, anyhow}; use candle_core::{D, IndexOp, Tensor}; use candle_nn::{ Activation, BatchNorm, BatchNormConfig, Conv1d, Conv1dConfig, Conv2d, Conv2dConfig, - ConvTranspose1d, ConvTranspose1dConfig, Embedding, LayerNorm, LayerNormConfig, Linear, Module, - ModuleT, RmsNorm, VarBuilder, batch_norm, conv1d, conv1d_no_bias, conv2d, conv2d_no_bias, - embedding, layer_norm, linear_b, linear_no_bias, ops::sigmoid, rms_norm, + ConvTranspose1d, ConvTranspose1dConfig, Embedding, GroupNorm, Init, LayerNorm, LayerNormConfig, + Linear, Module, ModuleT, RmsNorm, VarBuilder, batch_norm, conv1d, conv1d_no_bias, conv2d, + conv2d_no_bias, embedding, layer_norm, linear_b, linear_no_bias, ops::sigmoid, rms_norm, }; use rayon::iter::{IndexedParallelIterator, IntoParallelRefIterator, ParallelIterator}; @@ -464,9 +464,10 @@ impl NaiveAttnTwoLinearMLPBlock { linear2_pp_name, )?; - let input_layernorm = get_layer_norm(vb.pp(input_norm_pp_name), norm_eps, hidden_size)?; + let input_layernorm = + get_layer_norm(vb.pp(input_norm_pp_name), norm_eps, hidden_size, true)?; let post_attention_layernorm = - get_layer_norm(vb.pp(post_norm_pp_name), norm_eps, hidden_size)?; + get_layer_norm(vb.pp(post_norm_pp_name), norm_eps, hidden_size, true)?; Ok(Self { self_attn, mlp, @@ -689,16 +690,26 @@ pub fn get_conv1d( Ok(conv1d) } -pub fn get_layer_norm(vb: VarBuilder, eps: f64, dim: usize) -> Result { +pub fn get_layer_norm(vb: VarBuilder, eps: f64, dim: usize, affine: bool) -> Result { let ln_config = LayerNormConfig { eps, remove_mean: true, // true for layernorm, false for RMSNorm - affine: true, // true for with bias, false for without bias + affine, // true for with bias, false for without bias }; let norm = layer_norm(dim, ln_config, vb)?; Ok(norm) } +pub fn get_layer_norm_without_weight( + vb: VarBuilder, + eps: f64, + dim: usize, +) -> Result { + let weight = Tensor::ones(dim, vb.dtype(), vb.device())?; + let bias = Tensor::zeros(dim, vb.dtype(), vb.device())?; + Ok(LayerNorm::new(weight, bias, eps)) +} + pub fn get_batch_norm(vb: VarBuilder, eps: f64, dim: usize, affine: bool) -> Result { let bn_config = BatchNormConfig { eps, @@ -1167,3 +1178,9 @@ impl WNLinear { Ok(x) } } + +pub fn mish(xs: &Tensor) -> Result { + let tanh = xs.exp()?.affine(1.0, 1.0)?.log()?.tanh()?; + let xs = xs.mul(&tanh)?; + Ok(xs) +} diff --git a/src/models/deepseek_ocr/model.rs b/src/models/deepseek_ocr/model.rs index 493d214..6c9639f 100644 --- a/src/models/deepseek_ocr/model.rs +++ b/src/models/deepseek_ocr/model.rs @@ -252,7 +252,7 @@ impl Block { window_size: usize, input_size: Option<(usize, usize)>, ) -> Result { - let norm1 = get_layer_norm(vb.pp("norm1"), eps, dim)?; + let norm1 = get_layer_norm(vb.pp("norm1"), eps, dim, true)?; let input_size = if window_size == 0 { input_size } else { @@ -266,7 +266,7 @@ impl Block { use_rel_pos, input_size, )?; - let norm2 = get_layer_norm(vb.pp("norm2"), eps, dim)?; + let norm2 = get_layer_norm(vb.pp("norm2"), eps, dim, true)?; let mlp_dim = (dim as f32 * mlp_ratio) as usize; let mlp = TwoLinearMLP::new(vb.pp("mlp"), dim, mlp_dim, dim, act, true, "lin1", "lin2")?; Ok(Self { @@ -674,8 +674,8 @@ impl NoTPTransformerBlock { ) -> Result { let self_attn = NoTPAttention::new(vb.pp("self_attn"), hidden_size, num_heads)?; let mlp = NoTPFeedForward::new(vb.pp("mlp"), hidden_size, ffn_hidden_size)?; - let layer_norm1 = get_layer_norm(vb.pp("layer_norm1"), eps, hidden_size)?; - let layer_norm2 = get_layer_norm(vb.pp("layer_norm2"), eps, hidden_size)?; + let layer_norm1 = get_layer_norm(vb.pp("layer_norm1"), eps, hidden_size, true)?; + let layer_norm2 = get_layer_norm(vb.pp("layer_norm2"), eps, hidden_size, true)?; Ok(Self { self_attn, mlp, @@ -764,7 +764,7 @@ impl VitModel { ffn_hidden_size, eps, )?; - let pre_layrnorm = get_layer_norm(vb.pp("pre_layrnorm"), eps, hidden_size)?; + let pre_layrnorm = get_layer_norm(vb.pp("pre_layrnorm"), eps, hidden_size, true)?; Ok(Self { embeddings, transformer, diff --git a/src/models/fun_asr_nano/model.rs b/src/models/fun_asr_nano/model.rs index 755ecff..779cc38 100644 --- a/src/models/fun_asr_nano/model.rs +++ b/src/models/fun_asr_nano/model.rs @@ -227,8 +227,8 @@ impl EncoderLayerSANM { "w_1", "w_2", )?; - let norm1 = get_layer_norm(vb.pp("norm1"), 1e-5, in_dim)?; - let norm2 = get_layer_norm(vb.pp("norm2"), 1e-5, hidden_dim)?; + let norm1 = get_layer_norm(vb.pp("norm1"), 1e-5, in_dim, true)?; + let norm2 = get_layer_norm(vb.pp("norm2"), 1e-5, hidden_dim, true)?; let concat_linear = if concat_after { let lin = linear(hidden_dim * 2, hidden_dim, vb.pp("concat_linear"))?; Some(lin) @@ -400,8 +400,8 @@ impl SenseVoiceEncoderSmall { )?; tp_encoders.push(tp_blocks_i); } - let after_norm = get_layer_norm(vb.pp("after_norm"), 1e-5, output_size)?; - let tp_norm = get_layer_norm(vb.pp("tp_norm"), 1e-5, output_size)?; + let after_norm = get_layer_norm(vb.pp("after_norm"), 1e-5, output_size, true)?; + let tp_norm = get_layer_norm(vb.pp("tp_norm"), 1e-5, output_size, true)?; let scaling = (output_size as f64).powf(0.5); Ok(Self { embed, @@ -468,8 +468,8 @@ impl AdaptorEncoderLayer { "w_1", "w_2", )?; - let norm1 = get_layer_norm(vb.pp("norm1"), 1e-5, llm_dim)?; - let norm2 = get_layer_norm(vb.pp("norm2"), 1e-5, llm_dim)?; + let norm1 = get_layer_norm(vb.pp("norm1"), 1e-5, llm_dim, true)?; + let norm2 = get_layer_norm(vb.pp("norm2"), 1e-5, llm_dim, true)?; let concat_linear = if concat_after { let lin = linear(llm_dim * 2, llm_dim, vb.pp("concat_linear"))?; Some(lin) diff --git a/src/models/glm_asr_nano/model.rs b/src/models/glm_asr_nano/model.rs index 21e5023..6f96f96 100644 --- a/src/models/glm_asr_nano/model.rs +++ b/src/models/glm_asr_nano/model.rs @@ -129,11 +129,12 @@ impl GlmAsrEncoderLayer { "fc2", )?; let input_layernorm = - get_layer_norm(vb.pp("input_layernorm"), 1e-5, audio_cfg.hidden_size)?; + get_layer_norm(vb.pp("input_layernorm"), 1e-5, audio_cfg.hidden_size, true)?; let post_attention_layernorm = get_layer_norm( vb.pp("post_attention_layernorm"), 1e-5, audio_cfg.hidden_size, + true, )?; Ok(Self { self_attn, @@ -202,7 +203,7 @@ impl GlmAsrEncoder { let layer_i = GlmAsrEncoderLayer::new(vb_layers.pp(i), audio_cfg)?; layers.push(layer_i); } - let norm = get_layer_norm(vb.pp("norm"), 1e-5, audio_cfg.hidden_size)?; + let norm = get_layer_norm(vb.pp("norm"), 1e-5, audio_cfg.hidden_size, true)?; let dim = (audio_cfg.head_dim as f64 * audio_cfg.partial_rotary_factor) as usize; let rotary_emb = RoPE::new(dim, audio_cfg.rope_parameters.rope_theta, vb.device())?; Ok(Self { diff --git a/src/models/index_tts2/generate.rs b/src/models/index_tts2/generate.rs index 923e7a3..5c4f63d 100644 --- a/src/models/index_tts2/generate.rs +++ b/src/models/index_tts2/generate.rs @@ -1,14 +1,29 @@ use aha_openai_dive::v1::resources::chat::ChatCompletionParameters; -use anyhow::Result; +use anyhow::{Result, anyhow}; use candle_core::{DType, Device}; +use sentencepiece::SentencePieceProcessor; use crate::{ - models::index_tts2::{config::IndexTTS2Config, processor::IndexTTS2Processor}, - utils::{get_default_save_dir, get_device, get_dtype}, + models::index_tts2::{ + config::IndexTTS2Config, + model::IndexTTS2Model, + processor::IndexTTS2Processor, + utils::{TextNormalizer, tokenize_by_cjk_char}, + }, + tokenizer::sentencepiece_encode, + utils::{ + audio_utils::extract_audio_url, extract_user_text, get_default_save_dir, get_device, + get_dtype, + }, }; + pub struct IndexTTS2Generate { processor: IndexTTS2Processor, + tokenizer: SentencePieceProcessor, config: IndexTTS2Config, + cache_spk_audio_prompt: Option, + model: IndexTTS2Model, + device: Device, } impl IndexTTS2Generate { @@ -18,12 +33,46 @@ impl IndexTTS2Generate { let config: IndexTTS2Config = serde_yaml::from_slice(&std::fs::read(config_path)?)?; let device = get_device(device); let dtype = get_dtype(dtype, "bf16"); - let processor = IndexTTS2Processor::new(path, &save_dir, &config, &device, dtype)?; - - Ok(Self { config, processor }) + let processor = IndexTTS2Processor::new(&device)?; + let bpe_path = path.to_string() + "/bpe.model"; + let tokenizer = SentencePieceProcessor::open(bpe_path) + .map_err(|e| anyhow!(format!("load bpe,model file error:{}", e)))?; + let model = IndexTTS2Model::new(path, &save_dir, &config, &device, dtype)?; + Ok(Self { + processor, + tokenizer, + config, + cache_spk_audio_prompt: None, + model, + device, + }) } + + pub fn use_prompt(&self, mes: &ChatCompletionParameters) -> bool { + if let Some(cache) = &self.cache_spk_audio_prompt { + let audio_vec = extract_audio_url(mes); + if audio_vec.len() == 0 { + true + } else { + if cache.eq(&audio_vec[0]) { true } else { false } + } + } else { + false + } + } + pub fn generate(&mut self, mes: ChatCompletionParameters) -> Result<()> { - let _ = self.processor.process_info(&mes)?; + let text = extract_user_text(&mes)?; + let text = tokenize_by_cjk_char(&text, true); + let input_ids = sentencepiece_encode(&text, &self.tokenizer, &self.device)?; + let (audio_22k, audio_16k) = if self.use_prompt(&mes) { + (None, None) + } else { + let (audio_22k, audio_16k, prompt) = self.processor.process_info(&mes)?; + self.cache_spk_audio_prompt = Some(prompt); + (Some(audio_22k), Some(audio_16k)) + }; + let _ = self.model.forward(&input_ids, audio_22k.as_ref(), audio_16k.as_ref())?; Ok(()) } } diff --git a/src/models/index_tts2/model.rs b/src/models/index_tts2/model.rs index 3f9b692..a87473f 100644 --- a/src/models/index_tts2/model.rs +++ b/src/models/index_tts2/model.rs @@ -1,20 +1,33 @@ -use anyhow::Result; -use candle_core::{D, Tensor}; +use anyhow::{Result, anyhow}; +use candle_core::{D, DType, Device, IndexOp, Tensor, pickle::read_all_with_key}; use candle_nn::{ - Conv1d, Embedding, LayerNorm, Linear, Module, RmsNorm, VarBuilder, embedding, linear, linear_b, - ops::sigmoid, rms_norm, + Activation, Conv1d, Embedding, GroupNorm, Init, LayerNorm, Linear, Module, RmsNorm, VarBuilder, + embedding, group_norm, linear, linear_b, ops::sigmoid, rms_norm, }; use crate::{ models::{ + campplus::CAMPPlus, common::{ GateUpDownMLP, QKVCatAttention, TwoLinearMLP, WNConv1d, WNLinear, get_conv1d, - get_layer_norm, + get_layer_norm, get_layer_norm_without_weight, mish, }, - index_tts2::config::{DiTModelArgs, S2MelConfig}, + feature_extractor::seamless_m4t_feature_extractor::SeamlessM4TFeatureExtractor, + index_tts2::config::{DiTModelArgs, IndexTTS2Config, PreprocessParams, S2MelConfig}, + mask_gct::model::RepCodec, + w2v_bert_2_0::model::W2VBert2_0Model, }, position_embed::rope::RoPE, - utils::tensor_utils::{pad_reflect_last_dim, split_tensor_with_size}, + utils::{ + audio_utils::{ + create_hann_window, get_waveform_and_window_properties, kaldi_fbank, + kaldi_get_mel_banks, mel_filter_bank, torch_stft, + }, + get_vb_model_path, read_pth_tensor_info_cycle, + tensor_utils::{ + interpolate_nearest_1d, pad_reflect_last_dim, sequence_mask, split_tensor_with_size, + }, + }, }; pub struct AdaptiveLayerNorm { project_layer: Linear, @@ -219,7 +232,7 @@ impl TimestepEmbedder { candle_nn::Activation::Silu, true, "0", - "1", + "2", )?; let scale = 1000.0; let half = frequency_embedding_size / 2; @@ -338,7 +351,7 @@ impl Wavenet { let in_layer = SConv1d::new( vb_layers.pp(i), hidden_c, - 1 * hidden_c, + 2 * hidden_c, ks, 1, dilation, @@ -385,7 +398,7 @@ impl Wavenet { Ok(acts) } - pub fn forward(self, xs: &Tensor, x_mask: &Tensor, g: Option<&Tensor>) -> Result { + pub fn forward(&self, xs: &Tensor, x_mask: &Tensor, g: Option<&Tensor>) -> Result { let mut output = xs.zeros_like()?; let g = if let Some(g) = g && let Some(cond_layer) = &self.cond_layer @@ -432,7 +445,7 @@ impl FinalLayer { patch_size: usize, out_c: usize, ) -> Result { - let norm_final = get_layer_norm(vb.pp("norm_final"), 1e-6, hidden_size)?; + let norm_final = get_layer_norm_without_weight(vb.pp("norm_final"), 1e-6, hidden_size)?; let linear = WNLinear::new( vb.pp("linear"), hidden_size, @@ -485,6 +498,8 @@ pub struct DiT { time_as_token: bool, style_as_token: bool, uvit_skip_connection: bool, + transformer_style_condition: bool, + long_skip_connection: bool, } impl DiT { @@ -594,24 +609,451 @@ impl DiT { time_as_token, style_as_token, uvit_skip_connection, + transformer_style_condition: config.di_t.style_condition, + long_skip_connection: config.di_t.long_skip_connection, }) } + + pub fn forward( + &self, + xs: &Tensor, + prompt_xs: &Tensor, + x_lens: &Tensor, + t: &Tensor, + style: Option<&Tensor>, + cond: &Tensor, + ) -> Result { + let (_, _, t_dim) = xs.dims3()?; + let t1 = self.t_embedder.forward(t)?; + let cond = self.cond_projection.forward(cond)?; + let xs = xs.transpose(1, 2)?; + let prompt_xs = prompt_xs.transpose(1, 2)?; + let mut x_in = Tensor::cat(&[&xs, &prompt_xs, &cond], D::Minus1)?; + if self.transformer_style_condition + && !self.style_as_token + && let Some(style) = style + { + let style = style.unsqueeze(1)?.repeat((1, t_dim, 1))?; + x_in = Tensor::cat(&[&x_in, &style], D::Minus1)?; + } + x_in = self.cond_x_merge_linear.forward(&x_in)?; + // if self.style_as_token + // && let Some(style_in) = self.style_in.as_ref() + // { + // let style = style_in.forward(style)?.unsqueeze(1)?; + // x_in = Tensor::cat(&[&style, &x_in], 1)?; + // } + // if self.time_as_token { + // let t1 = t1.unsqueeze(1)?; + // x_in = Tensor::cat(&[&t1, &x_in], 1)?; + // } + // let mut x_lens = x_lens.clone(); + // if self.style_as_token { + // x_lens = x_lens.affine(1.0, 1.0)?; + // } + // if self.time_as_token { + // x_lens = x_lens.affine(1.0, 0.0)?; + // } + let x_mask = sequence_mask(&x_lens, Some(x_in.dim(1)? as u32))? + .to_device(xs.device())? + .unsqueeze(1)?; + let mut x_res = self.transformer.forward(&x_in, &t1.unsqueeze(1)?, None)?; + // if self.time_as_token { + // let last_dim = x_res.dim(D::Minus1)?; + // x_res = x_res.narrow(D::Minus1, 1, last_dim - 1)?; + // } + // if self.style_as_token { + // let last_dim = x_res.dim(D::Minus1)?; + // x_res = x_res.narrow(D::Minus1, 1, last_dim - 1)?; + // } + if self.long_skip_connection { + x_res = self + .skip_linear + .forward(&Tensor::cat(&[&x_res, &xs], D::Minus1)?)?; + } + let xs = self.conv1.forward(&x_res)?; + let xs = xs.transpose(1, 2)?; + let t2 = self.t_embedder2.forward(t)?; + let xs = self + .wavenet + .forward(&xs, &x_mask, Some(&t2.unsqueeze(2)?))? + .transpose(1, 2)? + .broadcast_add(&self.res_projection.forward(&x_res)?)?; + let xs = self.final_layer.forward(&xs, &t1)?.transpose(1, 2)?; + let xs = self.conv2.forward(&xs)?; + Ok(xs) + } } pub struct CFM { + in_channels: usize, estimator: DiT, + // criterion: l1Loss + sigma_min: f32, +} + +impl CFM { + pub fn new(vb: VarBuilder, config: &S2MelConfig) -> Result { + let in_channels = config.di_t.in_channels; + let sigma_min = 1e-6; + let estimator = DiT::new(vb.pp("estimator"), config)?; + Ok(Self { + in_channels, + estimator, + sigma_min, + }) + } +} + +pub struct InterpolateModule { + conv1d: Conv1d, + norm: GroupNorm, +} + +impl InterpolateModule { + pub fn new(vb: &VarBuilder, index: usize, channels: usize, groups: usize) -> Result { + let start_index = index * 3; + let conv1d = get_conv1d(vb.pp(start_index), channels, channels, 3, 1, 1, 1, 1, true)?; + let norm = group_norm(groups, channels, 1e-5, vb.pp(start_index + 1))?; + Ok(Self { conv1d, norm }) + } + + pub fn forward(&self, xs: &Tensor) -> Result { + let xs = self.conv1d.forward(xs)?; + let xs = self.norm.forward(&xs)?; + let xs = mish(&xs)?; + Ok(xs) + } +} + +pub struct InterpolateRegulator { + sampling_ratios: Vec, + out_channels: usize, + model0_11: Vec, + model_12: Conv1d, + embedding: Embedding, + mask_token: Tensor, + quantizer_dropout: f32, + content_in_proj: Linear, + n_codebooks: usize, + interpolate: bool, +} + +impl InterpolateRegulator { + pub fn new( + vb: VarBuilder, + channels: usize, + sampling_ratios: Vec, + is_discrete: bool, + in_channels: usize, + vector_quantize: bool, + codebook_size: usize, + out_channels: Option, + groups: usize, + n_codebooks: usize, + quantizer_dropout: f32, + f0_condition: bool, + n_f0_bins: usize, + ) -> Result { + let out_channels = out_channels.unwrap_or(channels); + let vb_model = vb.pp("model"); + let interpolate = true; + let mut model0_11 = vec![]; + for (index, _) in sampling_ratios.iter().enumerate() { + let inter = InterpolateModule::new(&vb_model, index, channels, groups)?; + model0_11.push(inter); + } + let model_12 = get_conv1d( + vb_model.pp("12"), + channels, + out_channels, + 1, + 0, + 1, + 1, + 1, + true, + )?; + let embedding = embedding(codebook_size, channels, vb.pp("embedding"))?; + let mask_token = vb.get_with_hints((1, channels), "mask_token", Init::Const(0.0))?; + let content_in_proj = linear(in_channels, channels, vb.pp("content_in_proj"))?; + Ok(Self { + sampling_ratios, + out_channels, + model0_11, + model_12, + embedding, + mask_token, + quantizer_dropout, + content_in_proj, + n_codebooks, + interpolate, + }) + } + pub fn forward(&self, x: &Tensor, y_lens: &Tensor) -> Result { + let mut xs = self.content_in_proj.forward(x)?; + xs = xs.transpose(1, 2)?.contiguous()?; + if self.interpolate { + let size = y_lens.max_all()?.to_scalar::()? as usize; + xs = interpolate_nearest_1d(&xs, size)?; + } + for model_i in self.model0_11.iter() { + xs = model_i.forward(&xs)?; + } + xs = self.model_12.forward(&xs)?.transpose(1, 2)?.contiguous()?; + Ok(xs) + } } pub struct MyModel { cfm: CFM, + length_regulator: InterpolateRegulator, } -pub struct IndexTTS2 { - cache_spk_cond: Option, - cache_s2mel_style: Option, - cache_s2mel_prompt: Option, - cache_spk_audio_prompt: Option, - cache_emo_cond: Option, - cache_emo_audio_prompt: Option, - cache_mel: Option, +impl MyModel { + pub fn new( + model_path: &str, + config: &S2MelConfig, + dtype: DType, + device: &Device, + ) -> Result { + let s2mel_path = model_path.to_string() + "/s2mel.pth"; + let length_regulator_dict = + read_pth_tensor_info_cycle(s2mel_path.clone(), Some("net.length_regulator"))?; + let length_regulator_vb = VarBuilder::from_tensors(length_regulator_dict, dtype, device); + let length_regulator = InterpolateRegulator::new( + length_regulator_vb, + config.length_regulator.channels, + config.length_regulator.sampling_ratios.clone(), + config.length_regulator.is_discrete, + config.length_regulator.in_channels, + config.length_regulator.vector_quantize, + config.length_regulator.content_codebook_size, + None, + 1, + config.length_regulator.n_codebooks, + config.length_regulator.quantizer_dropout, + config.length_regulator.f0_condition, + config.length_regulator.n_f0_bins, + )?; + let cfm_dict = read_pth_tensor_info_cycle(s2mel_path.clone(), Some("net.cfm"))?; + let cfm_vb = VarBuilder::from_tensors(cfm_dict, dtype, device); + let cfm = CFM::new(cfm_vb, config)?; + Ok(Self { + cfm, + length_regulator, + }) + } + + pub fn length_regulator_forward( + &self, + s_ori: &Tensor, + target_lengths: &Tensor, + ) -> Result { + let xs = self.length_regulator.forward(s_ori, target_lengths)?; + Ok(xs) + } +} + +pub struct IndexTTS2Cache { + pub cache_spk_cond: Tensor, + pub cache_s2mel_style: Tensor, + pub cache_s2mel_prompt: Tensor, + pub cache_mel: Tensor, + // pub cache_emo_cond: Tensor, + // pub cache_emo_audio_prompt: Tensor, +} + +pub struct IndexTTS2Model { + cache: Option, + feature_extractor: SeamlessM4TFeatureExtractor, + semantic_model: W2VBert2_0Model, + semantic_mean: Tensor, + semantic_std: Tensor, + semantic_codec: RepCodec, + s2mel_filters: Tensor, + s2mel_windows: Tensor, + s2mel_preprocess_params: PreprocessParams, + window_shift: usize, + window_size: usize, + padded_window_size: usize, + mel_energies: Tensor, + campplus_model: CAMPPlus, + s2mel: MyModel, +} + +impl IndexTTS2Model { + pub fn new( + path: &str, + save_dir: &str, + config: &IndexTTS2Config, + device: &Device, + dtype: DType, + ) -> Result { + let feature_extractor = SeamlessM4TFeatureExtractor::new( + // 80, + 80, + crate::utils::tensor_utils::PaddingSide::Right, + 1.0, + 16000, + 2, + device, + )?; + let w2vbert2_path = save_dir.to_string() + "/facebook/w2v-bert-2.0"; + let semantic_model = W2VBert2_0Model::init(&w2vbert2_path, device, dtype)?; + let semantic_mean_var_path = path.to_string() + "/" + &config.w2v_stat; + let dict = read_all_with_key(semantic_mean_var_path, None)?; + let mut semantic_mean = Tensor::new(0.0, device)?.to_dtype(dtype)?; + let mut semantic_std = Tensor::new(1.0, device)?.to_dtype(dtype)?; + for (k, v) in dict { + if k.eq("mean") { + semantic_mean = v.to_device(device)?.to_dtype(dtype)?; + } else if k.eq("var") { + semantic_std = v.to_device(device)?.to_dtype(dtype)?.sqrt()?; + } + } + + let semantic_codec_path = + save_dir.to_string() + "/amphion/MaskGCT/semantic_codec/model.safetensors"; + let vb = + unsafe { VarBuilder::from_mmaped_safetensors(&[semantic_codec_path], dtype, &device)? }; + let semantic_codec = RepCodec::new(vb, &config.semantic_codec)?; + let s2mel_filters = mel_filter_bank( + config.s2mel.preprocess_params.spect_params.n_fft / 2 + 1, + config.s2mel.preprocess_params.spect_params.n_mels, + config.s2mel.preprocess_params.spect_params.fmin as f32, + config + .s2mel + .preprocess_params + .spect_params + .fmax + .unwrap_or(config.s2mel.preprocess_params.sr / 2) as f32, + config.s2mel.preprocess_params.sr as f32, + Some("slaney"), + crate::utils::audio_utils::MelScale::Slaney, + false, + device, + )? + .t()?; + let s2mel_windows = create_hann_window( + config.s2mel.preprocess_params.spect_params.win_length, + dtype, + device, + )?; + let (window_shift, window_size, padded_window_size) = + get_waveform_and_window_properties(16000, 10.0, 25.0, true)?; + let (mel_energies, _) = + kaldi_get_mel_banks(80, padded_window_size, 16000 as f32, 20.0, 0.0, device)?; + let mel_energies = mel_energies.pad_with_zeros(D::Minus1, 0, 1)?.t()?; + let campplus_model_path = save_dir.to_string() + + "/iic/speech_campplus_sv_zh-cn_16k-common/campplus_cn_common.bin"; + let campplus_vb = get_vb_model_path(campplus_model_path, dtype, device.clone(), None)?; + let campplus_model = CAMPPlus::new(campplus_vb, 80, 192, 32, 4, 128)?; + let s2mel = MyModel::new(path, &config.s2mel, dtype, device)?; + Ok(Self { + cache: None, + feature_extractor, + semantic_model, + semantic_mean, + semantic_std, + semantic_codec, + s2mel_filters, + s2mel_windows, + s2mel_preprocess_params: config.s2mel.preprocess_params.clone(), + window_shift, + window_size, + padded_window_size, + mel_energies, + campplus_model, + s2mel, + }) + } + + pub fn get_emb( + &self, + input_features: &Tensor, + attention_mask: Option<&Tensor>, + ) -> Result { + let output = + self.semantic_model + .forward(input_features, attention_mask, Some(17), false)?; + let feature = &output.specify_layer_id_hidden_state.unwrap(); + let feature = feature + .broadcast_sub(&self.semantic_mean)? + .broadcast_div(&self.semantic_std)?; + Ok(feature) + } + + pub fn s2mel_spectrogram(&self, waveform: &Tensor) -> Result { + let pad = (self.s2mel_preprocess_params.spect_params.n_fft + - self.s2mel_preprocess_params.spect_params.hop_length) + / 2; + let pad_audio_22k = pad_reflect_last_dim(&waveform, (pad, pad))?; + let spec = torch_stft( + &pad_audio_22k, + self.s2mel_preprocess_params.spect_params.n_fft, + self.s2mel_preprocess_params.spect_params.hop_length, + &self.s2mel_windows, + )? + .transpose(1, 2)?; + let spec = self.s2mel_filters.broadcast_matmul(&spec)?; + let spec = spec.clamp(1e-5, f64::INFINITY)?.log()?; + Ok(spec) + } + + pub fn forward( + &mut self, + input_ids: &Tensor, + audio_22k: Option<&Tensor>, + audio_16k: Option<&Tensor>, + ) -> Result<()> { + if (audio_22k.is_none() || audio_16k.is_none()) && self.cache.is_none() { + return Err(anyhow!( + "Missing required audio input: must provide either audio_22k, audio_16k, or have cached prompt data available" + )); + } + let (spk_cond_emb, style, prompt_condition, ref_mel) = if let Some(audio_22k) = audio_22k + && let Some(audio_16k) = audio_16k + { + let (audio_16k_features, audio_16k_mask) = + self.feature_extractor.call(&audio_16k, 16000, true, true)?; + let spk_cond_emb = self.get_emb(&audio_16k_features, audio_16k_mask.as_ref())?; + let (_, s_ref) = self.semantic_codec.quantize(&spk_cond_emb)?; + let ref_mel = self.s2mel_spectrogram(&audio_22k)?; + let ref_target_lengths = Tensor::new(ref_mel.dim(2)? as u32, ref_mel.device())?; + let feat = kaldi_fbank( + &audio_16k, + &self.mel_energies, + self.window_shift, + self.window_size, + self.padded_window_size, + 0.0, + )?; + let feat = feat.broadcast_sub(&feat.mean_keepdim(1)?)?; + let style = self.campplus_model.forward(&feat)?; + let prompt_condition = self + .s2mel + .length_regulator + .forward(&s_ref, &ref_target_lengths)?; + let cache = IndexTTS2Cache { + cache_spk_cond: spk_cond_emb.clone(), + cache_s2mel_style: style.clone(), + cache_s2mel_prompt: prompt_condition.clone(), + cache_mel: ref_mel.clone(), + }; + self.cache = Some(cache); + (spk_cond_emb, style, prompt_condition, ref_mel) + } else { + let cache = self.cache.as_ref().unwrap(); + ( + cache.cache_spk_cond.clone(), + cache.cache_s2mel_style.clone(), + cache.cache_s2mel_prompt.clone(), + cache.cache_mel.clone(), + ) + }; + + + Ok(()) + } } diff --git a/src/models/index_tts2/processor.rs b/src/models/index_tts2/processor.rs index 88c8e28..23839fc 100644 --- a/src/models/index_tts2/processor.rs +++ b/src/models/index_tts2/processor.rs @@ -5,13 +5,16 @@ use candle_nn::VarBuilder; use crate::{ models::{ - campplus::CAMPPlus, feature_extractor::seamless_m4t_feature_extractor::SeamlessM4TFeatureExtractor, index_tts2::config::{IndexTTS2Config, PreprocessParams}, mask_gct::model::RepCodec, w2v_bert_2_0::model::W2VBert2_0Model + campplus::CAMPPlus, + feature_extractor::seamless_m4t_feature_extractor::SeamlessM4TFeatureExtractor, + index_tts2::config::{IndexTTS2Config, PreprocessParams}, + mask_gct::model::RepCodec, + w2v_bert_2_0::model::W2VBert2_0Model, }, utils::{ audio_utils::{ create_hann_window, extract_audio_url, get_waveform_and_window_properties, kaldi_fbank, - kaldi_get_mel_banks, load_audio, mel_filter_bank, resample_simple, - torch_stft, + kaldi_get_mel_banks, load_audio, mel_filter_bank, resample_simple, torch_stft, }, get_vb_model_path, tensor_utils::pad_reflect_last_dim, @@ -21,104 +24,104 @@ use crate::{ pub struct IndexTTS2Processor { device: Device, max_audio_length_seconds: usize, - feature_extractor: SeamlessM4TFeatureExtractor, - semantic_model: W2VBert2_0Model, - semantic_mean: Tensor, - semantic_std: Tensor, - semantic_codec: RepCodec, - s2mel_filters: Tensor, - s2mel_windows: Tensor, - s2mel_preprocess_params: PreprocessParams, - window_shift: usize, - window_size: usize, - padded_window_size: usize, - mel_energies: Tensor, - campplus_model: CAMPPlus, + // feature_extractor: SeamlessM4TFeatureExtractor, + // semantic_model: W2VBert2_0Model, + // semantic_mean: Tensor, + // semantic_std: Tensor, + // semantic_codec: RepCodec, + // s2mel_filters: Tensor, + // s2mel_windows: Tensor, + // s2mel_preprocess_params: PreprocessParams, + // window_shift: usize, + // window_size: usize, + // padded_window_size: usize, + // mel_energies: Tensor, + // campplus_model: CAMPPlus, } impl IndexTTS2Processor { pub fn new( - path: &str, - save_dir: &str, - config: &IndexTTS2Config, + // path: &str, + // save_dir: &str, + // config: &IndexTTS2Config, device: &Device, - dtype: DType, + // dtype: DType, ) -> Result { - let feature_extractor = SeamlessM4TFeatureExtractor::new( - // 80, - 80, - crate::utils::tensor_utils::PaddingSide::Right, - 1.0, - 16000, - 2, - device, - )?; - let w2vbert2_path = save_dir.to_string() + "/facebook/w2v-bert-2.0"; - let semantic_model = W2VBert2_0Model::init(&w2vbert2_path, device, dtype)?; - let semantic_mean_var_path = path.to_string() + "/" + &config.w2v_stat; - let dict = read_all_with_key(semantic_mean_var_path, None)?; - let mut semantic_mean = Tensor::new(0.0, device)?.to_dtype(dtype)?; - let mut semantic_std = Tensor::new(1.0, device)?.to_dtype(dtype)?; - for (k, v) in dict { - if k.eq("mean") { - semantic_mean = v.to_device(device)?.to_dtype(dtype)?; - } else if k.eq("var") { - semantic_std = v.to_device(device)?.to_dtype(dtype)?.sqrt()?; - } - } + // let feature_extractor = SeamlessM4TFeatureExtractor::new( + // // 80, + // 80, + // crate::utils::tensor_utils::PaddingSide::Right, + // 1.0, + // 16000, + // 2, + // device, + // )?; + // let w2vbert2_path = save_dir.to_string() + "/facebook/w2v-bert-2.0"; + // let semantic_model = W2VBert2_0Model::init(&w2vbert2_path, device, dtype)?; + // let semantic_mean_var_path = path.to_string() + "/" + &config.w2v_stat; + // let dict = read_all_with_key(semantic_mean_var_path, None)?; + // let mut semantic_mean = Tensor::new(0.0, device)?.to_dtype(dtype)?; + // let mut semantic_std = Tensor::new(1.0, device)?.to_dtype(dtype)?; + // for (k, v) in dict { + // if k.eq("mean") { + // semantic_mean = v.to_device(device)?.to_dtype(dtype)?; + // } else if k.eq("var") { + // semantic_std = v.to_device(device)?.to_dtype(dtype)?.sqrt()?; + // } + // } - let semantic_codec_path = - save_dir.to_string() + "/amphion/MaskGCT/semantic_codec/model.safetensors"; - let vb = - unsafe { VarBuilder::from_mmaped_safetensors(&[semantic_codec_path], dtype, &device)? }; - let semantic_codec = RepCodec::new(vb, &config.semantic_codec)?; - let s2mel_filters = mel_filter_bank( - config.s2mel.preprocess_params.spect_params.n_fft / 2 + 1, - config.s2mel.preprocess_params.spect_params.n_mels, - config.s2mel.preprocess_params.spect_params.fmin as f32, - config - .s2mel - .preprocess_params - .spect_params - .fmax - .unwrap_or(config.s2mel.preprocess_params.sr / 2) as f32, - config.s2mel.preprocess_params.sr as f32, - Some("slaney"), - crate::utils::audio_utils::MelScale::Slaney, - false, - device, - )? - .t()?; - let s2mel_windows = create_hann_window( - config.s2mel.preprocess_params.spect_params.win_length, - dtype, - device, - )?; - let (window_shift, window_size, padded_window_size) = - get_waveform_and_window_properties(16000, 10.0, 25.0, true)?; - let (mel_energies, _) = - kaldi_get_mel_banks(80, padded_window_size, 16000 as f32, 20.0, 0.0, device)?; - let mel_energies = mel_energies.pad_with_zeros(D::Minus1, 0, 1)?.t()?; - let campplus_model_path = save_dir.to_string() - + "/iic/speech_campplus_sv_zh-cn_16k-common/campplus_cn_common.bin"; - let campplus_vb = get_vb_model_path(campplus_model_path, dtype, device.clone(), None)?; - let campplus_model = CAMPPlus::new(campplus_vb, 80, 192, 32, 4, 128)?; + // let semantic_codec_path = + // save_dir.to_string() + "/amphion/MaskGCT/semantic_codec/model.safetensors"; + // let vb = + // unsafe { VarBuilder::from_mmaped_safetensors(&[semantic_codec_path], dtype, &device)? }; + // let semantic_codec = RepCodec::new(vb, &config.semantic_codec)?; + // let s2mel_filters = mel_filter_bank( + // config.s2mel.preprocess_params.spect_params.n_fft / 2 + 1, + // config.s2mel.preprocess_params.spect_params.n_mels, + // config.s2mel.preprocess_params.spect_params.fmin as f32, + // config + // .s2mel + // .preprocess_params + // .spect_params + // .fmax + // .unwrap_or(config.s2mel.preprocess_params.sr / 2) as f32, + // config.s2mel.preprocess_params.sr as f32, + // Some("slaney"), + // crate::utils::audio_utils::MelScale::Slaney, + // false, + // device, + // )? + // .t()?; + // let s2mel_windows = create_hann_window( + // config.s2mel.preprocess_params.spect_params.win_length, + // dtype, + // device, + // )?; + // let (window_shift, window_size, padded_window_size) = + // get_waveform_and_window_properties(16000, 10.0, 25.0, true)?; + // let (mel_energies, _) = + // kaldi_get_mel_banks(80, padded_window_size, 16000 as f32, 20.0, 0.0, device)?; + // let mel_energies = mel_energies.pad_with_zeros(D::Minus1, 0, 1)?.t()?; + // let campplus_model_path = save_dir.to_string() + // + "/iic/speech_campplus_sv_zh-cn_16k-common/campplus_cn_common.bin"; + // let campplus_vb = get_vb_model_path(campplus_model_path, dtype, device.clone(), None)?; + // let campplus_model = CAMPPlus::new(campplus_vb, 80, 192, 32, 4, 128)?; Ok(Self { device: device.clone(), max_audio_length_seconds: 15, - feature_extractor, - semantic_model, - semantic_mean, - semantic_std, - semantic_codec, - s2mel_filters, - s2mel_windows, - s2mel_preprocess_params: config.s2mel.preprocess_params.clone(), - window_shift, - window_size, - padded_window_size, - mel_energies, - campplus_model, + // feature_extractor, + // semantic_model, + // semantic_mean, + // semantic_std, + // semantic_codec, + // s2mel_filters, + // s2mel_windows, + // s2mel_preprocess_params: config.s2mel.preprocess_params.clone(), + // window_shift, + // window_size, + // padded_window_size, + // mel_energies, + // campplus_model, }) } @@ -137,63 +140,64 @@ impl IndexTTS2Processor { &self, mes: &ChatCompletionParameters, device: &Device, - ) -> Result<(Tensor, usize)> { + ) -> Result<(Tensor, usize, String)> { let audio_url_vec = extract_audio_url(mes); let (audio, sr) = load_audio(&audio_url_vec[0], device)?; - self.cut_audio(&audio, sr) + let (audio, sr) = self.cut_audio(&audio, sr)?; + Ok((audio, sr, audio_url_vec[0].clone())) } - pub fn get_emb( - &self, - input_features: &Tensor, - attention_mask: Option<&Tensor>, - ) -> Result { - let output = - self.semantic_model - .forward(input_features, attention_mask, Some(17), false)?; - let feature = &output.specify_layer_id_hidden_state.unwrap(); - let feature = feature - .broadcast_sub(&self.semantic_mean)? - .broadcast_div(&self.semantic_std)?; - Ok(feature) - } + // pub fn get_emb( + // &self, + // input_features: &Tensor, + // attention_mask: Option<&Tensor>, + // ) -> Result { + // let output = + // self.semantic_model + // .forward(input_features, attention_mask, Some(17), false)?; + // let feature = &output.specify_layer_id_hidden_state.unwrap(); + // let feature = feature + // .broadcast_sub(&self.semantic_mean)? + // .broadcast_div(&self.semantic_std)?; + // Ok(feature) + // } - pub fn s2mel_spectrogram(&self, waveform: &Tensor) -> Result { - let pad = (self.s2mel_preprocess_params.spect_params.n_fft - - self.s2mel_preprocess_params.spect_params.hop_length) - / 2; - let pad_audio_22k = pad_reflect_last_dim(&waveform, (pad, pad))?; - let spec = torch_stft( - &pad_audio_22k, - self.s2mel_preprocess_params.spect_params.n_fft, - self.s2mel_preprocess_params.spect_params.hop_length, - &self.s2mel_windows, - )? - .transpose(1, 2)?; - let spec = self.s2mel_filters.broadcast_matmul(&spec)?; - let spec = spec.clamp(1e-5, f64::INFINITY)?.log()?; - Ok(spec) - } - pub fn process_info(&self, mes: &ChatCompletionParameters) -> Result<()> { - let (audio, sr) = self.extract_audio_and_cut(mes, &self.device)?; + // pub fn s2mel_spectrogram(&self, waveform: &Tensor) -> Result { + // let pad = (self.s2mel_preprocess_params.spect_params.n_fft + // - self.s2mel_preprocess_params.spect_params.hop_length) + // / 2; + // let pad_audio_22k = pad_reflect_last_dim(&waveform, (pad, pad))?; + // let spec = torch_stft( + // &pad_audio_22k, + // self.s2mel_preprocess_params.spect_params.n_fft, + // self.s2mel_preprocess_params.spect_params.hop_length, + // &self.s2mel_windows, + // )? + // .transpose(1, 2)?; + // let spec = self.s2mel_filters.broadcast_matmul(&spec)?; + // let spec = spec.clamp(1e-5, f64::INFINITY)?.log()?; + // Ok(spec) + // } + pub fn process_info(&self, mes: &ChatCompletionParameters) -> Result<(Tensor, Tensor, String)> { + let (audio, sr, audio_url) = self.extract_audio_and_cut(mes, &self.device)?; let audio_22k = resample_simple(&audio, sr as i64, 22050)?; let audio_16k = resample_simple(&audio, sr as i64, 16000)?; - let (audio_16k_features, audio_16k_mask) = - self.feature_extractor.call(&audio_16k, 16000, true, true)?; - let spk_cond_emb = self.get_emb(&audio_16k_features, audio_16k_mask.as_ref())?; - let (_, s_ref) = self.semantic_codec.quantize(&spk_cond_emb)?; - let ref_mel = self.s2mel_spectrogram(&audio_22k)?; - let feat = kaldi_fbank( - &audio_16k, - &self.mel_energies, - self.window_shift, - self.window_size, - self.padded_window_size, - 0.0, - )?; - let feat = feat.broadcast_sub(&feat.mean_keepdim(1)?)?; - let style = self.campplus_model.forward(&feat)?; - println!("style: {}", style); - Ok(()) + // let (audio_16k_features, audio_16k_mask) = + // self.feature_extractor.call(&audio_16k, 16000, true, true)?; + // let spk_cond_emb = self.get_emb(&audio_16k_features, audio_16k_mask.as_ref())?; + // let (_, s_ref) = self.semantic_codec.quantize(&spk_cond_emb)?; + // let ref_mel = self.s2mel_spectrogram(&audio_22k)?; + // let feat = kaldi_fbank( + // &audio_16k, + // &self.mel_energies, + // self.window_shift, + // self.window_size, + // self.padded_window_size, + // 0.0, + // )?; + // let feat = feat.broadcast_sub(&feat.mean_keepdim(1)?)?; + // let style = self.campplus_model.forward(&feat)?; + // println!("style: {}", style); + Ok((audio_22k, audio_16k, audio_url)) } } diff --git a/src/models/index_tts2/utils.rs b/src/models/index_tts2/utils.rs index e5c45fa..d80ecfb 100644 --- a/src/models/index_tts2/utils.rs +++ b/src/models/index_tts2/utils.rs @@ -1,18 +1,171 @@ +use std::collections::HashMap; + +use anyhow::{Result, anyhow}; +use regex::Regex; + use crate::utils::{download_model, get_default_save_dir}; -pub async fn download_index_tts2_need_model(save_dir: Option<&str>) -> anyhow::Result<()> { +pub async fn download_index_tts2_need_model(save_dir: Option<&str>) -> Result<()> { let save_dir = match save_dir { Some(dir) => dir.to_string(), None => get_default_save_dir().expect("Failed to get home directory"), }; let w2v_bert2_0 = "facebook/w2v-bert-2.0"; - let mask_gct= "amphion/MaskGCT"; + let mask_gct = "amphion/MaskGCT"; // let campplus= "funasr/campplus"; // huggingface - let campplus = "iic/speech_campplus_sv_zh-cn_16k-common"; // modelscope + let campplus = "iic/speech_campplus_sv_zh-cn_16k-common"; // modelscope download_model(w2v_bert2_0, &save_dir, 3).await?; download_model(mask_gct, &save_dir, 3).await?; download_model(campplus, &save_dir, 3).await?; - + Ok(()) -} \ No newline at end of file +} + +#[derive(Debug, Clone)] +pub struct TextNormalizer { + // char_rep_map: HashMap, + // zh_char_rep_map: HashMap, + pinyin_tone_pattern: Regex, + // name_pattern: Regex, + // tech_term_pattern: Regex, + // english_contraction_pattern: Regex, + email_pattern: Regex, + cjk_range_pattern: Regex, +} + +impl TextNormalizer { + pub fn new() -> Result { + let pinyin_tone_pattern: Regex = Regex::new( + // r"(?i)(? bool { + self.email_pattern.is_match(email) + } + + pub fn use_chinese(&self, s: &str) -> bool { + let has_chinese = s.chars().any(|c| ('\u{4e00}'..='\u{9fff}').contains(&c)); + let has_alpha = s.chars().any(|c| c.is_alphabetic()); + let is_email = self.match_email(s); + + if has_chinese || !has_alpha || is_email { + return true; + } + + self.pinyin_tone_pattern.is_match(s) + } + + pub fn tokenize_by_cjk_char(&self, line: &str, do_upper_case: bool) -> String { + // Split the line by CJK characters + let parts: Vec<&str> = self.cjk_range_pattern.split(line.trim()).collect(); + // Process each part and join with spaces + let mut result_parts = Vec::new(); + for part in parts { + if !part.trim().is_empty() { + if do_upper_case { + result_parts.push(part.trim().to_uppercase()); + } else { + result_parts.push(part.trim().to_string()); + } + } + } + // Join the parts with spaces + result_parts.join(" ") + } +} + +pub fn tokenize_by_cjk_char(line: &str, do_upper_case: bool) -> String { + let mut result_parts = Vec::new(); + for ch in line.chars() { + if ('\u{1100}'..='\u{11ff}').contains(&ch) + || ('\u{2e80}'..='\u{a4cf}').contains(&ch) + || ('\u{a840}'..='\u{d7af}').contains(&ch) + || ('\u{f900}'..='\u{faff}').contains(&ch) + || ('\u{fe30}'..='\u{fe4f}').contains(&ch) + || ('\u{ff65}'..='\u{ffdc}').contains(&ch) + || ('\u{20000}'..='\u{2ffff}').contains(&ch) + || ('\u{4e00}'..='\u{9fff}').contains(&ch) + { + // CJK 字符 + if do_upper_case { + result_parts.push(ch.to_uppercase().collect::()); + } else { + result_parts.push(ch.to_string()); + } + } else { + // 非 CJK 字符，保持在一起 + if do_upper_case { + result_parts.push(ch.to_uppercase().to_string()); + } else { + result_parts.push(ch.to_string()); + } + } + } + + result_parts.join(" ") +} diff --git a/src/models/mask_gct/model.rs b/src/models/mask_gct/model.rs index 7dd9c8f..a212d52 100644 --- a/src/models/mask_gct/model.rs +++ b/src/models/mask_gct/model.rs @@ -28,7 +28,7 @@ impl ConvNeXtBlock { // layer_scale_init_value: f32, ) -> Result { let dwconv = get_conv1d(vb.pp("dwconv"), dim, dim, 7, 3, 1, 1, dim, true)?; - let norm = get_layer_norm(vb.pp("norm"), 1e-6, dim)?; + let norm = get_layer_norm(vb.pp("norm"), 1e-6, dim, true)?; let pwconv1 = linear(dim, intermediate_dim, vb.pp("pwconv1"))?; let pwconv2 = linear(intermediate_dim, dim, vb.pp("pwconv2"))?; let gamma = vb.get_with_hints(dim, "gamma", Init::Const(1.0))?; @@ -73,14 +73,14 @@ impl VocosBackbone { // layer_scale_init_value: Option, ) -> Result { let embed = get_conv1d(vb.pp("embed"), input_channels, dim, 7, 3, 1, 1, 1, true)?; - let norm = get_layer_norm(vb.pp("norm"), 1e-6, dim)?; + let norm = get_layer_norm(vb.pp("norm"), 1e-6, dim, true)?; let vb_convnext = vb.pp("convnext"); let mut convnext = vec![]; for i in 0..num_layers { let layer = ConvNeXtBlock::new(vb_convnext.pp(i), dim, intermediate_dim)?; convnext.push(layer); } - let final_layer_norm = get_layer_norm(vb.pp("final_layer_norm"), 1e-6, dim)?; + let final_layer_norm = get_layer_norm(vb.pp("final_layer_norm"), 1e-6, dim, true)?; Ok(Self { embed, norm, diff --git a/src/models/paddleocr_vl/model.rs b/src/models/paddleocr_vl/model.rs index 4f68f68..50847e8 100644 --- a/src/models/paddleocr_vl/model.rs +++ b/src/models/paddleocr_vl/model.rs @@ -37,6 +37,7 @@ impl Projector { vb.pp("pre_norm"), config.rms_norm_eps, config.vision_config.hidden_size, + true, )?; let linear_1 = linear(hidden_size, hidden_size, vb.pp("linear_1"))?; let linear_2 = linear(hidden_size, config.hidden_size, vb.pp("linear_2"))?; @@ -278,6 +279,7 @@ impl SiglipVisionModel { vb.pp("post_layernorm"), config.layer_norm_eps, config.hidden_size, + true, )?; Ok(Self { embeddings, diff --git a/src/models/qwen3_asr/model.rs b/src/models/qwen3_asr/model.rs index 7aae11a..188c863 100644 --- a/src/models/qwen3_asr/model.rs +++ b/src/models/qwen3_asr/model.rs @@ -50,11 +50,12 @@ impl Qwen3ASRAudioEncoderLayer { Some("out_proj"), )?; let self_attn_layer_norm = - get_layer_norm(vb.pp("self_attn_layer_norm"), 1e-5, config.d_model)?; + get_layer_norm(vb.pp("self_attn_layer_norm"), 1e-5, config.d_model, true)?; let activation_fn = config.activation_function; let fc1 = linear(config.d_model, config.encoder_ffn_dim, vb.pp("fc1"))?; let fc2 = linear(config.encoder_ffn_dim, config.d_model, vb.pp("fc2"))?; - let final_layer_norm = get_layer_norm(vb.pp("final_layer_norm"), 1e-5, config.d_model)?; + let final_layer_norm = + get_layer_norm(vb.pp("final_layer_norm"), 1e-5, config.d_model, true)?; Ok(Self { self_attn, self_attn_layer_norm, @@ -105,7 +106,7 @@ impl Qwen3ASRAudioEncoder { let layer = Qwen3ASRAudioEncoderLayer::new(vb_layers.pp(i), config)?; layers.push(layer); } - let ln_post = get_layer_norm(vb.pp("ln_post"), 1e-5, config.d_model)?; + let ln_post = get_layer_norm(vb.pp("ln_post"), 1e-5, config.d_model, true)?; let conv2d1 = get_conv2d( vb.pp("conv2d1"), 1, diff --git a/src/models/qwen3vl/model.rs b/src/models/qwen3vl/model.rs index 6bea3d4..b3bd750 100644 --- a/src/models/qwen3vl/model.rs +++ b/src/models/qwen3vl/model.rs @@ -90,7 +90,7 @@ impl Qwen3VLVisionPatchMerger { } else { config.hidden_size }; - let norm = get_layer_norm(vb.pp("norm"), 1e-6, norm_size)?; + let norm = get_layer_norm(vb.pp("norm"), 1e-6, norm_size, true)?; let linear_fc1 = linear(hidden_size, hidden_size, vb.pp("linear_fc1"))?; let act_fn = Activation::Gelu; let linear_fc2 = linear(hidden_size, config.out_hidden_size, vb.pp("linear_fc2"))?; @@ -200,8 +200,8 @@ pub struct Qwen3VLVisionBlock { impl Qwen3VLVisionBlock { pub fn new(config: Qwen3VLVisionConfig, vb: VarBuilder) -> Result { - let norm1 = get_layer_norm(vb.pp("norm1"), 1e-6, config.hidden_size)?; - let norm2 = get_layer_norm(vb.pp("norm2"), 1e-6, config.hidden_size)?; + let norm1 = get_layer_norm(vb.pp("norm1"), 1e-6, config.hidden_size, true)?; + let norm2 = get_layer_norm(vb.pp("norm2"), 1e-6, config.hidden_size, true)?; let attn = Qwen3VLVisionAttention::new(config.clone(), vb.pp("attn"))?; let mlp = TwoLinearMLP::new( vb.pp("mlp"), diff --git a/src/models/rmbg2_0/model.rs b/src/models/rmbg2_0/model.rs index f47f8d0..b3af2d2 100644 --- a/src/models/rmbg2_0/model.rs +++ b/src/models/rmbg2_0/model.rs @@ -42,7 +42,7 @@ impl PatchEmbed { true, )?; let norm = if patch_norm { - Some(get_layer_norm(vb.pp("norm"), 1e-5, embed_dim)?) + Some(get_layer_norm(vb.pp("norm"), 1e-5, embed_dim, true)?) } else { None }; @@ -242,7 +242,7 @@ impl SwinTransformerBlock { window_size: usize, shift_size: usize, ) -> Result { - let norm1 = get_layer_norm(vb.pp("norm1"), 1e-5, dim)?; + let norm1 = get_layer_norm(vb.pp("norm1"), 1e-5, dim, true)?; let attn = WindowAttention::new( vb.pp("attn"), @@ -251,7 +251,7 @@ impl SwinTransformerBlock { qkv_bias, (window_size, window_size), )?; - let norm2 = get_layer_norm(vb.pp("norm2"), 1e-5, dim)?; + let norm2 = get_layer_norm(vb.pp("norm2"), 1e-5, dim, true)?; let mlp_dim = (dim as f32 * mlp_ratio) as usize; let mlp = TwoLinearMLP::new(vb.pp("mlp"), dim, mlp_dim, dim, act, true, "fc1", "fc2")?; Ok(Self { @@ -323,7 +323,7 @@ struct PatchMerging { impl PatchMerging { pub fn new(vb: VarBuilder, dim: usize) -> Result { let reduction = linear_no_bias(4 * dim, 2 * dim, vb.pp("reduction"))?; - let norm = get_layer_norm(vb.pp("norm"), 1e-5, 4 * dim)?; + let norm = get_layer_norm(vb.pp("norm"), 1e-5, 4 * dim, true)?; Ok(Self { reduction, norm }) } @@ -518,7 +518,7 @@ impl SwinTransformer { } let mut norms = vec![]; for i in out_indices.clone() { - let layer_i = get_layer_norm(vb.pp(format!("norm{i}")), 1e-5, num_features[i])?; + let layer_i = get_layer_norm(vb.pp(format!("norm{i}")), 1e-5, num_features[i], true)?; norms.push(layer_i); } Ok(Self { diff --git a/src/models/w2v_bert_2_0/model.rs b/src/models/w2v_bert_2_0/model.rs index fd4ddfd..e9fa3f9 100644 --- a/src/models/w2v_bert_2_0/model.rs +++ b/src/models/w2v_bert_2_0/model.rs @@ -7,9 +7,7 @@ use candle_nn::{ use crate::{ models::{ - common::{ - GLU, TwoLinearMLP, eager_attention_forward, get_conv1d, get_layer_norm, - }, + common::{GLU, TwoLinearMLP, eager_attention_forward, get_conv1d, get_layer_norm}, w2v_bert_2_0::config::W2VBert2_0Config, }, position_embed::rope::{RoPE, apply_rotary_pos_emb}, @@ -27,6 +25,7 @@ impl Wav2Vec2BertFeatureProjection { vb.pp("layer_norm"), config.layer_norm_eps, config.feature_projection_input_dim, + true, )?; let projection = linear( config.feature_projection_input_dim, @@ -248,6 +247,7 @@ impl Wav2Vec2BertConvolutionModule { vb.pp("layer_norm"), config.layer_norm_eps, config.hidden_size, + true, )?; let pointwise_conv1 = get_conv1d( vb.pp("pointwise_conv1"), @@ -277,6 +277,7 @@ impl Wav2Vec2BertConvolutionModule { vb.pp("depthwise_layer_norm"), config.layer_norm_eps, config.hidden_size, + true, )?; let pointwise_conv2 = get_conv1d( vb.pp("pointwise_conv2"), @@ -341,6 +342,7 @@ impl Wav2Vec2BertEncoderLayer { vb.pp("ffn1_layer_norm"), config.layer_norm_eps, config.hidden_size, + true, )?; let ffn1 = TwoLinearMLP::new( vb.pp("ffn1"), @@ -356,6 +358,7 @@ impl Wav2Vec2BertEncoderLayer { vb.pp("self_attn_layer_norm"), config.layer_norm_eps, config.hidden_size, + true, )?; let self_attn = Wav2Vec2BertSelfAttention::new(vb.pp("self_attn"), config, false)?; let conv_module = Wav2Vec2BertConvolutionModule::new(vb.pp("conv_module"), config)?; @@ -363,6 +366,7 @@ impl Wav2Vec2BertEncoderLayer { vb.pp("ffn2_layer_norm"), config.layer_norm_eps, config.hidden_size, + true, )?; let ffn2 = TwoLinearMLP::new( vb.pp("ffn2"), @@ -378,6 +382,7 @@ impl Wav2Vec2BertEncoderLayer { vb.pp("final_layer_norm"), config.layer_norm_eps, config.hidden_size, + true, )?; Ok(Self { ffn1_layer_norm, @@ -488,7 +493,7 @@ impl Wav2Vec2BertEncoder { for (i, layer) in (&self.layers).iter().enumerate() { if output_hidden_states { hidden_states.push(xs.clone()); - } + } if let Some(id) = layer_id && id == i { @@ -500,7 +505,7 @@ impl Wav2Vec2BertEncoder { sin.as_ref(), attention_mask.as_ref(), conv_attention_mask, - )?; + )?; } let hidden_states = if hidden_states.len() > 0 { Some(hidden_states) diff --git a/src/tokenizer/mod.rs b/src/tokenizer/mod.rs index 35f08bc..bcd0811 100644 --- a/src/tokenizer/mod.rs +++ b/src/tokenizer/mod.rs @@ -1,5 +1,6 @@ -use anyhow::{Ok, Result, anyhow}; +use anyhow::{Result, anyhow}; use candle_core::{Device, Tensor}; +use sentencepiece::SentencePieceProcessor; use serde_json::Value; use tokenizers::{ AddedToken, Tokenizer, decoders::byte_level::ByteLevel as ByteLevelDecoder, models::bpe::BPE, @@ -114,3 +115,16 @@ impl TokenizerModel { Ok(decode) } } + +pub fn sentencepiece_encode( + text: &str, + tokenizer: &SentencePieceProcessor, + device: &Device, +) -> Result { + let tokens = tokenizer + .encode(text) + .map_err(|e| anyhow!(format!("tokenizer encode error:{}", e)))?; + let token_ids = tokens.iter().map(|p| p.id).collect::>(); + let tokens_t = Tensor::new(token_ids, device)?.unsqueeze(0)?; + Ok(tokens_t) +} diff --git a/src/utils/mod.rs b/src/utils/mod.rs index ab64cee..e14b690 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -219,7 +219,7 @@ pub fn crate_tensor_from_reader( pub fn read_pth_tensor_info_cycle>( path: P, key: Option<&str>, -) -> Result> { +) -> Result> { let file = std::fs::File::open(path.as_ref())?; let zip_reader = std::io::BufReader::new(file); let mut zip = zip::ZipArchive::new(zip_reader)?; @@ -356,7 +356,11 @@ pub fn read_pth_tensor_info_cycle>( } }; } - Ok(tensors) + let mut dict_to_hashmap = HashMap::new(); + for (k, v) in tensors { + dict_to_hashmap.insert(k, v); + } + Ok(dict_to_hashmap) } pub fn round_by_factor(num: u32, factor: u32) -> u32 { diff --git a/src/utils/tensor_utils.rs b/src/utils/tensor_utils.rs index 025d82a..75e35cd 100644 --- a/src/utils/tensor_utils.rs +++ b/src/utils/tensor_utils.rs @@ -111,7 +111,7 @@ pub fn split_tensor_with_size( // "input tensor dim size % splits_size must be equal to 0" // ); for (i, split) in (0..dim_size).step_by(splits_size).enumerate() { - let size = splits_size.min(dim_size - i*splits_size); + let size = splits_size.min(dim_size - i * splits_size); split_res.push(t.narrow(dim, split, size)?); } Ok(split_res) @@ -1060,3 +1060,11 @@ pub fn float_range_normalize(t: &Tensor) -> Result { t = t.clamp(-1.0, 1.0)?; Ok(t) } + +pub fn sequence_mask(length: &Tensor, max_length: Option) -> Result { + let max_length = max_length.unwrap_or(length.max_all()?.to_scalar::()?); + let x = Tensor::arange(0, max_length, length.device())?.unsqueeze(0)?; + let length = length.unsqueeze(1)?; + let mask = x.broadcast_lt(&length)?; + Ok(mask) +} diff --git a/tests/messy_test.rs b/tests/messy_test.rs index d7cfe61..f9b0d9c 100644 --- a/tests/messy_test.rs +++ b/tests/messy_test.rs @@ -2,22 +2,33 @@ use std::time::Instant; -use aha::utils::{tensor_utils::interpolate_nearest_1d}; -use anyhow::Result; -use candle_core::{Tensor}; +use aha::utils::tensor_utils::interpolate_nearest_1d; +use anyhow::{Result, anyhow}; +use candle_core::Tensor; +use sentencepiece::SentencePieceProcessor; // use symphonia::core::io::MediaSourceStream; #[test] fn messy_test() -> Result<()> { - // RUST_BACKTRACE=1 cargo test -F cuda,ffmpeg messy_test -r -- --nocapture + // RUST_BACKTRACE=1 cargo test -F cuda messy_test -r -- --nocapture let device = &candle_core::Device::Cpu; - let t = Tensor::arange(0.0f32, 40.0, device)?.broadcast_as((1, 40, 40))?; - println!("t: {}", t); - let i_start = Instant::now(); - let t_inter = interpolate_nearest_1d(&t, 20)?; - let i_duration = i_start.elapsed(); - println!("Time elapsed in interpolate_nearest_1d is: {:?}", i_duration); - println!("t_inter: {}", t_inter); + let save_dir = + 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 bpe_path = model_path.to_string() + "/bpe.model"; + let tokenizer = SentencePieceProcessor::open(bpe_path) + .map_err(|e| anyhow!(format!("load bpe.model file error:{}", e)))?; + let tokens = tokenizer + .encode("你好啊") + .map_err(|e| anyhow!(format!("tokenizer encode error:{}", e)))?; + println!("tokens: {:?}", tokens); + // let t = Tensor::arange(0.0f32, 40.0, device)?.broadcast_as((1, 40, 40))?; + // println!("t: {}", t); + // let i_start = Instant::now(); + // let t_inter = interpolate_nearest_1d(&t, 20)?; + // let i_duration = i_start.elapsed(); + // println!("Time elapsed in interpolate_nearest_1d is: {:?}", i_duration); + // println!("t_inter: {}", t_inter); // let url = "https://sis-sample-audio.obs.cn-north-1.myhuaweicloud.com/16k16bit.mp3"; // let client = reqwest::blocking::Client::new(); // let response = client.get(url).send()?; diff --git a/tests/test_fun_asr_nano.rs b/tests/test_fun_asr_nano.rs index a24001e..552917c 100644 --- a/tests/test_fun_asr_nano.rs +++ b/tests/test_fun_asr_nano.rs @@ -21,7 +21,7 @@ fn fun_asr_nano_generate() -> Result<()> { "type": "audio", "audio_url": { - "url": "https://package-release.coderbox.cn/aiway/test/other/%E5%93%AA%E5%90%92.wav" + "url": "file://./assets/audio/zh.mp3" } }, { diff --git a/tests/weight_test.rs b/tests/weight_test.rs index f433397..0ab8894 100644 --- a/tests/weight_test.rs +++ b/tests/weight_test.rs @@ -200,6 +200,7 @@ fn qwen3_weight() -> Result<()> { #[test] 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);