tckm.yan/aha
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

update doc

Browse Source
This commit is contained in:
jhqxxx
2026-05-24 19:26:47 +08:00
parent 925e5dd7c0
commit 754145246c
8 changed files with 291 additions and 101 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/models/voxcpm_refact/generate.rs
+159 -9
View File
@@ -1,11 +1,13 @@
use anyhow::{Result, anyhow};
use base64::{Engine, prelude::BASE64_STANDARD};
use candle_core::{DType, Device, Tensor, pickle::read_all_with_key};
use candle_nn::VarBuilder;
use rocket::futures::Stream;
use rocket::futures::{Stream, stream};
use std::collections::HashMap;
use crate::{
models::{
GenerateModel,
voxcpm::{
audio_vae::AudioVAE,
config::{AudioVaeConfig, VoxCPMConfig},
@@ -13,7 +15,12 @@ use crate::{
},
voxcpm_refact::{model::VoxCPMModelRefact, processor::VoxCPMProcessor},
},
utils::{find_type_files, get_device, get_dtype},
params::chat::{ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse},
utils::{
audio_utils::{extract_audio_url, get_audio_wav_u8},
extract_metadata_value, extract_user_text, find_type_files, get_device, get_dtype,
response_utils::build_audio_completion_response,
},
};
pub struct VoxCPMGenerateRefact {
@@ -23,7 +30,7 @@ pub struct VoxCPMGenerateRefact {
processor: VoxCPMProcessor,
prompt_cache: Option<HashMap<String, Tensor>>,
out_sample_rate: usize,
// model_name: String,
model_name: String,
}
impl VoxCPMGenerateRefact {
@@ -55,11 +62,11 @@ impl VoxCPMGenerateRefact {
sr_bin_boundaries: None,
},
};
// let model_name = std::path::Path::new(path)
// .file_name()
// .and_then(|s| s.to_str())
// .unwrap_or("VoxCPM")
// .to_string();
let model_name = std::path::Path::new(path)
.file_name()
.and_then(|s| s.to_str())
.unwrap_or("VoxCPM")
.to_string();
let audio_vae = AudioVAE::new(
vb_vae,
audio_config.encoder_dim,
@@ -118,7 +125,7 @@ impl VoxCPMGenerateRefact {
processor,
prompt_cache: None,
out_sample_rate,
// model_name,
model_name,
})
}
@@ -126,6 +133,73 @@ impl VoxCPMGenerateRefact {
self.out_sample_rate
}
pub fn inference(
&mut self,
target_text: String,
prompt_text: Option<String>,
prompt_wav_path: Option<String>,
min_len: usize,
max_len: usize,
inference_timesteps: usize,
cfg_value: f64,
retry_badcase: bool,
retry_badcase_ratio_threshold: f64,
) -> Result<Tensor> {
let (text_token, audio_feat, audio_mask) = self.processor.processor(
target_text,
prompt_text,
prompt_wav_path,
&self.tokenizer,
&self.audio_vae,
)?;
let target_text_length = if let Some(mask) = &audio_mask {
text_token.dim(1)? - (mask.sum_all()?.to_scalar::<u32>()? as usize)
} else {
text_token.dim(1)?
};
let max_len = if retry_badcase {
(target_text_length as f64 * retry_badcase_ratio_threshold + 10.0) as usize
} else {
max_len
};
let audio = self.voxcpm.inference(
&text_token,
audio_feat.as_ref(),
audio_mask.as_ref(),
min_len,
max_len,
inference_timesteps,
cfg_value,
&self.audio_vae,
)?;
self.voxcpm.clear_kv_cache();
Ok(audio)
}
pub fn generate_with_prompt_simple(
&mut self,
target_text: String,
prompt_text: Option<String>,
prompt_wav_path: Option<String>,
) -> Result<Tensor> {
let audio = self.inference(
target_text,
prompt_text,
prompt_wav_path,
2,
1000,
10,
2.0,
false,
6.0,
)?;
Ok(audio)
}
pub fn generate_simple(&mut self, target_text: String) -> Result<Tensor> {
let audio = self.inference(target_text, None, None, 2, 100, 10, 2.0, false, 6.0)?;
Ok(audio)
}
pub fn build_prompt_cache(
&mut self,
prompt_text: String,
@@ -225,3 +299,79 @@ impl VoxCPMGenerateRefact {
}
}
}
impl GenerateModel for VoxCPMGenerateRefact {
fn generate(&mut self, mes: ChatCompletionParameters) -> Result<ChatCompletionResponse> {
let prompt_text = extract_metadata_value::<String>(&mes.metadata, "prompt_text");
let control_instruction =
extract_metadata_value::<String>(&mes.metadata, "control_instruction");
let min_len = extract_metadata_value::<usize>(&mes.metadata, "min_len").unwrap_or(2);
let max_len = extract_metadata_value::<usize>(&mes.metadata, "max_len").unwrap_or(4096);
let inference_timesteps =
extract_metadata_value::<usize>(&mes.metadata, "inference_timesteps").unwrap_or(10);
let cfg_value = extract_metadata_value::<f64>(&mes.metadata, "cfg_value").unwrap_or(2.0);
let retry_badcase_ratio_threshold =
extract_metadata_value::<f64>(&mes.metadata, "retry_badcase_ratio_threshold")
.unwrap_or(6.0);
let prompt_wav = extract_audio_url(&mes);
let prompt_wav_path = if !prompt_wav.is_empty() {
Some(prompt_wav[0].clone())
} else {
None
};
if !self.model_name.contains("2") && prompt_wav_path.is_some() && prompt_text.is_none() {
return Err(anyhow!(
"reference mode is only supported with VoxCPM2 models"
));
}
let mut target_text = extract_user_text(&mes)?;
if let Some(instruction) = control_instruction
&& self.model_name.contains("2")
{
target_text = format!("({instruction}){target_text}");
}
let audio = self
.inference(
target_text,
prompt_text,
prompt_wav_path,
min_len,
max_len,
inference_timesteps,
cfg_value,
true,
retry_badcase_ratio_threshold,
)
.inspect_err(|_| {
self.voxcpm.clear_kv_cache();
})?;
let wav_u8 = get_audio_wav_u8(&audio, self.out_sample_rate as u32)?;
// let wave_u8_str = String::from_utf8(wav_u8)?;
let base64_audio = BASE64_STANDARD.encode(wav_u8);
let response = build_audio_completion_response(&base64_audio, &self.model_name);
self.voxcpm.clear_kv_cache();
Ok(response)
}
#[allow(unused_variables)]
fn generate_stream(
&mut self,
mes: ChatCompletionParameters,
) -> Result<
Box<
dyn Stream<Item = Result<ChatCompletionChunkResponse, anyhow::Error>>
+ Send
+ Unpin
+ '_,
>,
> {
let error_stream = stream::once(async {
Err(anyhow::anyhow!(format!(
"{} model not support stream",
self.model_name
))) as Result<ChatCompletionChunkResponse, anyhow::Error>
});
Ok(Box::new(Box::pin(error_stream)))
}
}
src/models/voxcpm_refact/model.rs
-3
View File
@@ -196,9 +196,6 @@ impl VoxCPMModelRefact {
(text_embed, prefix_feat_cond, None)
};
let mut pred_feat_seq = Vec::new();
// if feat_mask.i((1, t-1))?.to_scalar::<f32>()? == 0.0 {
// // TODO for stream
// }
let mut position_id = 0;
let mut seq_len = t;
let enc_outputs = self