pass appropriate arguments to styletts inference command

This commit is contained in:
Kasisnu committed 2024-01-12 17:58:12 +05:30
1 parent c6a6676f08
commit ea3a6d3c0b
4 files changed
+45 -147

No files matched your search

@@ -1,11 +1,7 @@
use crate::util::model_downloader::ModelDownloader;
pub struct StyleTTS2Downloaders {
pub styletts2: StyleTTS2ModelMapping,
pub whisper_lg: WhisperLargeModelMapping,
pub whisper_md: WhisperMediumModelMapping,
pub vocos_encodec_24khz: VocosEncodecModelMapping,
pub vocos_config: VocosEncodecConfigMapping
// pub styletts2: StyleTTS2ModelMapping,
}
crate::impl_model_downloader!(
@@ -54,18 +50,15 @@ crate::impl_model_downloader!(
);
impl StyleTTS2Downloaders {
// TODO(KS): All checkpoints are in the container, so we don't need to download them here.
// Maybe they should be here instead
pub fn build_all_from_env() -> Self {
Self {
styletts2: StyleTTS2ModelMapping::from_env(),
whisper_lg: WhisperLargeModelMapping::from_env(),
whisper_md: WhisperMediumModelMapping::from_env(),
vocos_encodec_24khz: VocosEncodecModelMapping::from_env(),
vocos_config: VocosEncodecConfigMapping::from_env(),
}
}
pub fn all_downloaders(&self) -> Vec<&dyn ModelDownloader> {
vec![&self.styletts2, &self.whisper_lg, &self.whisper_md, &self.vocos_encodec_24khz, &self.vocos_config]
vec![]
}
}
@@ -113,20 +113,10 @@ pub async fn process_create_voice(
info!("Dataset length info: {}", dataset.len());
// let temp_extension = String::from(".bin");
// let temp_prefix: String;
//
// if !deps.job.info.container.is_on_prem {
// temp_prefix = String::from("sample_"); // this is for seed in local dev to download the samples
// } else {
// temp_prefix = String::from(BUCKET_FILE_PREFIX_CREATE);
// }
let mut downloaded_dataset: Vec<PathBuf> = Vec::new();
for (index, record) in dataset.iter().enumerate() {
//https://storage.googleapis.com/dev-vocodes-public/media/5/3/3/w/8/533w8zs0fy11nv7gkcna7p7vt03h8nda/dev_zs_533w8zs0fy11nv7gkcna7p7vt03h8nda.bin <-- where the file actually is
let prefix: Option<&str> = record.maybe_public_bucket_prefix.as_ref().map(|s| s.as_str());
let extension: Option<&str> = record.maybe_public_bucket_extension.as_ref().map(|s| s.as_str());
@@ -166,9 +156,16 @@ pub async fn process_create_voice(
info!("FilePath to clone voice: {}", file_path.to_string_lossy());
downloaded_dataset.push(file_path.clone());
// TODO(KS): This is a hack to get only the first file to download
// Ideally we could have improved logic to pick the "best" media file from a dataset
if (index > 0) {
break;
}
}
info!("Dataset Length {}", downloaded_dataset.len());
info!("Dataset Length downloaded: {}", downloaded_dataset.len());
// STEP 4 Download the models
info!("Download models (if not present)...");
@@ -193,9 +190,7 @@ pub async fn process_create_voice(
info!("Files to process: {:?}", audio_files);
// Name of the output file
// NB: don't use the extension... for the inference since the container will add the extension.
let output_file_name = PathBuf::from("temp");
let output_embedding_path = work_temp_dir.path().join("style.npz");
let stderr_output_file = work_temp_dir.path().join("zero_shot_create_voice_err.txt");
let inference_start_time = Instant::now();
@@ -204,9 +199,8 @@ pub async fn process_create_voice(
let command_exit_status =
model_dependencies.create_embedding_command.execute_inference(
job::job_types::tts::styletts2::styletts2_inference_command::CreateVoiceInferenceArgs {
output_embedding_path: &workdir,
output_embedding_name: &output_file_name,
audio_files,
output_embedding_path: &output_embedding_path,
audio_file: audio_files,
stderr_output_file: &stderr_output_file,
}
);
@@ -220,15 +214,6 @@ pub async fn process_create_voice(
if let Ok(contents) = read_to_string(&stderr_output_file) {
warn!("Captured stderr output: {}", contents);
// Re-categorize error?
//match categorize_error(&contents) {
// Some(ProcessSingleJobError::FaceDetectionFailure) => {
// warn!("Face not detected in source image");
// error = ProcessSingleJobError::FaceDetectionFailure;
// }
// _ => {}
//}
}
//thread::sleep(Duration::from_secs(300));
@@ -246,13 +231,7 @@ pub async fn process_create_voice(
info!("Inference success!");
//info!("Success; waiting...");
//thread::sleep(Duration::from_secs(300));
// STEP 4. Download dataset each audio file
info!("Uploading Media ...");
let embedding_bucket_location = ZeroShotVoiceEmbeddingBucketPath::generate_new(
ModelCategory::Tts,
ModelType::StyleTTS2,
@@ -264,7 +243,7 @@ pub async fn process_create_voice(
// Get Finished File
let mut finished_file = work_temp_dir.path().to_path_buf();
let output_bucket_file_name = String::from("temp.npz"); // use extension for bucket upload.
let output_bucket_file_name = String::from("style.npz"); // use extension for bucket upload.
finished_file.push(&output_bucket_file_name);
info!("Upload File Path: {:?}", finished_file);
@@ -16,6 +16,7 @@ use hashing::sha256::sha256_hash_file::sha256_hash_file;
use mysql_queries::queries::media_files::create::insert_media_file_from_zero_shot_tts::insert_media_file_from_zero_shot;
use mysql_queries::queries::media_files::create::insert_media_file_from_zero_shot_tts::InsertArgs;
use mysql_queries::queries::voice_designer::voices::get_voice::{get_voice_by_token, ZsVoice};
use tts_common::clean_symbols::clean_symbols;
use crate::job::job_loop::job_success_result::JobSuccessResult;
use crate::job::job_loop::job_success_result::ResultEntity;
@@ -52,7 +53,9 @@ pub async fn process_inference_voice(
// get job args
let text = match job.maybe_raw_inference_text.clone() {
Some(value) => { value }
Some(value) => {
clean_symbols(value.as_str())
}
None => {
return Err(ProcessSingleJobError::InvalidJob(anyhow!("Missing Text for Inference")));
}
@@ -90,17 +93,6 @@ pub async fn process_inference_voice(
}
}
// Might not need this for inference.
// let creator_user_token:UserToken;
// match &job.maybe_creator_user_token {
// Some(token) => {
// creator_user_token = UserToken::new_from_str(token);
// },
// None => {
// return Err(ProcessSingleJobError::InvalidJob(anyhow!("Missing Creator User Token")));
// }
// }
// run inference
let work_temp_dir = format!("/tmp/temp_zeroshot_inference_{}", job.id.0);
@@ -110,7 +102,7 @@ pub async fn process_inference_voice(
.map_err(|e| ProcessSingleJobError::from_io_error(e))?;
let workdir = work_temp_dir.path().to_path_buf();
let filename = "weights.npz".to_string();
let filename = "style.npz".to_string();
let mut downloaded_weights_path = work_temp_dir.path().to_path_buf();
downloaded_weights_path.push(&filename);
@@ -129,20 +121,23 @@ pub async fn process_inference_voice(
.log_status("running inference")
.map_err(|e| ProcessSingleJobError::Other(e))?;
let text_input_fs_path = work_temp_dir.path().join("inference_input.txt");
std::fs::write(&text_input_fs_path, &text)
.map_err(|e| ProcessSingleJobError::from_io_error(e))?;
let inference_start_time = Instant::now();
let output_file_name = String::from("output.wav");
let output_file_name = work_temp_dir.path().join("output.wav");
let stderr_output_file = work_temp_dir.path().join("zero_shot_inference.txt");
// Run Inference
let command_exit_status =
model_dependencies.inference_command.execute_inference(
InferenceArgs {
input_embedding_path: &workdir,
input_embedding_name: filename,
input_text: String::from(text), // text
output_file_name: output_file_name.clone(), // output file name in the output folder
input_embedding_file_path: &workdir,
input_text_file_path: &text_input_fs_path, // text
output_file_name: &output_file_name, // output file name in the output folder
stderr_output_file: &stderr_output_file,
}
);
@@ -156,19 +151,8 @@ pub async fn process_inference_voice(
if let Ok(contents) = read_to_string(&stderr_output_file) {
warn!("Captured stderr output: {}", contents);
// Re-categorize error?
//match categorize_error(&contents) {
// Some(ProcessSingleJobError::FaceDetectionFailure) => {
// warn!("Face not detected in source image");
// error = ProcessSingleJobError::FaceDetectionFailure;
// }
// _ => {}
//}
}
//thread::sleep(Duration::from_secs(300));
// Clean up temp files
//safe_delete_temp_file(&audio_path.filesystem_path);
//safe_delete_temp_file(&image_path.filesystem_path);
@@ -39,17 +39,15 @@ static IGNORED_ENVIRONMENT_VARS: Lazy<HashSet<String>> = Lazy::new(|| {
// so we have all the files downloaded for the inputs.
pub struct InferenceArgs<P: AsRef<Path>> {
/// --driven_audio: path to the input audio
pub input_embedding_path: P, // name of the embedding.npz in the work dir
pub input_embedding_name: String,
pub input_text: String, // name of the text
pub output_file_name: String, // output file name in the output folder
pub input_embedding_file_path: P, // name of the embedding.npz in the work dir
pub input_text_file_path: P, // name of the text
pub output_file_name: P, // output file name in the output folder
pub stderr_output_file: P,
}
pub struct CreateVoiceInferenceArgs<P: AsRef<Path>> {
pub output_embedding_path: P,
pub output_embedding_name: P,
pub audio_files: String,
pub audio_file: String,
pub stderr_output_file: P,
}
@@ -195,46 +193,19 @@ impl StyleTTS2InferenceCommand {
// ===== Begin Python Args =====
command.push_str(" --text ");
command.push_str(&format!("\"{}\"", &args.input_text));
command.push_str(" --input ");
command.push_str(&path_to_string(args.input_text_file_path));
command.push_str(" --prompt-name ");
command.push_str(&path_to_string(args.input_embedding_name));
command.push_str(" --input-style-npz ");
command.push_str(&path_to_string(args.input_embedding_file_path));
command.push_str(" --prompt-path ");
command.push_str(&path_to_string(&args.input_embedding_path));
command.push_str(" --audio-name ");
command.push_str(" --output ");
command.push_str(&path_to_string(args.output_file_name));
command.push_str(" --audio-path ");
command.push_str(&path_to_string(&args.input_embedding_path));
command.push_str(" --mode ");
command.push_str(&path_to_string("0"));
// TODO improve and use a better model for premium users
command.push_str(" --whisper-model ");
command.push_str(&path_to_string(Path::new("medium")));
command.push_str(" --whisper-folder-path ");
command.push_str(
&path_to_string(Path::new("/tmp/downloads/zero_shot_tts/vall-e-x_1.0/whisper/"))
);
command.push_str(" --vocos-folder-path ");
command.push_str(
&path_to_string(Path::new("/tmp/downloads/zero_shot_tts/vall-e-x_1.0/vocos-encodec/"))
);
command.push_str(" --StyleTTS2-path ");
command.push_str(
&path_to_string(
Path::new("/tmp/downloads/zero_shot_tts/vall-e-x_1.0/StyleTTS2-checkpoint.pt")
)
);
command.push_str(" --tmp-work-dir ");
command.push_str(&path_to_string(Path::new("/tmp/downloads/zero_shot_tts/")));
//TODO(KS): Make this a parameter
command.push_str(" --vcsteps ");
command.push_str(format!("{}", 20).as_str());
// ===== End Python Args =====
@@ -446,40 +417,11 @@ impl StyleTTS2CreateEmbeddingCommand {
// ===== Begin Python Args =====
command.push_str(" --prompt-name ");
command.push_str(&path_to_string(args.output_embedding_name));
command.push_str(" --prompt-path ");
command.push_str(" --output-style-npz ");
command.push_str(&path_to_string(&args.output_embedding_path));
command.push_str(" --audio-wav-files ");
command.push_str(&args.audio_files);
command.push_str(" --mode ");
command.push_str(&path_to_string("1"));
// TODO improve and use a better model for premium users
command.push_str(" --whisper-model ");
command.push_str(&path_to_string(Path::new("medium")));
command.push_str(" --whisper-folder-path ");
command.push_str(
&path_to_string(Path::new("/tmp/downloads/zero_shot_tts/vall-e-x_1.0/whisper/"))
);
command.push_str(" --vocos-folder-path ");
command.push_str(
&path_to_string(Path::new("/tmp/downloads/zero_shot_tts/vall-e-x_1.0/vocos-encodec/"))
);
command.push_str(" --StyleTTS2-path ");
command.push_str(
&path_to_string(
Path::new("/tmp/downloads/zero_shot_tts/vall-e-x_1.0/StyleTTS2-checkpoint.pt")
)
);
command.push_str(" --tmp-work-dir ");
command.push_str(&path_to_string(Path::new("/tmp/downloads/zero_shot_tts/")));
command.push_str(" --voice ");
command.push_str(&args.audio_file);
// ===== End Python Args =====