update voxcpm
This commit is contained in:
jhqxxx
Generated
-61
@@ -36,7 +36,6 @@ dependencies = [
|
||||
"openai_dive",
|
||||
"reqwest",
|
||||
"rocket",
|
||||
"rubato",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"tokenizers",
|
||||
@@ -2626,15 +2625,6 @@ dependencies = [
|
||||
"zerocopy",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "primal-check"
|
||||
version = "0.3.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "dc0d895b311e3af9902528fbb8f928688abbd95872819320517cc24ca6b2bd08"
|
||||
dependencies = [
|
||||
"num-integer",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "proc-macro-crate"
|
||||
version = "3.4.0"
|
||||
@@ -2918,15 +2908,6 @@ dependencies = [
|
||||
"crossbeam-utils",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "realfft"
|
||||
version = "3.5.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f821338fddb99d089116342c46e9f1fbf3828dba077674613e734e01d6ea8677"
|
||||
dependencies = [
|
||||
"rustfft",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "reborrow"
|
||||
version = "0.5.5"
|
||||
@@ -3153,18 +3134,6 @@ dependencies = [
|
||||
"uncased",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rubato"
|
||||
version = "0.16.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "5258099699851cfd0082aeb645feb9c084d9a5e1f1b8d5372086b989fc5e56a1"
|
||||
dependencies = [
|
||||
"num-complex",
|
||||
"num-integer",
|
||||
"num-traits",
|
||||
"realfft",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rustc-demangle"
|
||||
version = "0.1.26"
|
||||
@@ -3177,20 +3146,6 @@ version = "2.1.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "357703d41365b4b27c590e3ed91eabb1b663f07c4c084095e60cbed4362dff0d"
|
||||
|
||||
[[package]]
|
||||
name = "rustfft"
|
||||
version = "6.4.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "21db5f9893e91f41798c88680037dba611ca6674703c1a18601b01a72c8adb89"
|
||||
dependencies = [
|
||||
"num-complex",
|
||||
"num-integer",
|
||||
"num-traits",
|
||||
"primal-check",
|
||||
"strength_reduce",
|
||||
"transpose",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rustix"
|
||||
version = "1.1.2"
|
||||
@@ -3523,12 +3478,6 @@ version = "1.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f"
|
||||
|
||||
[[package]]
|
||||
name = "strength_reduce"
|
||||
version = "0.2.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "fe895eb47f22e2ddd4dabc02bce419d2e643c8e3b585c78158b349195bc24d82"
|
||||
|
||||
[[package]]
|
||||
name = "strsim"
|
||||
version = "0.11.1"
|
||||
@@ -4042,16 +3991,6 @@ dependencies = [
|
||||
"tracing-log",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "transpose"
|
||||
version = "0.2.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1ad61aed86bc3faea4300c7aee358b4c6d0c8d6ccc36524c96e4c92ccf26e77e"
|
||||
dependencies = [
|
||||
"num-integer",
|
||||
"strength_reduce",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "try-lock"
|
||||
version = "0.2.5"
|
||||
|
||||
@@ -27,7 +27,6 @@ chrono = "0.4.42"
|
||||
rocket = "0.5.1"
|
||||
tokio = "1.47.1"
|
||||
hound = "3.5.1"
|
||||
rubato = "0.16.2"
|
||||
|
||||
[features]
|
||||
flash-attn=["candle-flash-attn"]
|
||||
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -1,7 +1,7 @@
|
||||
use anyhow::{Error, Ok, Result};
|
||||
use candle_core::{D, IndexOp, Tensor};
|
||||
use anyhow::{Ok, Result};
|
||||
use candle_core::{D, Tensor};
|
||||
use candle_nn::{Conv1d, Conv1dConfig, ConvTranspose1d, ConvTranspose1dConfig, Module, VarBuilder};
|
||||
use std::result::Result::Ok as StdOk;
|
||||
use std::{result::Result::Ok as StdOk, thread, time};
|
||||
|
||||
pub struct CausalConv1d {
|
||||
conv1d: Conv1d,
|
||||
@@ -9,13 +9,9 @@ pub struct CausalConv1d {
|
||||
}
|
||||
|
||||
impl CausalConv1d {
|
||||
// CausalConv1d::new(scaled_weight, bias, padding, dilation, stride)?;
|
||||
pub fn new(
|
||||
weight: Tensor,
|
||||
bias: Option<Tensor>,
|
||||
// in_c: usize,
|
||||
// out_c: usize,
|
||||
// kernel_size: usize,
|
||||
padding: usize,
|
||||
dilation: usize,
|
||||
groups: usize,
|
||||
@@ -59,7 +55,7 @@ impl CausalConvTranspose1d {
|
||||
) -> Result<Self> {
|
||||
let config = ConvTranspose1dConfig {
|
||||
padding: 0,
|
||||
output_padding,
|
||||
output_padding: 0,
|
||||
stride,
|
||||
dilation,
|
||||
groups,
|
||||
@@ -74,23 +70,10 @@ impl CausalConvTranspose1d {
|
||||
})
|
||||
}
|
||||
pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
|
||||
println!("transpose conv input x: {:?}", x);
|
||||
println!("transpose conv config stride: {:?}", self.config.stride);
|
||||
println!("transpose conv config padding: {:?}", self.config.padding);
|
||||
println!("transpose conv config output_padding: {:?}", self.config.output_padding);
|
||||
println!("transpose conv config groups: {:?}", self.config.groups);
|
||||
println!("transpose conv config dilation: {:?}", self.config.dilation);
|
||||
println!("transpose conv config weight: {:?}", self.conv_transpose1d.weight());
|
||||
|
||||
let x = self.conv_transpose1d.forward(x)?;
|
||||
println!("transpose conv after x: {:?}", x);
|
||||
println!("transpose conv after self.padding: {:?}", self.padding);
|
||||
println!("transpose conv after self.output_padding: {:?}", self.output_padding);
|
||||
let last_dim = x.dim(D::Minus1)?;
|
||||
let select_num = last_dim - (self.padding * 2 - self.output_padding);
|
||||
println!("transpose conv after select_num: {:?}", select_num);
|
||||
let x = x.narrow(D::Minus1, 0, select_num)?;
|
||||
println!("transpose conv after x: {:?}", x);
|
||||
Ok(x)
|
||||
}
|
||||
}
|
||||
@@ -116,17 +99,14 @@ impl WNCausalConv1d {
|
||||
StdOk(b) => Some(b),
|
||||
Err(_) => None,
|
||||
};
|
||||
let weight_norm = weight_v.sqr()?.sum_keepdim(D::Minus1)?.sqrt()?;
|
||||
let weight_norm = weight_v.sqr()?.sum_keepdim(1)?.sum_keepdim(2)?.sqrt()?;
|
||||
let normalized_weight = weight_v.broadcast_div(&weight_norm)?;
|
||||
let scaled_weight = normalized_weight.broadcast_mul(&weight_g)?;
|
||||
let conv = CausalConv1d::new(scaled_weight, bias, padding, dilation, groups, stride)?;
|
||||
Ok(Self { conv })
|
||||
}
|
||||
pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
|
||||
println!("conv1d: x: {:?}", x);
|
||||
println!("conv weight: : {:?}", self.conv.conv1d.weight());
|
||||
let x = self.conv.forward(x)?;
|
||||
println!("conv1d: WN causal x: {:?}", x);
|
||||
Ok(x)
|
||||
}
|
||||
}
|
||||
@@ -154,7 +134,7 @@ impl WNCausalConvTranspose1d {
|
||||
StdOk(b) => Some(b),
|
||||
Err(_) => None,
|
||||
};
|
||||
let weight_norm = weight_v.sqr()?.sum_keepdim(D::Minus1)?.sqrt()?;
|
||||
let weight_norm = weight_v.sqr()?.sum_keepdim(1)?.sum_keepdim(2)?.sqrt()?;
|
||||
let normalized_weight = weight_v.broadcast_div(&weight_norm)?;
|
||||
let scaled_weight = normalized_weight.broadcast_mul(&weight_g)?;
|
||||
let conv_transpose = CausalConvTranspose1d::new(
|
||||
@@ -230,7 +210,6 @@ impl CausalResidualUnit {
|
||||
}
|
||||
|
||||
pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
|
||||
println!("causal residual unit x: {:?}", x);
|
||||
// let orig_dim = x.dims();
|
||||
let last_dim_x = x.dim(D::Minus1)?;
|
||||
let mut res_x = x.clone();
|
||||
@@ -238,11 +217,8 @@ impl CausalResidualUnit {
|
||||
let y = self.block1.forward(&y)?;
|
||||
let y = self.block2.forward(&y)?;
|
||||
let y = self.block3.forward(&y)?;
|
||||
println!("causal residual unit y: {:?}", y);
|
||||
// let dim = y.dims();
|
||||
let last_dim_y = y.dim(D::Minus1)?;
|
||||
println!("last_dim_x: {:?}", last_dim_x);
|
||||
println!("last_dim_y: {:?}", last_dim_y);
|
||||
let pad = (last_dim_x - last_dim_y) / 2;
|
||||
if pad > 0 {
|
||||
res_x = res_x.narrow(D::Minus1, pad, last_dim_y)?;
|
||||
@@ -415,17 +391,11 @@ impl CausalDecoderBlock {
|
||||
}
|
||||
|
||||
pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
|
||||
println!("decoder block x : {:?}", x);
|
||||
let x = self.block0.forward(x)?;
|
||||
println!("decoder block0 x : {:?}", x);
|
||||
let x = self.block1.forward(&x)?;
|
||||
println!("decoder block1 x : {:?}", x);
|
||||
let x = self.block2.forward(&x)?;
|
||||
println!("decoder block2 x : {:?}", x);
|
||||
let x = self.block3.forward(&x)?;
|
||||
println!("decoder block3 x : {:?}", x);
|
||||
let x = self.block4.forward(&x)?;
|
||||
println!("decoder block4 x : {:?}", x);
|
||||
Ok(x)
|
||||
}
|
||||
}
|
||||
@@ -456,7 +426,7 @@ impl CausalDecoder {
|
||||
input_channel,
|
||||
1,
|
||||
)?;
|
||||
let model1 = WNCausalConv1d::new(vb.pp("model.1"), input_channel, channels, 1, 1, 1, 1, 1)?;
|
||||
let model1 = WNCausalConv1d::new(vb.pp("model.1"), input_channel, channels, 1, 1, 0, 1, 1)?;
|
||||
let vb_model = vb.pp("model");
|
||||
let mut output_dim = channels;
|
||||
let mut model2_5 = Vec::new();
|
||||
@@ -485,9 +455,7 @@ impl CausalDecoder {
|
||||
}
|
||||
|
||||
pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
|
||||
print!("audio_vae decoder input x shape: {:?}", x);
|
||||
let x = self.model0.forward(x)?;
|
||||
print!("audio_vae decoder model0 x shape: {:?}", x);
|
||||
let mut x = self.model1.forward(&x)?;
|
||||
for model_i in &self.model2_5 {
|
||||
x = model_i.forward(&x)?;
|
||||
|
||||
+10
-48
@@ -72,10 +72,10 @@ impl SinusoidalPosEmb {
|
||||
.to_dtype(x.dtype())?;
|
||||
|
||||
let emb = x
|
||||
.unsqueeze(D::Minus1)?
|
||||
.unsqueeze(1)?
|
||||
.contiguous()?
|
||||
.matmul(&emb.unsqueeze(0)?.contiguous()?)?
|
||||
.affine(scale as f64, 0.0)?;
|
||||
.affine(scale as f64, 0.0)?
|
||||
.matmul(&emb.unsqueeze(0)?.contiguous()?)?;
|
||||
let emb = Tensor::cat(&[emb.sin()?, emb.cos()?], D::Minus1)?;
|
||||
Ok(emb)
|
||||
}
|
||||
@@ -167,7 +167,7 @@ impl VoxCPMLocDiT {
|
||||
let cond = self
|
||||
.cond_proj
|
||||
.forward(&cond.transpose(1, 2)?.contiguous()?)?;
|
||||
let prefix = cond.dims()[1];
|
||||
let prefix = cond.dim(1)?;
|
||||
let t = self.time_embeddings.forward(t, 1000)?.to_dtype(x.dtype())?;
|
||||
let t = self.time_mlp.forward(&t)?;
|
||||
let dt = self
|
||||
@@ -233,7 +233,6 @@ impl UnifiedCFM {
|
||||
let z = Tensor::randn(0.0f32, 1.0, (b, self.in_channels, t), mu.device())?
|
||||
.to_dtype(dtype)?
|
||||
.affine(temperature, 0.0)?;
|
||||
println!("z: {}", z);
|
||||
let t_span = linspace(1.0, 0.0, n_timesteps + 1, mu.device())?.to_dtype(dtype)?;
|
||||
let t_span = t_span
|
||||
.affine(f64::consts::PI / 2.0, 0.0)?
|
||||
@@ -242,11 +241,6 @@ impl UnifiedCFM {
|
||||
.add(&t_span)?
|
||||
.affine(sway_sampling_coef, 0.0)?
|
||||
.add(&t_span)?;
|
||||
println!("t_span: {}", t_span);
|
||||
println!("mu: {}", mu);
|
||||
println!("cond: {}", cond);
|
||||
println!("cfg_value: {}", cfg_value);
|
||||
println!("use_cfg_zero_star: {}", use_cfg_zero_star);
|
||||
let x = self.solve_euler(&z, &t_span, mu, cond, cfg_value, use_cfg_zero_star)?;
|
||||
Ok(x)
|
||||
}
|
||||
@@ -274,7 +268,7 @@ impl UnifiedCFM {
|
||||
let mut t = t_span.i(0)?;
|
||||
let mut dt = t.sub(&t_span.i(1)?)?;
|
||||
let mut sol = Vec::new();
|
||||
let t_span_len = t_span.dims1()?;
|
||||
let t_span_len = t_span.dim(0)?;
|
||||
let zero_init_steps = max(1, (t_span_len as f32 * 0.04) as usize);
|
||||
let mut dphi_dt = Tensor::zeros(1, t_span.dtype(), t_span.device())?;
|
||||
let mut x = x.clone();
|
||||
@@ -320,7 +314,8 @@ impl UnifiedCFM {
|
||||
dt = t.sub(&t_span.i(step + 1)?)?;
|
||||
}
|
||||
}
|
||||
Ok(sol[sol.len() - 1].clone())
|
||||
let ret = sol[sol.len() - 1].clone();
|
||||
Ok(ret)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -333,8 +328,6 @@ pub struct VoxCPMLocEnc {
|
||||
|
||||
impl VoxCPMLocEnc {
|
||||
pub fn new(vb: VarBuilder, config: VoxMiniCPM4Config, input_dim: usize) -> Result<Self> {
|
||||
// let special_token = Tensor::randn(0.0f32, 1.0, (1, 1, 1, config.hidden_size), vb.device())?
|
||||
// .to_dtype(vb.dtype())?;
|
||||
let special_token = vb.get((1, 1, 1, config.hidden_size), "special_token")?;
|
||||
let in_proj = linear(input_dim, config.hidden_size, vb.pp("in_proj"))?;
|
||||
assert_eq!(
|
||||
@@ -354,16 +347,12 @@ impl VoxCPMLocEnc {
|
||||
pub fn forward(&mut self, x: &Tensor) -> Result<Tensor> {
|
||||
let (b, t, p, d) = x.dims4()?;
|
||||
let x = self.in_proj.forward(x)?;
|
||||
println!("VoxCPMLocEnc: in_proj: {}", x);
|
||||
let special_tokens = self.special_token.expand((b, t, 1, self.hidden_size))?;
|
||||
let x = Tensor::cat(&[special_tokens, x], 2)?;
|
||||
println!("VoxCPMLocEnc: cat: {}", x);
|
||||
let (b, t, p, c) = x.dims4()?;
|
||||
let x = x.reshape((b * t, p, c))?;
|
||||
let outputs = self.encoder.forward(&x, 0, false)?;
|
||||
println!("VoxCPMLocEnc: encoder: {}", outputs);
|
||||
let cls_output = outputs.i((.., 0, ..))?;
|
||||
println!("VoxCPMLocEnc: cls_output: {}", cls_output);
|
||||
let cls_output = cls_output.reshape((b, t, c))?;
|
||||
Ok(cls_output)
|
||||
}
|
||||
@@ -537,10 +526,8 @@ impl VoxCPMModel {
|
||||
let audio_feat = audio_feat
|
||||
.reshape((self.audio_vae.latent_dim, (), self.patch_size))?
|
||||
.permute((1, 2, 0))?;
|
||||
let dim0 = audio_feat.dim(0)?;
|
||||
println!("audio_feat: {:?}", audio_feat);
|
||||
let dim0 = audio_feat.dim(0)? - 1;
|
||||
let audio_feat = audio_feat.i(..dim0)?;
|
||||
println!("audio_feat --: {:?}", audio_feat);
|
||||
let audio_length = audio_feat.dim(0)?;
|
||||
let text_pad_token = Tensor::zeros(audio_length, DType::U32, &self.device)?;
|
||||
let text_token = Tensor::cat(&[text_token, text_pad_token], D::Minus1)?;
|
||||
@@ -594,7 +581,6 @@ impl VoxCPMModel {
|
||||
.squeeze(1)?;
|
||||
let decode_audio_len = decode_audio.dim(D::Minus1)? - 640 - 640;
|
||||
let decode_audio = decode_audio.narrow(D::Minus1, 640, decode_audio_len)?;
|
||||
println!("decode_audio: {}", decode_audio);
|
||||
Ok(decode_audio)
|
||||
}
|
||||
|
||||
@@ -609,21 +595,15 @@ impl VoxCPMModel {
|
||||
inference_timesteps: usize,
|
||||
cfg_value: f64,
|
||||
) -> Result<Tensor> {
|
||||
println!("text: {}", text);
|
||||
println!("text_mask: {}", text_mask);
|
||||
println!("feat: {}", feat);
|
||||
println!("feat_mask: {}", feat_mask);
|
||||
let (b, t, p, d) = feat.dims4()?;
|
||||
let feat_embed = self.feat_encoder.forward(feat)?; // [b, t, h_feat]
|
||||
println!("feat_embed: {}", feat_embed);
|
||||
let feat_embed = self.enc_to_lm_proj.forward(&feat_embed)?;
|
||||
println!("feat_embed: {}", feat_embed);
|
||||
|
||||
let scale_emb = if self.config.lm_config.use_mup {
|
||||
self.config.lm_config.scale_emb
|
||||
} else {
|
||||
1.0
|
||||
};
|
||||
|
||||
let text_embed = self
|
||||
.base_lm
|
||||
.embed_tokens
|
||||
@@ -631,41 +611,32 @@ impl VoxCPMModel {
|
||||
.unwrap()
|
||||
.forward(text)?
|
||||
.affine(scale_emb as f64, 0.0)?;
|
||||
println!("text_embed: {}", text_embed);
|
||||
let combined_embed = text_mask
|
||||
.unsqueeze(D::Minus1)?
|
||||
.broadcast_mul(&text_embed)?
|
||||
.add(&feat_mask.unsqueeze(D::Minus1)?.broadcast_mul(&feat_embed)?)?;
|
||||
println!("combined_embed: {}", combined_embed);
|
||||
|
||||
let mut prefix_feat_cond = feat.i((.., t - 1, ..))?;
|
||||
let mut pred_feat_seq = Vec::new();
|
||||
let mut position_id = 0;
|
||||
let mut seq_len = t;
|
||||
let enc_outputs = self.base_lm.forward_step(&combined_embed, position_id)?;
|
||||
println!("base_lm enc_outputs: {}", enc_outputs);
|
||||
let enc_outputs = self
|
||||
.fsq_layer
|
||||
.forward(&enc_outputs)?
|
||||
.broadcast_mul(&feat_mask.unsqueeze(D::Minus1)?)?
|
||||
.add(&enc_outputs.broadcast_mul(&text_mask.unsqueeze(D::Minus1)?)?)?;
|
||||
println!("fsq_layer enc_outputs: {}", enc_outputs);
|
||||
|
||||
let mut lm_hidden = enc_outputs.i((.., t - 1, ..))?;
|
||||
println!("lm_hidden shape: {:?}", lm_hidden);
|
||||
|
||||
let input_embeds =
|
||||
enc_outputs.add(&feat_mask.unsqueeze(D::Minus1)?.broadcast_mul(&feat_embed)?)?;
|
||||
let residual_enc_outputs = self.residual_lm.forward_step(&input_embeds, position_id)?;
|
||||
println!("residual_lm residual_enc_outputs: {}", residual_enc_outputs);
|
||||
let mut residual_hidden = residual_enc_outputs.i((.., t - 1, ..))?;
|
||||
|
||||
for i in 0..max_len {
|
||||
let dit_hidden_1 = self.lm_to_dit_proj.forward(&lm_hidden)?; // [b, h_dit]
|
||||
println!("dit_hidden_1: {}", dit_hidden_1);
|
||||
let dit_hidden_2 = self.res_to_dit_proj.forward(&residual_hidden)?; // [b, h_dit]
|
||||
println!("dit_hidden_2: {}", dit_hidden_2);
|
||||
let dit_hidden = dit_hidden_1.add(&dit_hidden_2)?;
|
||||
println!("dit_hidden: {}", dit_hidden);
|
||||
let cond = prefix_feat_cond.transpose(1, 2)?.contiguous()?;
|
||||
|
||||
let pred_feat = self
|
||||
@@ -681,24 +652,18 @@ impl VoxCPMModel {
|
||||
true,
|
||||
)?
|
||||
.transpose(1, 2)?; // [b, p, d]
|
||||
println!("pred_feat: {}", pred_feat);
|
||||
|
||||
let curr_embed = self.feat_encoder.forward(&pred_feat.unsqueeze(1)?)?; // [b, 1, c]
|
||||
let curr_embed = self.enc_to_lm_proj.forward(&curr_embed)?;
|
||||
println!("curr_embed: {}", curr_embed);
|
||||
pred_feat_seq.push(pred_feat.unsqueeze(1)?);
|
||||
|
||||
prefix_feat_cond = pred_feat;
|
||||
println!("lm_hidden: {}", lm_hidden);
|
||||
let stop_flag = self.stop_proj.forward(&lm_hidden)?.silu()?;
|
||||
println!("stop_flag: {}", stop_flag);
|
||||
let stop_flag = self
|
||||
.stop_head
|
||||
.forward(&stop_flag)?
|
||||
.argmax(D::Minus1)?
|
||||
.i(0)?
|
||||
.to_scalar::<u32>()?;
|
||||
println!("i: {}, stop_flag: {}", i, stop_flag);
|
||||
if i > min_len && stop_flag == 1 {
|
||||
break;
|
||||
}
|
||||
@@ -716,15 +681,12 @@ impl VoxCPMModel {
|
||||
}
|
||||
let pred_seq = Tensor::cat(&pred_feat_seq, 1)?; // (b, t, p, d)
|
||||
let (b, t, p, d) = pred_seq.dims4()?;
|
||||
println!("pred_seq: {:?}", pred_seq);
|
||||
let feat_pred = pred_seq
|
||||
.permute((0, 3, 1, 2))?
|
||||
.reshape((b, d, ()))?
|
||||
.contiguous()?;
|
||||
println!("feat_pred: {:?}", feat_pred);
|
||||
self.base_lm.clear_kv_cache();
|
||||
self.residual_lm.clear_kv_cache();
|
||||
|
||||
Ok(feat_pred)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -45,7 +45,7 @@ impl SingleChineseTokenizer {
|
||||
.encode(text, false)
|
||||
.map_err(|e| anyhow!(format!("tokenizer encode error: {}", e)))?;
|
||||
let tokens = encode.get_tokens();
|
||||
println!("tokens: {:?}", tokens);
|
||||
// println!("tokens: {:?}", tokens);
|
||||
let mut split_character = Vec::new();
|
||||
for token in tokens {
|
||||
let clean_token = token.replace("▁", "to");
|
||||
@@ -56,7 +56,7 @@ impl SingleChineseTokenizer {
|
||||
split_character.push(token.clone());
|
||||
}
|
||||
}
|
||||
println!("split_character: {:?}", split_character);
|
||||
// println!("split_character: {:?}", split_character);
|
||||
let ids: Vec<u32> = split_character
|
||||
.iter()
|
||||
.filter_map(|c| self.tokenizer.token_to_id(c))
|
||||
|
||||
+62
-27
@@ -1,11 +1,9 @@
|
||||
use anyhow::{Result, anyhow};
|
||||
use candle_core::{D, DType, Device, Tensor};
|
||||
use candle_nn::{conv1d_no_bias, Conv1d, Conv1dConfig, Module};
|
||||
use candle_nn::{Conv1d, Conv1dConfig, Module, conv1d_no_bias};
|
||||
use hound::{SampleFormat, WavReader};
|
||||
use rocket::futures::future::ok;
|
||||
use rubato::{
|
||||
Resampler, SincFixedIn, SincInterpolationParameters, SincInterpolationType, WindowFunction,
|
||||
};
|
||||
|
||||
use std::f64::consts::PI;
|
||||
use std::path::Path;
|
||||
|
||||
@@ -97,7 +95,10 @@ pub fn get_sinc_resample_kernel(
|
||||
let sqrt_dims = sqrt_arg.dims();
|
||||
let sqrt_arg_vec = sqrt_arg.flatten_all()?.to_vec1::<f32>()?;
|
||||
|
||||
let window_val:Vec<f32> = sqrt_arg_vec.iter().map(|x| i0(beta_val * x) / i0_beta).collect();
|
||||
let window_val: Vec<f32> = sqrt_arg_vec
|
||||
.iter()
|
||||
.map(|x| i0(beta_val * x) / i0_beta)
|
||||
.collect();
|
||||
let window = Tensor::new(window_val, device)?.reshape(sqrt_dims)?;
|
||||
window
|
||||
}
|
||||
@@ -130,20 +131,21 @@ pub fn apply_sinc_resample_kernel(
|
||||
|
||||
// 获取波形形状
|
||||
let dims = waveform.dims();
|
||||
let waveform_flat = waveform.reshape(((), dims[dims.len()-1]))?;
|
||||
let waveform_flat = waveform.reshape(((), dims[dims.len() - 1]))?;
|
||||
|
||||
let (num_wavs, length) = waveform_flat.dims2()?;
|
||||
let padded_waveform = waveform.pad_with_zeros(D::Minus1, width as usize, (width+orig_freq) as usize)?;
|
||||
let padded_waveform =
|
||||
waveform.pad_with_zeros(D::Minus1, width as usize, (width + orig_freq) as usize)?;
|
||||
|
||||
// 添加通道维度 [batch_size, 1, padded_length]
|
||||
let waveform_3d = padded_waveform.unsqueeze(1)?;
|
||||
let config = Conv1dConfig {
|
||||
padding: 0,
|
||||
stride: orig_freq as usize,
|
||||
dilation: 1,
|
||||
groups: 1,
|
||||
cudnn_fwd_algo: None,
|
||||
};
|
||||
padding: 0,
|
||||
stride: orig_freq as usize,
|
||||
dilation: 1,
|
||||
groups: 1,
|
||||
cudnn_fwd_algo: None,
|
||||
};
|
||||
|
||||
let conv1d = Conv1d::new(kernel.clone(), None, config);
|
||||
// 执行卷积
|
||||
@@ -155,14 +157,13 @@ pub fn apply_sinc_resample_kernel(
|
||||
let conv_transposed = conv_output.transpose(1, 2)?.reshape((num_wavs, ()))?;
|
||||
|
||||
// 计算目标长度
|
||||
let target_length =
|
||||
((new_freq as f64 * length as f64) / orig_freq as f64).ceil() as usize;
|
||||
let target_length = ((new_freq as f64 * length as f64) / orig_freq as f64).ceil() as usize;
|
||||
|
||||
// 截取目标长度
|
||||
let resampled_flat =
|
||||
conv_transposed.narrow(1, 0, target_length.min(conv_transposed.dim(1)?))?;
|
||||
let mut new_dims = dims.to_vec();
|
||||
let last_dim = new_dims.len()-1;
|
||||
let last_dim = new_dims.len() - 1;
|
||||
new_dims[last_dim] = resampled_flat.dim(1)?;
|
||||
// 恢复原始批次形状
|
||||
|
||||
@@ -182,9 +183,7 @@ pub fn resample(
|
||||
beta: Option<f32>,
|
||||
) -> Result<Tensor> {
|
||||
if orig_freq <= 0 || new_freq <= 0 {
|
||||
return Err(anyhow!(
|
||||
"Frequencies must be positive".to_string(),
|
||||
));
|
||||
return Err(anyhow!("Frequencies must be positive".to_string(),));
|
||||
}
|
||||
|
||||
if orig_freq == new_freq {
|
||||
@@ -227,10 +226,26 @@ pub fn load_audio<P: AsRef<Path>>(path: P, device: Device) -> Result<(Tensor, us
|
||||
let samples: Vec<f32> = match spec.sample_format {
|
||||
SampleFormat::Int => {
|
||||
// 将整数样本转换为浮点数 [-1.0, 1.0]
|
||||
let max_value = match spec.bits_per_sample {
|
||||
8 => i8::MAX as f32,
|
||||
16 => i16::MAX as f32,
|
||||
24 => 8388607.0,
|
||||
println!("spec.bits_per_sample: {}", spec.bits_per_sample);
|
||||
let samples = match spec.bits_per_sample {
|
||||
8 => {
|
||||
reader
|
||||
.samples::<i8>()
|
||||
.map(|s| s.map(|sample| sample as f32 / i8::MAX as f32))
|
||||
.collect::<Result<Vec<_>, _>>()?
|
||||
},
|
||||
16 => {
|
||||
reader
|
||||
.samples::<i16>()
|
||||
.map(|s| s.map(|sample| sample as f32 / i16::MAX as f32))
|
||||
.collect::<Result<Vec<_>, _>>()?
|
||||
},
|
||||
24 => {
|
||||
reader
|
||||
.samples::<i32>()
|
||||
.map(|s| s.map(|sample| sample as f32 / 8388607.0))
|
||||
.collect::<Result<Vec<_>, _>>()?
|
||||
},
|
||||
_ => {
|
||||
return Err(anyhow::anyhow!(
|
||||
"Unsupported bit depth: {}",
|
||||
@@ -238,10 +253,7 @@ pub fn load_audio<P: AsRef<Path>>(path: P, device: Device) -> Result<(Tensor, us
|
||||
));
|
||||
}
|
||||
};
|
||||
reader
|
||||
.samples::<i16>()
|
||||
.map(|s| s.map(|sample| sample as f32 / max_value))
|
||||
.collect::<Result<Vec<_>, _>>()?
|
||||
samples
|
||||
}
|
||||
SampleFormat::Float => {
|
||||
// 直接读取浮点数样本
|
||||
@@ -258,6 +270,7 @@ pub fn load_audio<P: AsRef<Path>>(path: P, device: Device) -> Result<(Tensor, us
|
||||
&device,
|
||||
)?
|
||||
.t()?;
|
||||
// println!("audio channels: {}", spec.channels);
|
||||
if spec.channels > 1 {
|
||||
// 对channel通道求平均， channel维度变为1
|
||||
audio_tensor = audio_tensor.mean_keepdim(0)?;
|
||||
@@ -277,3 +290,25 @@ pub fn load_audio_with_resample<P: AsRef<Path>>(
|
||||
}
|
||||
Ok(audio)
|
||||
}
|
||||
|
||||
pub fn save_wav(audio: &Tensor, save_path: &str) -> Result<()> {
|
||||
let spec = hound::WavSpec {
|
||||
channels: 1,
|
||||
sample_rate: 16000,
|
||||
bits_per_sample: 16,
|
||||
sample_format: hound::SampleFormat::Int,
|
||||
};
|
||||
assert_eq!(audio.dim(0)?, 1, "audio channel must be 1");
|
||||
let max = audio.abs()?.max_all()?;
|
||||
let max = max.to_scalar::<f32>()?;
|
||||
let ratio = if max > 1.0 { 32767.0 / max } else { 32767.0 };
|
||||
let audio = audio.squeeze(0)?;
|
||||
let audio_vec = audio.to_vec1::<f32>()?;
|
||||
let mut writer = hound::WavWriter::create(save_path, spec).unwrap();
|
||||
for i in audio_vec {
|
||||
let sample_i16 = (i * ratio).round() as i16;
|
||||
writer.write_sample(sample_i16).unwrap();
|
||||
}
|
||||
writer.finalize().unwrap();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
+39
-7
@@ -1,13 +1,22 @@
|
||||
use std::collections::HashMap;
|
||||
use anyhow::{Ok, Result};
|
||||
use std::collections::HashMap;
|
||||
|
||||
use aha::{models::voxcpm::{audio_vae::AudioVAE, config::VoxCPMConfig, model::VoxCPMModel, tokenizer::SingleChineseTokenizer}, utils::utils::{find_type_files, get_device}};
|
||||
use aha::{
|
||||
models::voxcpm::{
|
||||
audio_vae::AudioVAE, config::VoxCPMConfig, model::VoxCPMModel,
|
||||
tokenizer::SingleChineseTokenizer,
|
||||
},
|
||||
utils::{
|
||||
audio_utils::save_wav,
|
||||
utils::{find_type_files, get_device},
|
||||
},
|
||||
};
|
||||
use candle_core::pickle::read_all_with_key;
|
||||
use candle_nn::VarBuilder;
|
||||
|
||||
|
||||
#[test]
|
||||
fn voxcpm_generate() -> Result<()> {
|
||||
// RUST_BACKTRACE=1 cargo test -F cuda,flash-attn voxcpm_generate -- --nocapture
|
||||
let model_path = "/home/jhq/huggingface_model/openbmb/VoxCPM-0.5B/";
|
||||
let model_list = find_type_files(&model_path, "pth")?;
|
||||
println!(" pth model_list: {:?}", model_list);
|
||||
@@ -19,11 +28,22 @@ fn voxcpm_generate() -> Result<()> {
|
||||
dtype = dict[0].1.dtype();
|
||||
for (k, v) in dict {
|
||||
// println!("key: {}, tensor shape: {:?}", k, v);
|
||||
// if k.contains("decoder.model.2.block.1") {
|
||||
// println!("val: {}", v);
|
||||
// }
|
||||
dict_to_hashmap.insert(k, v);
|
||||
}
|
||||
}
|
||||
let vb = VarBuilder::from_tensors(dict_to_hashmap, dtype, &dev);
|
||||
let audio_vae = AudioVAE::new(vb, 128, vec![2, 5, 8, 8], Some(64), 1536, vec![8, 8, 5, 2], 16000)?;
|
||||
let audio_vae = AudioVAE::new(
|
||||
vb,
|
||||
128,
|
||||
vec![2, 5, 8, 8],
|
||||
Some(64),
|
||||
1536,
|
||||
vec![8, 8, 5, 2],
|
||||
16000,
|
||||
)?;
|
||||
println!("audio vae load down");
|
||||
let model_list = find_type_files(&model_path, "bin")?;
|
||||
println!(" bin model_list: {:?}", model_list);
|
||||
@@ -41,9 +61,21 @@ fn voxcpm_generate() -> Result<()> {
|
||||
let config: VoxCPMConfig = serde_json::from_slice(&std::fs::read(config_path)?)?;
|
||||
let tokenizer = SingleChineseTokenizer::new(model_path)?;
|
||||
let mut voxcpm = VoxCPMModel::new(vb_vox, config, tokenizer, audio_vae)?;
|
||||
let generate = voxcpm.generate("你好啊，这是初始测试语句".to_string(), None, None, 2, 30, 10, 2.0, false, 3, 6.0)?;
|
||||
// let audio_path = "./assets/audio/example.wav";
|
||||
|
||||
let generate = voxcpm.generate(
|
||||
"太阳当空照，花儿对我笑，小鸟说早早早".to_string(),
|
||||
Some("啥子小师叔，打狗还要看主人，你再要继续，我，就是你的对手".to_string()),
|
||||
Some("./assets/audio/voice_01.wav".to_string()),
|
||||
// Some("一定被灰太狼给吃了，我已经为他准备好了花圈了".to_string()),
|
||||
// Some("./assets/audio/voice_05.wav".to_string()),
|
||||
2,
|
||||
100,
|
||||
10,
|
||||
2.0,
|
||||
false,
|
||||
3,
|
||||
6.0,
|
||||
)?;
|
||||
let _ = save_wav(&generate, "voxcpm_init.wav")?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
||||
Binary file not shown.
Reference in New Issue
Block a user