From 68048659d4a5ebb865b38c0debc8d6752b515cb5 Mon Sep 17 00:00:00 2001 From: Michael Chung Date: Tue, 17 Oct 2023 22:00:14 -0400 Subject: [PATCH] Testing inference_job --- .../tts/vall_e_x/model_downloaders.rs | 21 +++- .../job/job_types/tts/vall_e_x/process_job.rs | 102 +++++++++--------- .../vall_e_x/vall_e_x_inference_command.rs | 26 +++-- .../service/web/storyteller_web/src/routes.rs | 2 +- 4 files changed, 89 insertions(+), 62 deletions(-) diff --git a/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/model_downloaders.rs b/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/model_downloaders.rs index a18d817b98..007a3282d3 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/model_downloaders.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/model_downloaders.rs @@ -5,6 +5,7 @@ pub struct VallEXDownloaders { pub whisper_lg: WhisperLargeModelMapping, pub whisper_md: WhisperMediumModelMapping, pub vocos_encodec_24khz: VocosEncodecModelMapping, + pub vocos_config: VocosEncodecConfigMapping } crate::impl_model_downloader!( @@ -29,18 +30,27 @@ crate::impl_model_downloader!( WhisperMediumModelMapping, "whisper Medium Model Mapping", "WHISPER_MEDIUM_TENSOR_BUCKET_PATH", - "/dependencies/zero_shot_tts/vall-e-x_1.0/whisper-medium/model.safetensors", + "/dependencies/zero_shot_tts/vall-e-x_1.0/whisper/medium.pt", "WHISPER_MEDIUM_TENSOR_FILESYSTEM_PATH", - "/tmp/downloads/zero_shot_tts/vall-e-x_1.0/whisper-medium/model.safetensors" + "/tmp/downloads/zero_shot_tts/vall-e-x_1.0/whisper/medium.pt" ); crate::impl_model_downloader!( VocosEncodecModelMapping, "Vocos Encodec Model Mapping", "VOCOS_ENCODEC_TENSOR_BUCKET_PATH", - "/dependencies/zero_shot_tts/vall-e-x_1.0/encodec_pytorch_model.bin", + "/dependencies/zero_shot_tts/vall-e-x_1.0/vocos-encodec/encodec_pytorch_model.bin", "VOCOS_ENCODEC_TENSOR_FILESYSTEM_PATH", - "/tmp/downloads/zero_shot_tts/vall-e-x_1.0/encodec_pytorch_model.bin" + "/tmp/downloads/zero_shot_tts/vall-e-x_1.0/vocos-encodec/encodec_pytorch_model.bin" +); + +crate::impl_model_downloader!( + VocosEncodecConfigMapping, + "Vocos Config Mapping", + "VOCOS_ENCODEC_CONFIG_BUCKET_PATH", + "/dependencies/zero_shot_tts/vall-e-x_1.0/vocos-encodec/config.yaml", + "VOCOS_ENCODEC_CONFIG_FILESYSTEM_PATH", + "/tmp/downloads/zero_shot_tts/vall-e-x_1.0/vocos-encodec/config.yaml" ); impl VallEXDownloaders { @@ -50,11 +60,12 @@ impl VallEXDownloaders { 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.vall_e_x, &self.whisper_lg, &self.whisper_md, &self.vocos_encodec_24khz] + vec![&self.vall_e_x, &self.whisper_lg, &self.whisper_md, &self.vocos_encodec_24khz,&self.vocos_config] } } diff --git a/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/process_job.rs b/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/process_job.rs index 4bbb0819e0..a70ad73751 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/process_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/process_job.rs @@ -21,17 +21,17 @@ pub struct VoiceFile { pub async fn download_voice_embedding_from_hash( bucket_hash: &str, name: &str, - private_bucket_client: BucketClient + private_bucket_client: &BucketClient, + path: &str ) -> Result { - let unifer = BucketPathUnifier::default_paths(); let object_path = unifer.zero_shot_tts_speaker_encoding(bucket_hash, 0); - let filesystem_path = format!("{}_weights.npz", name); - + let filesystem_path = format!("{}/{}_weights.npz", path, name); + let result = private_bucket_client.download_file_to_disk(object_path, &filesystem_path).await; let voice_file = VoiceFile { - filesystem_path: PathBuf::from(&filesystem_path), + filesystem_path: PathBuf::from(&filesystem_path), }; Ok(voice_file) @@ -50,22 +50,24 @@ pub async fn process_job( let job = args.job; let deps = args.job_dependencies; let mysql_pool = &deps.mysql_pool; - + // get some globals - let job_progress_reporter = deps.job_progress_reporter + let mut job_progress_reporter = deps.job_progress_reporter .new_generic_inference(job.inference_job_token.as_str()) .map_err(|e| ProcessSingleJobError::Other(anyhow!(e)))?; // get job args let text = match job.maybe_raw_inference_text.clone() { - Some(value) => { value }, - None => { return Err(ProcessSingleJobError::InvalidJob(anyhow!("Missing Text for Inference"))) } + Some(value) => { value } + None => { + return Err(ProcessSingleJobError::InvalidJob(anyhow!("Missing Text for Inference"))); + } }; // get args token let jobArgs = validate_job(&job)?; // bubbles error up let voice_token = tokens::tokens::zs_voices::ZsVoiceToken(jobArgs.voice_token); - + // Get voice bucket hash - from voice token let voice_lookup_result = get_voice_by_token(&voice_token, false, &mysql_pool).await; @@ -73,7 +75,7 @@ pub async fn process_job( Ok(Some(voice)) => voice, Ok(None) => { warn!("Voice not found: {:?}", voice_token); - return Err(ProcessSingleJobError::Other(anyhow!("Voice not found: {:?}",voice_token))); + return Err(ProcessSingleJobError::Other(anyhow!("Voice not found: {:?}", voice_token))); } Err(err) => { warn!("Error looking up voice: {:?}", err); @@ -81,10 +83,6 @@ pub async fn process_job( } }; - //download_voice_embedding_from_hash( bucket_hash, name, private_bucket_client); - - - // Need to download the models info!("Download models (if not present)..."); @@ -100,57 +98,61 @@ pub async fn process_job( } } - // Download embeddings file using embedding token - // Create a temp dir to download things to - - // Run Inference - let command_exit_status = - args.job_dependencies.job_type_details.vall_e_x.inference_command.execute_inference( - InferenceArgs { - input_embedding: String::from(""), // name of the embedding.npz in the tmp dir - input_text: String::from(text), // name of the text - /// --result_file: path to final file output - output_file_name: String::from("output.wav"), // output file name in the output folder - stderr_output_file: String::from(""), - } - ); - - // upload audio to bucket - - // deletetemp dir files - - // save result ? - - // ==================== UPLOAD AUDIO TO BUCKET ==================== // - // Download embeddings file using embedding token - // run inference let work_temp_dir = format!("temp_zeroshot_inference_{}", job.id.0); + let work_temp_dir_string = work_temp_dir.clone(); // NB: TempDir exists until it goes out of scope, at which point it should delete from filesystem. let work_temp_dir = args.job_dependencies.fs.scoped_temp_dir_creator_for_work .new_tempdir(&work_temp_dir) .map_err(|e| ProcessSingleJobError::from_io_error(e))?; - // delete temp dir files + let voiceFile = download_voice_embedding_from_hash( + &voice.bucket_hash, + &voice.title, + &deps.private_bucket_client, + &work_temp_dir_string + ).await?; + // // Download embeddings file using embedding token + // // Create a temp dir to download things to + job_progress_reporter + .log_status("running inference") + .map_err(|e| ProcessSingleJobError::Other(e))?; - // ==================== UPLOAD AUDIO TO BUCKET ==================== // - let maybe_args = job.maybe_inference_args - .as_ref() - .map(|args| args.args.as_ref()) - .flatten(); + // Run Inference + let command_exit_status = + args.job_dependencies.job_type_details.vall_e_x.inference_command.execute_inference( + InferenceArgs { + input_embedding_path: format!("{}_weights.npz", &voice.title), // name of the embedding.npz in the tmp dir + input_text: String::from(text), // text + output_file_name: String::from("output.wav"), // output file name in the output folder + stderr_output_file: String::from("test.txt"), + } + ); + + // // upload audio to bucket + // // deletetemp dir files - let workdir = work_temp_dir.path().to_path_buf(); - let stderr_output_file = work_temp_dir.path().join("stderr.txt"); - let inference_start_time = Instant::now(); + // // ==================== UPLOAD AUDIO TO BUCKET ==================== // + + // // delete temp dir files + + // // ==================== UPLOAD AUDIO TO BUCKET ==================== // + // let maybe_args = job.maybe_inference_args + // .as_ref() + // .map(|args| args.args.as_ref()) + // .flatten(); + + // let workdir = work_temp_dir.path().to_path_buf(); + // let stderr_output_file = work_temp_dir.path().join("stderr.txt"); + // let inference_start_time = Instant::now(); // upload audio to public bucket ( for voice ) // return ok Err(ProcessSingleJobError::InvalidJob(anyhow!("this job flow is not yet implemented"))) } - mod tests { #[test] fn test_length_zero() { println!("Hello Test") } -} \ No newline at end of file +} diff --git a/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/vall_e_x_inference_command.rs b/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/vall_e_x_inference_command.rs index 5bb9a8a67e..62daae6923 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/vall_e_x_inference_command.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/vall_e_x_inference_command.rs @@ -43,7 +43,7 @@ static IGNORED_ENVIRONMENT_VARS : Lazy> = Lazy::new(|| { // so we have all the files downloaded for the inputs. pub struct InferenceArgs> { /// --driven_audio: path to the input audio - pub input_embedding: P, // name of the embedding.npz in the tmp dir + pub input_embedding_path: P, // name of the embedding.npz in the work dir pub input_text: P, // name of the text /// --result_file: path to final file output pub output_file_name: P, // output file name in the output folder @@ -97,7 +97,6 @@ impl VallEXInferenceCommand { let maybe_inference_command = easyenv::get_env_string_optional( "VALL_E_X_INFERENCE_COMMAND"); - // EVENTUALLY should remove the check // Optional, eg. `./infer.py`. Typically we'll use the command form instead. let maybe_inference_executable = easyenv::get_env_pathbuf_optional( @@ -187,15 +186,30 @@ impl VallEXInferenceCommand { // ===== Begin Python Args ===== - command.push_str(" -text "); + command.push_str(" --text "); command.push_str(&path_to_string(args.input_text)); - command.push_str(" -use_prompt_name "); - command.push_str(&path_to_string(args.input_embedding)); + command.push_str(" --prompt-name "); + command.push_str(&path_to_string(args.input_embedding_path)); - command.push_str(" -output_file_name "); + command.push_str(" --audio-name "); command.push_str(&path_to_string(args.output_file_name)); + 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-path "); + command.push_str(&path_to_string(Path::new("/tmp/downloads/zero_shot_tts/vall-e-x_1.0/medium.pt"))); + + 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(" --vallex-path "); + command.push_str(&path_to_string(Path::new("/tmp/downloads/zero_shot_tts/vall-e-x_1.0/vallex-checkpoint.pt"))); + // ===== End Python Args ===== if let Some(docker_options) = self.maybe_docker_options.as_ref() { diff --git a/crates/service/web/storyteller_web/src/routes.rs b/crates/service/web/storyteller_web/src/routes.rs index 15f2995719..a62c50cd24 100644 --- a/crates/service/web/storyteller_web/src/routes.rs +++ b/crates/service/web/storyteller_web/src/routes.rs @@ -171,7 +171,7 @@ pub fn add_routes (app: App) -> App app = add_subscription_routes(app); /* /v1/subscriptions/... */ // TODO find a long term feature flag solution, since this code is likely deployed into production we don't want the route found. - let enable_voice_designer_route = false; + let enable_voice_designer_route = true; if enable_voice_designer_route { app = add_voice_designer_routes(app); /* /v1/voice_designer */ }