From 9d4e426acef6c953561cd1fac7322c19b729a592 Mon Sep 17 00:00:00 2001 From: Michael Chung Date: Wed, 18 Oct 2023 18:02:30 -0400 Subject: [PATCH] It is now being tested then we work on the command side. --- .../generic_synthetic_ids/id_category.rs | 11 +++ .../media_file_origin_model_type.rs | 6 ++ .../config/inference-job.development.env | 2 +- .../job/job_types/tts/vall_e_x/process_job.rs | 67 ++++++++++++------- 4 files changed, 61 insertions(+), 25 deletions(-) diff --git a/crates/_types/enums/src/by_table/generic_synthetic_ids/id_category.rs b/crates/_types/enums/src/by_table/generic_synthetic_ids/id_category.rs index 2611d797ff..ebf6b0ec2f 100644 --- a/crates/_types/enums/src/by_table/generic_synthetic_ids/id_category.rs +++ b/crates/_types/enums/src/by_table/generic_synthetic_ids/id_category.rs @@ -21,6 +21,10 @@ pub enum IdCategory { #[serde(rename = "lipsync_animation")] LipsyncAnimation, + /// Results from the zero shot tts + #[serde(rename = "zs_audio_tts")] + ZeroShotTTS, + /// Zs dataset which lives in the zs_voice_datasets table #[serde(rename = "zs_dataset")] ZsDataset, @@ -42,6 +46,7 @@ impl IdCategory { Self::LipsyncAnimation => "lipsync_animation", Self::ZsDataset => "zs_dataset", Self::ZsVoice => "zs_voice", + Self::ZeroShotTTS => "zs_audio_tts" } } @@ -51,6 +56,7 @@ impl IdCategory { "lipsync_animation" => Ok(Self::LipsyncAnimation), "zs_dataset" => Ok(Self::ZsDataset), "zs_voice" => Ok(Self::ZsVoice), + "zs_audio_tts" => Ok(Self::ZeroShotTTS), _ => Err(format!("invalid value: {:?}", value)), } } @@ -63,6 +69,7 @@ impl IdCategory { Self::LipsyncAnimation, Self::ZsDataset, Self::ZsVoice, + Self::ZeroShotTTS ]) } } @@ -76,6 +83,7 @@ mod tests { fn test_serialization() { assert_serialization(IdCategory::MediaFile, "media_file"); assert_serialization(IdCategory::LipsyncAnimation, "lipsync_animation"); + assert_serialization(IdCategory::ZeroShotTTS, "zs_audio_tts"); } #[test] @@ -84,6 +92,7 @@ mod tests { assert_eq!(IdCategory::LipsyncAnimation.to_str(), "lipsync_animation"); assert_eq!(IdCategory::ZsDataset.to_str(), "zs_dataset"); assert_eq!(IdCategory::ZsVoice.to_str(), "zs_voice"); + assert_eq!(IdCategory::ZeroShotTTS.to_str(),"zs_audio_tts"); } #[test] @@ -92,6 +101,7 @@ mod tests { assert_eq!(IdCategory::from_str("lipsync_animation").unwrap(), IdCategory::LipsyncAnimation); assert_eq!(IdCategory::from_str("zs_dataset").unwrap(), IdCategory::ZsDataset); assert_eq!(IdCategory::from_str("zs_voice").unwrap(), IdCategory::ZsVoice); + assert_eq!(IdCategory::from_str("zs_audio_tts").unwrap(), IdCategory::ZeroShotTTS); } #[test] @@ -103,6 +113,7 @@ mod tests { assert_eq!(variants.pop_first(), Some(IdCategory::LipsyncAnimation)); assert_eq!(variants.pop_first(), Some(IdCategory::ZsDataset)); assert_eq!(variants.pop_first(), Some(IdCategory::ZsVoice)); + assert_eq!(variants.pop_first(), Some(IdCategory::ZeroShotTTS)); assert_eq!(variants.pop_first(), None); // Generated check diff --git a/crates/_types/enums/src/by_table/media_files/media_file_origin_model_type.rs b/crates/_types/enums/src/by_table/media_files/media_file_origin_model_type.rs index b401758897..eb5853fbc0 100644 --- a/crates/_types/enums/src/by_table/media_files/media_file_origin_model_type.rs +++ b/crates/_types/enums/src/by_table/media_files/media_file_origin_model_type.rs @@ -14,6 +14,7 @@ use strum::EnumIter; pub enum MediaFileOriginModelType { /// SadTalker -- v1, we may add another enum value for future versions SadTalker, + VallEX, } // TODO(bt, 2022-12-21): This desperately needs MySQL integration tests! @@ -25,12 +26,14 @@ impl MediaFileOriginModelType { pub fn to_str(&self) -> &'static str { match self { Self::SadTalker => "sad_talker", + Self::VallEX => "vall_e_x" } } pub fn from_str(value: &str) -> Result { match value { "sad_talker" => Ok(Self::SadTalker), + "vall_e_x" => Ok(Self::VallEX), _ => Err(format!("invalid value: {:?}", value)), } } @@ -44,16 +47,19 @@ mod tests { #[test] fn test_serialization() { assert_serialization(MediaFileOriginModelType::SadTalker, "sad_talker"); + assert_serialization(MediaFileOriginModelType::VallEX, "sad_talker"); } #[test] fn test_to_str() { assert_eq!(MediaFileOriginModelType::SadTalker.to_str(), "sad_talker"); + assert_eq!(MediaFileOriginModelType::VallEX.to_str(), "vall_e_x"); } #[test] fn test_from_str() { assert_eq!(MediaFileOriginModelType::from_str("sad_talker").unwrap(), MediaFileOriginModelType::SadTalker); + assert_eq!(MediaFileOriginModelType::from_str("vall_e_x").unwrap(), MediaFileOriginModelType::VallEX); assert!(MediaFileOriginModelType::from_str("foo").is_err()); } } diff --git a/crates/service/job/inference_job/config/inference-job.development.env b/crates/service/job/inference_job/config/inference-job.development.env index 303534d758..e580785afe 100644 --- a/crates/service/job/inference_job/config/inference-job.development.env +++ b/crates/service/job/inference_job/config/inference-job.development.env @@ -20,8 +20,8 @@ REDIS_FOR_KEEPALIVE_URL="redis://localhost/0" VALL_E_X_INFERENCE_ROOT_DIRECTORY = "/home/tensor/code/storyteller/storyteller-ml/tts/VALL-E-X" VALL_E_X_INFERENCE_COMMAND = "python3 main.py" VALL_E_X_INFERENCE_MAYBE_VENV_COMMAND = "source .venv/bin/activate" -# VALL_E_X_INFERENCE_MAYBE_DOCKER_IMAGE = "" # fill this out for dev? +# VALL_E_X_INFERENCE_MAYBE_DOCKER_IMAGE = "" # fill this out for dev? #VALL_E_X_INFERENCE_ROOT_DIRECTORY="%storyteller_ml%/tts/VALL-E-X" #VALL_E_X_INFERENCE_COMMAND="python3 main.py" #VALL_E_X_INFERENCE_MAYBE_VENV_COMMAND="source .env/bin/activate" 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 e1e1e2b983..dd9b22d8b5 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 @@ -1,4 +1,5 @@ use anyhow::anyhow; +use buckets::public::media_files::original_file::MediaFileBucketPath; use cloud_storage::bucket_client::BucketClient; use container_common::filesystem; use log::{ error, info, warn }; @@ -17,6 +18,9 @@ use crate::job_dependencies::JobDependencies; use crate::job::job_types::tts::vall_e_x::vall_e_x_inference_command::InferenceArgs; use cloud_storage::bucket_path_unifier::BucketPathUnifier; +const BUCKET_FILE_PREFIX: &str = "fakeyou_"; +const BUCKET_FILE_EXTENSION: &str = ".wav"; +const MIME_TYPE: &str = "audio/wav"; pub struct VoiceFile { pub filesystem_path: PathBuf, } @@ -35,9 +39,9 @@ pub async fn download_voice_embedding_from_hash( // let filesystem_path = format!("{}{}_weights.npz",path, name); let mut path = path.clone(); - let file_name = format!("{}",name); + let file_name = format!("{}", name); path.push(&file_name); - + let result = private_bucket_client.download_file_to_disk(object_path, &path).await; let voice_file = VoiceFile { @@ -115,7 +119,7 @@ pub async fn process_job( 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))?; - + let workdir = work_temp_dir.path().to_path_buf(); let file_name = format!("{}_weights.npz", &voice.title); @@ -128,7 +132,7 @@ pub async fn process_job( &workdir ).await?; - println!("voicefile path! {}",voiceFile.filesystem_path.to_string_lossy()); + println!("voicefile path! {}", voiceFile.filesystem_path.to_string_lossy()); // Download embeddings file using embedding token // Create a temp dir to download things to @@ -138,42 +142,58 @@ pub async fn process_job( let inference_start_time = Instant::now(); + let output_file_name = String::from("output.wav"); // Run Inference let command_exit_status = args.job_dependencies.job_type_details.vall_e_x.inference_command.execute_inference( InferenceArgs { - // TODO: zero shot tts folder to tmp - //input_embedding_path: "/tmp/downloads/zero_shot_tts", // name of the embedding.npz in the tmp dir input_embedding_path: &workdir, - input_embedding_name: file_name, + input_embedding_name: file_name, input_text: String::from(text), // text - output_file_name: String::from("output.wav"), // output file name in the output folder + output_file_name: output_file_name.clone(), // output file name in the output folder stderr_output_file: String::from("zero_shot.txt"), } ); - + // upload audio to bucket info!("Uploading media ..."); let result_bucket_location = MediaFileBucketPath::generate_new( Some(BUCKET_FILE_PREFIX), - Some(BUCKET_FILE_EXTENSION)); - let result_bucket_object_pathbuf = result_bucket_location.to_full_object_pathbuf(); + Some(BUCKET_FILE_EXTENSION) + ); + let result_bucket_object_pathbuf = result_bucket_location.to_full_object_pathbuf(); - args.job_dependencies.public_bucket_client.upload_filename_with_content_type( - &result_bucket_object_pathbuf, - &finished_file, - &mimetype) // TODO: We should check the mimetype to make sure bad payloads can't get uploaded - .await + let mut finished_file = work_temp_dir.path().to_path_buf(); + finished_file.push(&output_file_name); + + args.job_dependencies.public_bucket_client + .upload_filename_with_content_type( + &result_bucket_object_pathbuf, + &finished_file, + &MIME_TYPE + ).await // TODO: We should check the mimetype to make sure bad payloads can't get uploaded .map_err(|e| ProcessSingleJobError::Other(e))?; - // ==================== UPLOAD AUDIO TO BUCKET ==================== // - - // upload audio to public bucket ( for voice ) - // return ok - job_progress_reporter.log_status("done") - .map_err(|e| ProcessSingleJobError::Other(e))?; + // ==================== UPLOAD AUDIO TO BUCKET ==================== + + job_progress_reporter.log_status("done").map_err(|e| ProcessSingleJobError::Other(e))?; + + // insert into db the record + let (media_file_token, id) = insert_media_file_from_face_zero_shot(InsertArgs { + pool: &args.job_dependencies.mysql_pool, + job: &job, + maybe_mime_type: Some(&mimetype), + file_size_bytes, + sha256_checksum: &file_checksum, + public_bucket_directory_hash: result_bucket_location.get_object_hash(), + maybe_public_bucket_prefix: Some(BUCKET_FILE_PREFIX), + maybe_public_bucket_extension: Some(BUCKET_FILE_EXTENSION), + is_on_prem: args.job_dependencies.container.is_on_prem, + worker_hostname: &args.job_dependencies.container.hostname, + worker_cluster: &args.job_dependencies.container.cluster_name, + }).await.map_err(|e| ProcessSingleJobError::Other(e))?; info!("Job {:?} complete success! Downloaded, ran inference, and uploaded. Saved model record: {}, Result Token: {}", job.id, id, &media_file_token); @@ -184,9 +204,8 @@ pub async fn process_job( entity_token: media_file_token.to_string(), }), inference_duration, - }) + }); - //Err(ProcessSingleJobError::InvalidJob(anyhow!("this job flow is not yet implemented"))) } mod tests { #[test]