From c6a6676f088d946df3874f1798d410983dc520ab Mon Sep 17 00:00:00 2001 From: Kasisnu <5044120+kasisnu@users.noreply.github.com> Date: Fri, 5 Jan 2024 18:11:43 +0530 Subject: [PATCH 1/3] adds styletts2 as voice designer model --- .../src/private/zs_voices/bucket_directory.rs | 2 + .../inference_model_type.rs | 9 + .../src/by_table/zs_voices/model_type.rs | 10 +- .../src/job/job_types/tts/mod.rs | 1 + .../src/job/job_types/tts/styletts2/mod.rs | 7 + .../tts/styletts2/model_downloaders.rs | 94 +++ .../job_types/tts/styletts2/process_job.rs | 45 ++ .../tts/styletts2/process_job_create_voice.rs | 370 ++++++++++++ .../styletts2/process_job_inference_voice.rs | 318 ++++++++++ .../tts/styletts2/styletts2_dependencies.rs | 20 + .../styletts2/styletts2_inference_command.rs | 549 ++++++++++++++++++ .../job_types/tts/styletts2/validate_job.rs | 84 +++ .../src/job_specific_dependencies.rs | 11 +- 13 files changed, 1518 insertions(+), 2 deletions(-) create mode 100644 crates/service/job/inference_job/src/job/job_types/tts/styletts2/mod.rs create mode 100644 crates/service/job/inference_job/src/job/job_types/tts/styletts2/model_downloaders.rs create mode 100644 crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job.rs create mode 100644 crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_create_voice.rs create mode 100644 crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_inference_voice.rs create mode 100644 crates/service/job/inference_job/src/job/job_types/tts/styletts2/styletts2_dependencies.rs create mode 100644 crates/service/job/inference_job/src/job/job_types/tts/styletts2/styletts2_inference_command.rs create mode 100644 crates/service/job/inference_job/src/job/job_types/tts/styletts2/validate_job.rs diff --git a/crates/_types/buckets/src/private/zs_voices/bucket_directory.rs b/crates/_types/buckets/src/private/zs_voices/bucket_directory.rs index 9d4ecee304..0c95cf153a 100644 --- a/crates/_types/buckets/src/private/zs_voices/bucket_directory.rs +++ b/crates/_types/buckets/src/private/zs_voices/bucket_directory.rs @@ -26,12 +26,14 @@ impl ModelCategory { #[derive(Debug, Clone, Copy, Eq, PartialEq)] pub enum ModelType { VallEx, + StyleTTS2, } impl ModelType { pub fn to_str(&self) -> &'static str { match self { ModelType::VallEx => "vallex", + ModelType::StyleTTS2 => "styletts2", } } } diff --git a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_model_type.rs b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_model_type.rs index cd71c3246a..48d7de9dce 100644 --- a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_model_type.rs +++ b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_model_type.rs @@ -35,6 +35,8 @@ pub enum InferenceModelType { RerenderAVideo, #[serde(rename = "mocap_net")] MocapNet, + #[serde(rename = "styletts2")] + StyleTTS2, } // TODO(bt, 2022-12-21): This desperately needs MySQL integration tests! @@ -53,6 +55,7 @@ impl InferenceModelType { Self::VallEX => "vall_e_x", Self::RerenderAVideo => "rerender_a_video", Self::MocapNet => "mocap_net", + Self::StyleTTS2 => "styletts2", } } @@ -66,6 +69,7 @@ impl InferenceModelType { "vall_e_x" => Ok(Self::VallEX), "rerender_a_video" => Ok(Self::RerenderAVideo), "mocap_net" => Ok(Self::MocapNet), + "styletts2" => Ok(Self::StyleTTS2), _ => Err(format!("invalid value: {:?}", value)), } } @@ -82,6 +86,7 @@ impl InferenceModelType { InferenceModelType::VallEX, InferenceModelType::RerenderAVideo, InferenceModelType::MocapNet, + InferenceModelType::StyleTTS2, ]) } } @@ -101,6 +106,7 @@ mod tests { assert_serialization(InferenceModelType::VallEX, "vall_e_x"); assert_serialization(InferenceModelType::RerenderAVideo, "rerender_a_video"); assert_serialization(InferenceModelType::MocapNet, "mocap_net"); + assert_serialization(InferenceModelType::StyleTTS2, "styletts2"); } #[test] @@ -113,6 +119,7 @@ mod tests { assert_eq!(InferenceModelType::VallEX.to_str(), "vall_e_x"); assert_eq!(InferenceModelType::RerenderAVideo.to_str(), "rerender_a_video"); assert_eq!(InferenceModelType::MocapNet.to_str(), "mocap_net"); + assert_eq!(InferenceModelType::StyleTTS2.to_str(), "styletts2"); } #[test] @@ -125,6 +132,7 @@ mod tests { assert_eq!(InferenceModelType::from_str("vall_e_x").unwrap(), InferenceModelType::VallEX); assert_eq!(InferenceModelType::from_str("rerender_a_video").unwrap(), InferenceModelType::RerenderAVideo); assert_eq!(InferenceModelType::from_str("mocap_net").unwrap(), InferenceModelType::MocapNet); + assert_eq!(InferenceModelType::from_str("styletts2").unwrap(), InferenceModelType::StyleTTS2); } #[test] @@ -140,6 +148,7 @@ mod tests { assert_eq!(variants.pop_first(), Some(InferenceModelType::VallEX)); assert_eq!(variants.pop_first(), Some(InferenceModelType::RerenderAVideo)); assert_eq!(variants.pop_first(), Some(InferenceModelType::MocapNet)); + assert_eq!(variants.pop_first(), Some(InferenceModelType::StyleTTS2)); assert_eq!(variants.pop_first(), None); // Generated check diff --git a/crates/_types/enums/src/by_table/zs_voices/model_type.rs b/crates/_types/enums/src/by_table/zs_voices/model_type.rs index d8f022035d..a5e4dca813 100644 --- a/crates/_types/enums/src/by_table/zs_voices/model_type.rs +++ b/crates/_types/enums/src/by_table/zs_voices/model_type.rs @@ -16,6 +16,8 @@ pub enum ZsVoiceModelType { /// TTS-type zero shot models #[serde(rename = "vall-e-x")] VallEX, + #[serde(rename = "styletts2")] + StyleTTS2, } // TODO(bt, 2023-01-17): This desperately needs MySQL integration tests! @@ -27,12 +29,14 @@ impl ZsVoiceModelType { pub fn to_str(&self) -> &'static str { match self { Self::VallEX=> "vall-e-x", + Self::StyleTTS2 => "styletts2", } } pub fn from_str(value: &str) -> Result { match value { "vall-e-x" => Ok(Self::VallEX), + "styletts2" => Ok(Self::StyleTTS2), _ => Err(format!("invalid value: {:?}", value)), } } @@ -42,6 +46,7 @@ impl ZsVoiceModelType { // NB: BTreeSet::from() isn't const, but not worth using LazyStatic, etc. BTreeSet::from([ Self::VallEX, + Self::StyleTTS2, ]) } } @@ -66,11 +71,13 @@ mod tests { #[test] fn to_str() { assert_eq!(ZsVoiceModelType::VallEX.to_str(), "vall-e-x"); + assert_eq!(ZsVoiceModelType::StyleTTS2.to_str(), "styletts2"); } #[test] fn from_str() { assert_eq!(ZsVoiceModelType::from_str("vall-e-x").unwrap(), ZsVoiceModelType::VallEX); + assert_eq!(ZsVoiceModelType::from_str("styletts2").unwrap(), ZsVoiceModelType::StyleTTS2); assert!(ZsVoiceModelType::from_str("foo").is_err()); } } @@ -81,8 +88,9 @@ mod tests { #[test] fn all_variants() { let mut variants = ZsVoiceModelType::all_variants(); - assert_eq!(variants.len(), 1); + assert_eq!(variants.len(), 2); assert_eq!(variants.pop_first(), Some(ZsVoiceModelType::VallEX)); + assert_eq!(variants.pop_first(), Some(ZsVoiceModelType::StyleTTS2)); assert_eq!(variants.pop_first(), None); } } diff --git a/crates/service/job/inference_job/src/job/job_types/tts/mod.rs b/crates/service/job/inference_job/src/job/job_types/tts/mod.rs index 26d581e703..137d21bf61 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/mod.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/mod.rs @@ -1,4 +1,5 @@ pub mod process_single_tts_job; +pub mod styletts2; pub mod tacotron2_v2_early_fakeyou; pub mod vits; pub mod vall_e_x; diff --git a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/mod.rs b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/mod.rs new file mode 100644 index 0000000000..5623a593ff --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/mod.rs @@ -0,0 +1,7 @@ +pub mod model_downloaders; +pub mod process_job; +pub mod process_job_create_voice; +pub mod process_job_inference_voice; +pub mod validate_job; +pub mod styletts2_dependencies; +pub mod styletts2_inference_command; \ No newline at end of file diff --git a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/model_downloaders.rs b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/model_downloaders.rs new file mode 100644 index 0000000000..18bcc1c931 --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/model_downloaders.rs @@ -0,0 +1,94 @@ +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 +} + +crate::impl_model_downloader!( + StyleTTS2ModelMapping, + "Vall-e-x Model Mapping", + "VALL_E_X_MAPPING_BUCKET_PATH", + "/dependencies/zero_shot_tts/vall-e-x_1.0/StyleTTS2-checkpoint.pt", + "VALL_E_X_MAPPING_FILESYSTEM_PATH", + "/tmp/downloads/zero_shot_tts/vall-e-x_1.0/StyleTTS2-checkpoint.pt" +); + +crate::impl_model_downloader!( + WhisperLargeModelMapping, + "whisper Large Model Mapping", + "WHISPER_LARGE_MAPPING_BUCKET_PATH", + "/dependencies/zero_shot_tts/vall-e-x_1.0/whisper-large/model.safetensors", + "WHISPER_LARGE_MAPPING_FILESYSTEM_PATH", + "/tmp/downloads/zero_shot_tts/vall-e-x_1.0/whisper-large/model.safetensors" +); + +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.pt", + "WHISPER_MEDIUM_TENSOR_FILESYSTEM_PATH", + "/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/vocos-encodec/encodec_pytorch_model.bin", + "VOCOS_ENCODEC_TENSOR_FILESYSTEM_PATH", + "/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 StyleTTS2Downloaders { + 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] + } +} + +mod test { + // use std::path::Path; + + // use crate::job::job_types::lipsync::sad_talker::model_downloaders::{SadTalkerModelFirstMapping, SadTalkerModelSecondMapping}; + // use crate::util::model_downloader::ModelDownloader; + + // #[test] + // fn test_downloader_model_first_mapping() { + // let model = SadTalkerModelFirstMapping::default(); + // assert_eq!(model.get_model_name(), "sad talker model first mapping"); + // assert_eq!(model.get_cloud_bucket_path(), "/animation_sadtalker/model/mapping_00109-model.pth.tar"); + // assert_eq!(model.get_filesystem_path(), Path::new("/tmp/downloads/sadtalker/mapping_00109-model.pth.tar")); + // } + + // #[test] + // fn test_downloader_model_second_mapping() { + // let model = SadTalkerModelSecondMapping::default(); + // assert_eq!(model.get_model_name(), "sad talker model second mapping"); + // assert_eq!(model.get_cloud_bucket_path(), "/animation_sadtalker/model/mapping_00229-model.pth.tar"); + // assert_eq!(model.get_filesystem_path(), Path::new("/tmp/downloads/sadtalker/mapping_00229-model.pth.tar")); + // } + +} diff --git a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job.rs b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job.rs new file mode 100644 index 0000000000..e98c004db0 --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job.rs @@ -0,0 +1,45 @@ +use anyhow::anyhow; + +use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; + +use crate::job::job_loop::job_success_result::JobSuccessResult; +use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; +use crate::job::job_types::tts::styletts2::process_job_create_voice::process_create_voice; +use crate::job::job_types::tts::styletts2::process_job_inference_voice::process_inference_voice; +use crate::job::job_types::tts::styletts2::validate_job::validate_job; +use crate::job_dependencies::JobDependencies; + +use super::validate_job::JobType; + +// This will download everything get into the root host OS then ... will invoke inference using the pathes from the files invoked +pub struct StyleTTS2ProcessJobArgs<'a> { + pub job_dependencies: &'a JobDependencies, + pub job: &'a AvailableInferenceJob, +} + +// query using the token then grab the bucket hash +pub async fn process_job( + args: StyleTTS2ProcessJobArgs<'_> +) -> Result { + let job = args.job; + + // get args token + let jobArgs = validate_job(&job)?; // bubbles error up + + match jobArgs.job_type { + JobType::Create => { + if let Some(voice_dataset_token) = jobArgs.voice_dataset_token { + process_create_voice(args, voice_dataset_token).await + } else { + Err(ProcessSingleJobError::Other(anyhow!("Missing Dataset Token?"))) + } + } + JobType::Inference => { + if let Some(voice_token) = jobArgs.voice_token { + process_inference_voice(args, voice_token).await + } else { + Err(ProcessSingleJobError::Other(anyhow!("Missing Voice Token?"))) + } + } + } +} diff --git a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_create_voice.rs b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_create_voice.rs new file mode 100644 index 0000000000..d64cc4f809 --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_create_voice.rs @@ -0,0 +1,370 @@ +use std::fs::read_to_string; +use std::path::PathBuf; +use std::time::Instant; + +use anyhow::anyhow; +use log::{error, info, warn}; + +use buckets::private::zs_voices::bucket_directory::{ModelCategory, ModelType}; +use buckets::private::zs_voices::bucket_file_path::ZeroShotVoiceEmbeddingBucketPath; +use buckets::public::media_files::bucket_file_path::MediaFileBucketPath; +use cloud_storage::bucket_client::BucketClient; +use cloud_storage::bucket_path_unifier::BucketPathUnifier; +use enums::by_table::generic_inference_jobs::inference_result_type::InferenceResultType; +use enums::by_table::zs_voices::encoding_type::ZsVoiceEncodingType; +use enums::by_table::zs_voices::model_category::ZsVoiceModelCategory; +use enums::by_table::zs_voices::model_type::ZsVoiceModelType; +use enums::common::visibility::Visibility; +use mysql_queries::queries::voice_designer::datasets::get_dataset::get_dataset_by_token; +use mysql_queries::queries::voice_designer::datasets::get_dataset::ZsDataset; +use mysql_queries::queries::voice_designer::voice_samples::list_dataset_samples_for_dataset_token::DatasetSampleRecordForList; +use mysql_queries::queries::voice_designer::voice_samples::list_dataset_samples_for_dataset_token::list_dataset_samples_for_dataset_token; +use mysql_queries::queries::voice_designer::voices::create_voice::create_voice; +use mysql_queries::queries::voice_designer::voices::create_voice::CreateVoiceArgs; +use tokens::tokens::users::UserToken; + +use crate::job; +use crate::job::job_loop::job_success_result::JobSuccessResult; +use crate::job::job_loop::job_success_result::ResultEntity; +use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; +use crate::job::job_types::tts::styletts2::process_job::StyleTTS2ProcessJobArgs; + +const BUCKET_FILE_PREFIX_CREATE: &str = "fakeyou_"; +const BUCKET_FILE_EXTENSION_CREATE: &str = ".bin"; +const MIME_TYPE_CREATE: &str = "application/x-binary"; + +/// We'll increment this if we change the underlying StyleTTS2 code in an incompatible way. +const CURRENT_STYLE_TTS2_MODEL_VERSION: u64 = 0; + +pub async fn process_create_voice( + args: StyleTTS2ProcessJobArgs<'_>, + dataset_token: String +) -> Result { + let deps = args.job_dependencies; + let job = args.job; + let mysql_pool = &deps.db.mysql_pool; + + let model_dependencies = deps + .job + .job_specific_dependencies + .maybe_styletts2_dependencies + .as_ref() + .ok_or_else(|| ProcessSingleJobError::JobSystemMisconfiguration(Some("missing StyleTTS2 dependencies".to_string())))?; + + // get some globals + let mut job_progress_reporter = deps.clients.job_progress_reporter + .new_generic_inference(job.inference_job_token.as_str()) + .map_err(|e| ProcessSingleJobError::Other(anyhow!(e)))?; + + info!("token! {}", dataset_token); + let voice_dataset_token = tokens::tokens::zs_voice_datasets::ZsVoiceDatasetToken(dataset_token); + + let creator_ip_address = &job.creator_ip_address; + + let creator_user_token = match &job.maybe_creator_user_token { + Some(token) => { + UserToken::new_from_str(token) + }, + None => { + return Err(ProcessSingleJobError::InvalidJob(anyhow!("Missing Creator User Token"))); + } + }; + + // STEP 1. SETUP A TEMP DIRECTORY + let work_temp_dir = format!("/tmp/temp_zeroshot_create_voice_{}", job.id.0); + // 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))?; + + let workdir = work_temp_dir.path().to_path_buf(); + + // STEP 2. Get dataset for the title for the voice + let voice_dataset = get_dataset_by_token(&voice_dataset_token, false, &mysql_pool).await; + let single_dataset: ZsDataset; + match voice_dataset { + Ok(val) => { + match val { + Some(val) => { + single_dataset = val; + } + None => { + return Err(ProcessSingleJobError::InvalidJob(anyhow!("Missing Dataset"))); + } + } + } + Err(e) => { + return Err(ProcessSingleJobError::from_anyhow_error(e)); + } + } + info!("Title:{} Token:{}", &single_dataset.title, &single_dataset.token); + + // STEP 3. Download dataset each audio file + let result = list_dataset_samples_for_dataset_token( + &voice_dataset_token, + false, + &mysql_pool + ).await; + + let dataset: Vec = match result { + Ok(val) => val, + Err(e) => return Err(ProcessSingleJobError::from_anyhow_error(e)), + }; + + 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 = 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()); + + info!("Record=> hash:{} prefix:{:?} extension:{:?}", + record.public_bucket_directory_hash, + prefix, + extension + ); + + let audio_media_file = MediaFileBucketPath::from_object_hash( + &record.public_bucket_directory_hash, + prefix, + extension + ); + + info!("Download using audio_media_file_path: {:?}", audio_media_file.to_full_object_pathbuf()); + + let file_name_wav = format!("{}.wav", index); + + let mut file_path = PathBuf::new(); + file_path.push(workdir.clone()); + file_path.push(file_path.clone()); + file_path.push(file_name_wav); + + info!("Downloading to path: {:?}", file_path); + + // TODO: we might want to catch the error and not include the pathes into download dataset? + let result = deps.buckets.public_bucket_client.download_file_to_disk( + audio_media_file.to_full_object_pathbuf(), + &file_path + ).await; + + if let Err(err) = result { + error!("could not download sample: {:?}", err); + return Err(ProcessSingleJobError::from_anyhow_error(err)); + } + + info!("FilePath to clone voice: {}", file_path.to_string_lossy()); + downloaded_dataset.push(file_path.clone()); + } + + info!("Dataset Length {}", downloaded_dataset.len()); + + // STEP 4 Download the models + info!("Download models (if not present)..."); + for downloader in model_dependencies.downloaders.all_downloaders() { + let result = downloader.download_if_not_on_filesystem( + &args.job_dependencies.buckets.private_bucket_client, + &args.job_dependencies.fs.scoped_temp_dir_creator_for_short_lived_downloads + ).await; + if let Err(err) = result { + error!("could not download: {:?}", err); + return Err(err); + } + } + + job_progress_reporter + .log_status("running inference") + .map_err(|e| ProcessSingleJobError::Other(e))?; + + // Command line arg for a list of paths to insert the container + let audio_files = join_paths(downloaded_dataset); + + 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 stderr_output_file = work_temp_dir.path().join("zero_shot_create_voice_err.txt"); + + let inference_start_time = Instant::now(); + + // Run Inference + 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, + stderr_output_file: &stderr_output_file, + } + ); + + let inference_duration = Instant::now().duration_since(inference_start_time); + + if !command_exit_status.is_success() { + error!("Create Embedding Failed: {:?}", command_exit_status); + + let error = ProcessSingleJobError::Other(anyhow!("CommandExitStatus: {:?}", command_exit_status)); + + 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); + //safe_delete_temp_file(&usable_image_path); + //safe_delete_temp_file(&output_video_fs_path); + //safe_delete_temp_file(&stderr_output_file); + //safe_delete_temp_directory(&work_temp_dir); + + return Err(error); + } + + 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, + CURRENT_STYLE_TTS2_MODEL_VERSION + ); + + let embedding_bucket_object_pathbuf = embedding_bucket_location.to_full_object_pathbuf(); + + // 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. + finished_file.push(&output_bucket_file_name); + + info!("Upload File Path: {:?}", finished_file); + info!("Upload Bucket Path: {:?}", embedding_bucket_object_pathbuf); + + args.job_dependencies.buckets.private_bucket_client + .upload_filename_with_content_type( + &embedding_bucket_object_pathbuf, + &finished_file, + &MIME_TYPE_CREATE + ) + .await // TODO: We should check the mimetype to make sure bad payloads can't get uploaded + .map_err(|e| ProcessSingleJobError::Other(e))?; + + // CLEARIFY! these items + // 1.Should this be object hash? + // 2.As well as this what should the voice name be? + let voice_name = single_dataset.title; + + // insert record + let voice_token = create_voice(CreateVoiceArgs { + dataset_token: &voice_dataset_token, + model_category: ZsVoiceModelCategory::Tts, + model_type: ZsVoiceModelType::StyleTTS2, + model_version: CURRENT_STYLE_TTS2_MODEL_VERSION, + model_encoding_type: ZsVoiceEncodingType::Encodec, + voice_title: &voice_name, + bucket_hash: embedding_bucket_location.get_object_hash().clone(), + maybe_creator_user_token: Some(&creator_user_token), + creator_ip_address: &creator_ip_address, + creator_set_visibility: Visibility::Public, + mysql_pool, + }).await + .map_err(|e| ProcessSingleJobError::Other(e))?; + + Ok(JobSuccessResult { + maybe_result_entity: Some(ResultEntity { + entity_type: InferenceResultType::ZeroShotVoiceEmbedding, + entity_token: voice_token.to_string(), + }), + inference_duration, + }) +} + +pub struct AudioFile { + pub filesystem_path: PathBuf, +} + +pub async fn download_audio_from_hash( + bucket_hash: &str, + name: &str, + private_bucket_client: &BucketClient, + path: &PathBuf +) -> Result { + let unifer = BucketPathUnifier::default_paths(); + let object_path = unifer.zero_shot_tts_speaker_encoding(bucket_hash, 0); + + let mut path = path.clone(); + + let file_name = format!("{}", name); + path.push(&file_name); + + private_bucket_client.download_file_to_disk(object_path, &path) + .await + .map_err(|err| { + error!("Could not download audio file: {err}"); + ProcessSingleJobError::from_anyhow_error(anyhow!("Could not download audio file: {err}")) + })?; + + let audio_file = AudioFile { + filesystem_path: PathBuf::from(&path.clone()), + }; + + Ok(audio_file) +} + +fn join_paths(paths: Vec) -> String { + paths + .into_iter() + .map(|p| format!("\"{}\"", p.display())) + .collect::>() + .join(" ") +} + +#[cfg(test)] +mod tests { + use std::path::PathBuf; + + use crate::job::job_types::tts::styletts2::process_job_create_voice::join_paths; + + #[test] + fn test_path_build() { + let paths = vec![ + PathBuf::from("/home/tensor/code/TTSDockerContainer/Vall-E-mount/input/20.wav"), + PathBuf::from("/home/tensor/code/TTSDockerContainer/Vall-E-mount/input/21.wav") + ]; + + let value = join_paths(paths); + let expected = + "\"/home/tensor/code/TTSDockerContainer/Vall-E-mount/input/20.wav\" \"/home/tensor/code/TTSDockerContainer/Vall-E-mount/input/21.wav\""; + assert_eq!(value, expected); + } +} diff --git a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_inference_voice.rs b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_inference_voice.rs new file mode 100644 index 0000000000..cb8d476954 --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_inference_voice.rs @@ -0,0 +1,318 @@ +use std::fs::read_to_string; +use std::path::PathBuf; +use std::time::Instant; + +use anyhow::anyhow; +use log::{error, info, warn}; + +use buckets::private::zs_voices::bucket_directory::{ModelCategory, ModelType}; +use buckets::private::zs_voices::bucket_file_path::ZeroShotVoiceEmbeddingBucketPath; +use buckets::public::media_files::bucket_file_path::MediaFileBucketPath; +use cloud_storage::bucket_client::BucketClient; +use cloud_storage::bucket_path_unifier::BucketPathUnifier; +use enums::by_table::generic_inference_jobs::inference_result_type::InferenceResultType; +use filesys::file_size::file_size; +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 crate::job::job_loop::job_success_result::JobSuccessResult; +use crate::job::job_loop::job_success_result::ResultEntity; +use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; +use crate::job::job_types::tts::styletts2::process_job::StyleTTS2ProcessJobArgs; +use crate::job::job_types::tts::styletts2::styletts2_inference_command::InferenceArgs; + +const BUCKET_FILE_PREFIX: &str = "fakeyou_"; +const BUCKET_FILE_EXTENSION: &str = ".wav"; +const MIME_TYPE: &str = "audio/wav"; + + +pub async fn process_inference_voice( + args: StyleTTS2ProcessJobArgs<'_>, + voice_token: String +) -> Result { + let deps = args.job_dependencies; + let job = args.job; + let mysql_pool = &deps.db.mysql_pool; + + let model_dependencies = deps + .job + .job_specific_dependencies + .maybe_styletts2_dependencies + .as_ref() + .ok_or_else(|| ProcessSingleJobError::JobSystemMisconfiguration(Some("missing StyleTTS2 dependencies".to_string())))?; + + // get some globals + let mut job_progress_reporter = deps + .clients + .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"))); + } + }; + + let voice_token = tokens::tokens::zs_voices::ZsVoiceToken(voice_token); + + // Get voice bucket hash - from voice token + let voice_lookup_result = get_voice_by_token(&voice_token, false, &mysql_pool).await; + + let voice = match voice_lookup_result { + Ok(Some(voice)) => voice, + Ok(None) => { + warn!("Voice not found: {:?}", voice_token); + return Err(ProcessSingleJobError::Other(anyhow!("Voice not found: {:?}", voice_token))); + } + Err(err) => { + warn!("Error looking up voice: {:?}", err); + return Err(ProcessSingleJobError::Other(anyhow!("Error looking up voice: {:?}", err))); + } + }; + + // Need to download the models + info!("Download models (if not present)..."); + + for downloader in model_dependencies.downloaders.all_downloaders() { + let result = downloader.download_if_not_on_filesystem( + &args.job_dependencies.buckets.private_bucket_client, + &args.job_dependencies.fs.scoped_temp_dir_creator_for_short_lived_downloads + ).await; + + if let Err(err) = result { + error!("could not download: {:?}", err); + return Err(err); + } + } + + // 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); + + // 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))?; + + let workdir = work_temp_dir.path().to_path_buf(); + let filename = "weights.npz".to_string(); + + let mut downloaded_weights_path = work_temp_dir.path().to_path_buf(); + downloaded_weights_path.push(&filename); + + let voice_file = download_voice_embedding( + &voice, + &deps.buckets.private_bucket_client, + &downloaded_weights_path + ).await?; + + println!("voice file path! {}", voice_file.filesystem_path.to_string_lossy()); + + // 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))?; + + let inference_start_time = Instant::now(); + + let output_file_name = String::from("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 + stderr_output_file: &stderr_output_file, + } + ); + + let inference_duration = Instant::now().duration_since(inference_start_time); + + if !command_exit_status.is_success() { + error!("Zero shot TTS failed: {:?}", command_exit_status); + + let error = ProcessSingleJobError::Other(anyhow!("CommandExitStatus: {:?}", command_exit_status)); + + 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); + //safe_delete_temp_file(&usable_image_path); + //safe_delete_temp_file(&output_video_fs_path); + //safe_delete_temp_file(&stderr_output_file); + //safe_delete_temp_directory(&work_temp_dir); + + return Err(error); + } + + // 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(); + + // Finished file path + let mut finished_file = work_temp_dir.path().to_path_buf(); + finished_file.push(&output_file_name); + + info!("Upload Bucket Path: {:?}", result_bucket_object_pathbuf); + info!("Upload File Path: {:?}", finished_file); + + args.job_dependencies.buckets.public_bucket_client + .upload_filename_with_content_type( + &result_bucket_object_pathbuf, + &finished_file, + &MIME_TYPE + ) + .await + .map_err(|e| ProcessSingleJobError::Other(e))?; + + // ==================== UPLOAD AUDIO TO BUCKET ==================== + + info!("Calculating sha256..."); + + let file_checksum = sha256_hash_file(&finished_file).map_err(|err| { + ProcessSingleJobError::Other(anyhow!("Error hashing file: {:?}", err)) + })?; + + let file_size_bytes = file_size(&finished_file).map_err(|err| + ProcessSingleJobError::Other(err) + )?; + + 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_zero_shot(InsertArgs { + pool: &args.job_dependencies.db.mysql_pool, + job: &job, + maybe_mime_type: Some(&MIME_TYPE), + 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.job.info.container.is_on_prem, + worker_hostname: &args.job_dependencies.job.info.container.hostname, + worker_cluster: &args.job_dependencies.job.info.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 + ); + + Ok(JobSuccessResult { + maybe_result_entity: Some(ResultEntity { + entity_type: InferenceResultType::MediaFile, + entity_token: media_file_token.to_string(), + }), + inference_duration, + }) +} + +pub struct VoiceFile { + pub filesystem_path: PathBuf, +} + +pub async fn download_voice_embedding_from_hash( + bucket_hash: &str, + name: &str, + private_bucket_client: &BucketClient, + path: &PathBuf +) -> Result { + let unifer = BucketPathUnifier::default_paths(); + let object_path = unifer.zero_shot_tts_speaker_encoding(bucket_hash, 0); + + let mut path = path.clone(); + + let file_name = format!("{}", name); + path.push(&file_name); + + private_bucket_client.download_file_to_disk(object_path, &path) + .await + .map_err(|err| { + error!("Could not download embedding file: {err}"); + ProcessSingleJobError::from_anyhow_error(anyhow!("Could not download embedding file: {err}")) + })?; + + let voice_file = VoiceFile { + filesystem_path: PathBuf::from(&path.clone()), + }; + + Ok(voice_file) +} + +pub async fn download_voice_embedding( + voice: &ZsVoice, + private_bucket_client: &BucketClient, + download_path: &PathBuf, +) -> Result { + + let embedding_bucket_location = ZeroShotVoiceEmbeddingBucketPath::from_object_hash( + &voice.bucket_hash, + ModelCategory::Tts, + ModelType::StyleTTS2, + voice.model_version + ); + + info!("Downloading embedding from: {:?}", &embedding_bucket_location.to_full_object_pathbuf()); + info!("Downloading to filesystem location: {:?}", download_path); + + let result = private_bucket_client.download_file_to_disk( + embedding_bucket_location.to_full_object_pathbuf(), + &download_path + ).await; + + if let Err(err) = result { + error!("could not download embedding file: {:?}", err); + return Err(ProcessSingleJobError::from_anyhow_error(err)); + } + + let voice_file = VoiceFile { + filesystem_path: PathBuf::from(&download_path.clone()), + }; + + Ok(voice_file) +} + diff --git a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/styletts2_dependencies.rs b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/styletts2_dependencies.rs new file mode 100644 index 0000000000..eb96c87c7a --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/styletts2_dependencies.rs @@ -0,0 +1,20 @@ +use errors::AnyhowResult; + +use crate::job::job_types::tts::styletts2::model_downloaders::StyleTTS2Downloaders; +use crate::job::job_types::tts::styletts2::styletts2_inference_command::{StyleTTS2CreateEmbeddingCommand, StyleTTS2InferenceCommand}; + +pub struct StyleTTS2Dependencies { + pub create_embedding_command: StyleTTS2CreateEmbeddingCommand, + pub downloaders: StyleTTS2Downloaders, + pub inference_command: StyleTTS2InferenceCommand, +} + +impl StyleTTS2Dependencies { + pub fn setup() -> AnyhowResult { + Ok(Self { + create_embedding_command: StyleTTS2CreateEmbeddingCommand::from_env()?, + downloaders: StyleTTS2Downloaders::build_all_from_env(), + inference_command: StyleTTS2InferenceCommand::from_env()?, + }) + } +} diff --git a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/styletts2_inference_command.rs b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/styletts2_inference_command.rs new file mode 100644 index 0000000000..f23fe05c09 --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/styletts2_inference_command.rs @@ -0,0 +1,549 @@ +use std::collections::HashSet; +use std::env; +use std::ffi::OsString; +use std::fs::File; +use std::path::{Path, PathBuf}; +use std::time::Duration; + +use anyhow::anyhow; +use log::info; +use once_cell::sync::Lazy; +use subprocess::{Popen, PopenConfig, Redirection}; + +use errors::AnyhowResult; +use filesys::path_to_string::path_to_string; +use subprocess_common::docker_options::{DockerFilesystemMount, DockerGpu, DockerOptions}; + +use crate::job::job_loop::command_exit_status::CommandExitStatus; + +#[derive(Clone)] +pub enum ExecutableOrCommand { + /// Eg. `inference.py` + Executable(PathBuf), + /// Eg. `python3 inference.py` + Command(String), +} + +// These environment vars are not copied over to the subprocess +// TODO/FIXME(bt, 2023-05-28): This is horrific security! +static IGNORED_ENVIRONMENT_VARS: Lazy> = Lazy::new(|| { + let env_var_names = ["MYSQL_URL", "ACCESS_KEY", "SECRET_KEY", "NEWRELIC_API_KEY"]; + + env_var_names + .iter() + .map(|value| value.to_string()) + .collect::>() +}); + +// The inference command assumes everything has been setup to run inference +// so we have all the files downloaded for the inputs. +pub struct InferenceArgs> { + /// --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 stderr_output_file: P, +} + +pub struct CreateVoiceInferenceArgs> { + pub output_embedding_path: P, + pub output_embedding_name: P, + pub audio_files: String, + pub stderr_output_file: P, +} + +#[derive(Clone)] +pub struct StyleTTS2InferenceCommand { + /// Where the code lives + root_code_directory: PathBuf, + /// A single executable script or a much larger bash command. + executable_or_command: ExecutableOrCommand, + /// eg. `source python/bin/activate` + maybe_virtual_env_activation_command: Option, + /// Optional default config file to use + maybe_default_config_path: Option, + /// If this is run under Docker (eg. in development), these are the options. + maybe_docker_options: Option, + /// If the execution should be ended after a certain point. + maybe_execution_timeout: Option, + /// Inference arg. + /// --checkpoint_dir: optional location for checkpoints directory + pub alternate_checkpoint_dir: Option, +} + +impl StyleTTS2InferenceCommand { + pub fn new>( + root_code_directory: P, + executable_or_command: ExecutableOrCommand, + maybe_virtual_env_activation_command: Option<&str>, + maybe_default_config_path: Option

, + maybe_docker_options: Option, + maybe_execution_timeout: Option, + alternate_checkpoint_dir: Option + ) -> Self { + Self { + root_code_directory: root_code_directory.as_ref().to_path_buf(), + executable_or_command, + maybe_virtual_env_activation_command: maybe_virtual_env_activation_command.map(|s| + s.to_string() + ), + maybe_default_config_path: maybe_default_config_path.map(|p| p.as_ref().to_path_buf()), + maybe_docker_options, + maybe_execution_timeout, + alternate_checkpoint_dir, + } + } + + pub fn from_env() -> AnyhowResult { + let root_code_directory = easyenv::get_env_pathbuf_required( + "STYLE_TTS2_INFERENCE_ROOT_DIRECTORY" + )?; + + let maybe_inference_command = easyenv::get_env_string_optional( + "STYLE_TTS2_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( + "STYLE_TTS2_INFERENCE_EXECUTABLE" + ); + + let executable_or_command = match maybe_inference_command { + Some(command) => ExecutableOrCommand::Command(command), + None => + match maybe_inference_executable { + Some(executable) => ExecutableOrCommand::Executable(executable), + None => { + return Err(anyhow!("neither command nor executable passed")); + } + } + }; + + let maybe_virtual_env_activation_command = easyenv::get_env_string_optional( + "STYLE_TTS2_INFERENCE_MAYBE_VENV_COMMAND" + ); + + let maybe_default_config_path = easyenv::get_env_pathbuf_optional( + "STYLE_TTS2_INFERENCE_MAYBE_DEFAULT_CONFIG_PATH" + ); + + let maybe_execution_timeout = easyenv::get_env_duration_seconds_optional("TIMEOUT_SECONDS"); + + // Probably for local + let maybe_docker_options = easyenv + ::get_env_string_optional("STYLE_TTS2_INFERENCE_MAYBE_DOCKER_IMAGE") + .map(|image_name| { + DockerOptions { + image_name, + maybe_bind_mount: Some(DockerFilesystemMount::tmp_to_tmp()), + maybe_environment_variables: None, + maybe_gpu: Some(DockerGpu::All), + } + }); + + // Override for --checkpoint_dir at inference time + let alternate_checkpoint_dir = easyenv::get_env_pathbuf_optional( + "STYLE_TTS2_ALTERNATE_CHECKPOINT_PATH" + ); + + Ok(Self { + root_code_directory, + executable_or_command, + maybe_virtual_env_activation_command, + maybe_default_config_path, + maybe_docker_options, + maybe_execution_timeout, + alternate_checkpoint_dir, + }) + } + + pub fn execute_inference>(&self, args: InferenceArgs

) -> CommandExitStatus { + match self.do_execute_inference(args) { + Ok(exit_status) => exit_status, + Err(error) => + CommandExitStatus::FailureWithReason { reason: format!("error: {:?}", error) }, + } + } + + fn do_execute_inference>( + &self, + args: InferenceArgs

+ ) -> AnyhowResult { + let mut command = String::new(); + command.push_str(&format!("cd {}", path_to_string(&self.root_code_directory))); + + if let Some(venv_command) = self.maybe_virtual_env_activation_command.as_deref() { + command.push_str(" && "); + command.push_str(venv_command); + command.push_str(" "); + } + + command.push_str(" && "); + + match self.executable_or_command { + ExecutableOrCommand::Executable(ref executable) => { + command.push_str(&path_to_string(executable)); + command.push_str(" "); + } + ExecutableOrCommand::Command(ref cmd) => { + command.push_str(cmd); + command.push_str(" "); + } + } + + // ===== Begin Python Args ===== + + command.push_str(" --text "); + command.push_str(&format!("\"{}\"", &args.input_text)); + + command.push_str(" --prompt-name "); + command.push_str(&path_to_string(args.input_embedding_name)); + + command.push_str(" --prompt-path "); + command.push_str(&path_to_string(&args.input_embedding_path)); + + command.push_str(" --audio-name "); + 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/"))); + + // ===== End Python Args ===== + + if let Some(docker_options) = self.maybe_docker_options.as_ref() { + command = docker_options.to_command_string(&command); + } + + info!("Command: {:?}", command); + + let command_parts = ["bash", "-c", &command]; + + let mut env_vars = Vec::new(); + + // Copy all environment variables from the parent process. + // This is necessary to send all the kubernetes settings for Nvidia / CUDA. + for (env_key, env_value) in env::vars() { + if IGNORED_ENVIRONMENT_VARS.contains(&env_key) { + continue; + } + env_vars.push((OsString::from(env_key), OsString::from(env_value))); + } + + let mut config = PopenConfig::default(); + + info!("stderr will be written to file: {:?}", args.stderr_output_file.as_ref()); + + let stderr_file = File::create(&args.stderr_output_file)?; + config.stderr = Redirection::File(stderr_file); + + if !env_vars.is_empty() { + config.env = Some(env_vars); + } + + let mut p = Popen::create(&command_parts, config)?; + + info!("Subprocess PID: {:?}", p.pid()); + + match self.maybe_execution_timeout { + None => { + let exit_status = p.wait()?; + info!("Subprocess exit status: {:?}", exit_status); + Ok(CommandExitStatus::from_exit_status(exit_status)) + } + Some(timeout) => { + info!("Executing with timeout: {:?}", &timeout); + let exit_status = p.wait_timeout(timeout.clone())?; + + match exit_status { + None => { + // NB: If the program didn't successfully terminate, kill it. + info!( + "Subprocess didn't end after timeout: {:?}; terminating...", + &timeout + ); + let _r = p.terminate()?; + Ok(CommandExitStatus::Timeout) + } + Some(exit_status) => { + info!("Subprocess timed wait exit status: {:?}", exit_status); + Ok(CommandExitStatus::from_exit_status(exit_status)) + } + } + } + } + } +} + +#[derive(Clone)] +pub struct StyleTTS2CreateEmbeddingCommand { + /// Where the code lives + root_code_directory: PathBuf, + /// A single executable script or a much larger bash command. + executable_or_command: ExecutableOrCommand, + /// eg. `source python/bin/activate` + maybe_virtual_env_activation_command: Option, + /// Optional default config file to use + maybe_default_config_path: Option, + /// If this is run under Docker (eg. in development), these are the options. + maybe_docker_options: Option, + /// If the execution should be ended after a certain point. + maybe_execution_timeout: Option, + /// Inference arg. + /// --checkpoint_dir: optional location for checkpoints directory + pub alternate_checkpoint_dir: Option, +} + +impl StyleTTS2CreateEmbeddingCommand { + pub fn new>( + root_code_directory: P, + executable_or_command: ExecutableOrCommand, + maybe_virtual_env_activation_command: Option<&str>, + maybe_default_config_path: Option

, + maybe_docker_options: Option, + maybe_execution_timeout: Option, + alternate_checkpoint_dir: Option + ) -> Self { + Self { + root_code_directory: root_code_directory.as_ref().to_path_buf(), + executable_or_command, + maybe_virtual_env_activation_command: maybe_virtual_env_activation_command.map(|s| + s.to_string() + ), + maybe_default_config_path: maybe_default_config_path.map(|p| p.as_ref().to_path_buf()), + maybe_docker_options, + maybe_execution_timeout, + alternate_checkpoint_dir, + } + } + + // helper funciton to set up env + pub fn from_env() -> AnyhowResult { + // TODO: fix for now just for compilation + let root_code_directory = easyenv::get_env_pathbuf_required( + "STYLE_TTS2_INFERENCE_ROOT_DIRECTORY" + )?; + + let maybe_inference_command = easyenv::get_env_string_optional( + "STYLE_TTS2_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( + "STYLE_TTS2_INFERENCE_EXECUTABLE" + ); + + let executable_or_command = match maybe_inference_command { + Some(command) => ExecutableOrCommand::Command(command), + None => + match maybe_inference_executable { + Some(executable) => ExecutableOrCommand::Executable(executable), + None => { + return Err(anyhow!("neither command nor executable passed")); + } + } + }; + + let maybe_virtual_env_activation_command = easyenv::get_env_string_optional( + "STYLE_TTS2_INFERENCE_MAYBE_VENV_COMMAND" + ); + + let maybe_default_config_path = easyenv::get_env_pathbuf_optional( + "STYLE_TTS2_INFERENCE_MAYBE_DEFAULT_CONFIG_PATH" + ); + + let maybe_execution_timeout = easyenv::get_env_duration_seconds_optional("TIMEOUT_SECONDS"); + + // Probably for local + let maybe_docker_options = easyenv + ::get_env_string_optional("STYLE_TTS2_INFERENCE_MAYBE_DOCKER_IMAGE") + .map(|image_name| { + DockerOptions { + image_name, + maybe_bind_mount: Some(DockerFilesystemMount::tmp_to_tmp()), + maybe_environment_variables: None, + maybe_gpu: Some(DockerGpu::All), + } + }); + + // Override for --checkpoint_dir at inference time + let alternate_checkpoint_dir = easyenv::get_env_pathbuf_optional( + "STYLE_TTS2_CHECKPOINT_PATH" + ); + + Ok(Self { + root_code_directory, + executable_or_command, + maybe_virtual_env_activation_command, + maybe_default_config_path, + maybe_docker_options, + maybe_execution_timeout, + alternate_checkpoint_dir, + }) + } + + pub fn execute_inference>(&self, args: CreateVoiceInferenceArgs

) -> CommandExitStatus { + match self.do_execute_inference(args) { + Ok(exit_status) => exit_status, + Err(error) => + CommandExitStatus::FailureWithReason { reason: format!("error: {:?}", error) }, + } + } + + fn do_execute_inference>( + &self, + args: CreateVoiceInferenceArgs

+ ) -> AnyhowResult { + let mut command = String::new(); + command.push_str(&format!("cd {}", path_to_string(&self.root_code_directory))); + + if let Some(venv_command) = self.maybe_virtual_env_activation_command.as_deref() { + command.push_str(" && "); + command.push_str(venv_command); + command.push_str(" "); + } + + command.push_str(" && "); + + match self.executable_or_command { + ExecutableOrCommand::Executable(ref executable) => { + command.push_str(&path_to_string(executable)); + command.push_str(" "); + } + ExecutableOrCommand::Command(ref cmd) => { + command.push_str(cmd); + command.push_str(" "); + } + } + + // ===== 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(&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/"))); + + // ===== End Python Args ===== + + if let Some(docker_options) = self.maybe_docker_options.as_ref() { + command = docker_options.to_command_string(&command); + } + + info!("Command: {:?}", command); + + let command_parts = ["bash", "-c", &command]; + + let mut env_vars = Vec::new(); + + // Copy all environment variables from the parent process. + // This is necessary to send all the kubernetes settings for Nvidia / CUDA. + for (env_key, env_value) in env::vars() { + if IGNORED_ENVIRONMENT_VARS.contains(&env_key) { + continue; + } + env_vars.push((OsString::from(env_key), OsString::from(env_value))); + } + + let mut config = PopenConfig::default(); + + info!("stderr will be written to file: {:?}", args.stderr_output_file.as_ref()); + + let stderr_file = File::create(&args.stderr_output_file)?; + config.stderr = Redirection::File(stderr_file); + + if !env_vars.is_empty() { + config.env = Some(env_vars); + } + + let mut p = Popen::create(&command_parts, config)?; + + info!("Subprocess PID: {:?}", p.pid()); + + match self.maybe_execution_timeout { + None => { + let exit_status = p.wait()?; + info!("Subprocess exit status: {:?}", exit_status); + Ok(CommandExitStatus::from_exit_status(exit_status)) + } + Some(timeout) => { + info!("Executing with timeout: {:?}", &timeout); + let exit_status = p.wait_timeout(timeout.clone())?; + + match exit_status { + None => { + // NB: If the program didn't successfully terminate, kill it. + info!( + "Subprocess didn't end after timeout: {:?}; terminating...", + &timeout + ); + let _r = p.terminate()?; + Ok(CommandExitStatus::Timeout) + } + Some(exit_status) => { + info!("Subprocess timed wait exit status: {:?}", exit_status); + Ok(CommandExitStatus::from_exit_status(exit_status)) + } + } + } + } + } + +} diff --git a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/validate_job.rs b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/validate_job.rs new file mode 100644 index 0000000000..41b76feac1 --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/validate_job.rs @@ -0,0 +1,84 @@ +use anyhow::anyhow; + +use mysql_queries::payloads::generic_inference_args::generic_inference_args::{ + InferenceCategoryAbbreviated, + PolymorphicInferenceArgs, +}; +use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; + +use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; + +pub enum JobType { + Create, + Inference +} +// determines the job type voice inference or voice creation +pub struct JobArgs { + pub voice_token: Option, // infers + pub voice_dataset_token: Option, // creates + pub job_type: JobType +} + +pub fn validate_job(job: &AvailableInferenceJob) -> Result { + let inference_category = job.maybe_inference_args + .as_ref() + .map(|args| args.inference_category) + .flatten(); + + match inference_category { + Some(InferenceCategoryAbbreviated::TextToSpeech) => {} + Some(category) => { + return Err( + ProcessSingleJobError::from_anyhow_error( + anyhow!("wrong inference category for job: {:?}", category) + ) + ); + } + None => { + return Err( + ProcessSingleJobError::from_anyhow_error(anyhow!("no inference category for job!")) + ); + } + } + + let inference_args = job.maybe_inference_args + .as_ref() + .map(|args| args.args.as_ref()) + .flatten(); + + let inference_args = match inference_args { + Some(args) => args, + None => { + return Err( + ProcessSingleJobError::from_anyhow_error(anyhow!("no inference args for job!")) + ); + } + }; + + let ttsArgs = match inference_args { + PolymorphicInferenceArgs::Tts(inference_args) => inference_args, + _ => { + return Err( + ProcessSingleJobError::from_anyhow_error(anyhow!("wrong inner args for job!")) + ); + } + }; + + if let Some(voice_token) = &ttsArgs.voice_token { + return Ok::(JobArgs { + voice_token: Some(voice_token.clone()), + voice_dataset_token: None, + job_type: JobType::Inference + }); + } + + if let Some(dataset_token) = &ttsArgs.dataset_token { + return Ok::(JobArgs { + voice_token: None, + voice_dataset_token: Some(dataset_token.clone()), + job_type: JobType::Create + }); + } + + Err(ProcessSingleJobError::from_anyhow_error(anyhow!("Missing Voice Token!"))) +} diff --git a/crates/service/job/inference_job/src/job_specific_dependencies.rs b/crates/service/job/inference_job/src/job_specific_dependencies.rs index a17b4ef472..1263583493 100644 --- a/crates/service/job/inference_job/src/job_specific_dependencies.rs +++ b/crates/service/job/inference_job/src/job_specific_dependencies.rs @@ -2,6 +2,7 @@ use enums::by_table::generic_inference_jobs::inference_model_type::InferenceMode use errors::AnyhowResult; use crate::job::job_types::lipsync::sad_talker::sad_talker_dependencies::SadTalkerDependencies; +use crate::job::job_types::tts::styletts2::styletts2_dependencies::StyleTTS2Dependencies; use crate::job::job_types::tts::tacotron2_v2_early_fakeyou::tacotron2_dependencies::Tacotron2Dependencies; use crate::job::job_types::tts::vall_e_x::vall_e_x_dependencies::VallExDependencies; use crate::job::job_types::tts::vits::vits_dependencies::VitsDependencies; @@ -20,6 +21,7 @@ pub struct JobSpecificDependencies { pub maybe_vits_dependencies: Option, pub maybe_rerender_dependencies: Option, pub maybe_mocapnet_dependencies: Option, + pub maybe_styletts2_dependencies: Option, } impl JobSpecificDependencies { @@ -33,6 +35,7 @@ impl JobSpecificDependencies { let mut maybe_vits_dependencies = None; let mut maybe_rerender_dependencies = None; let mut maybe_mocapnet_dependencies = None; + let mut maybe_styletts2_dependencies = None; if scoped_execution.can_run_job(InferenceModelType::RvcV2) { print_with_space("Setting RVCv2 dependencies..."); @@ -59,6 +62,11 @@ impl JobSpecificDependencies { maybe_vall_e_x_dependencies = Some(VallExDependencies::setup()?); } + if scoped_execution.can_run_job(InferenceModelType::StyleTTS2) { + print_with_space("Setting StyleTTS2 dependencies..."); + maybe_styletts2_dependencies = Some(StyleTTS2Dependencies::setup()?); + } + if scoped_execution.can_run_job(InferenceModelType::Vits) { print_with_space("Setting Vits dependencies..."); maybe_vits_dependencies = Some(VitsDependencies::setup()?); @@ -82,7 +90,8 @@ impl JobSpecificDependencies { maybe_vall_e_x_dependencies, maybe_vits_dependencies, maybe_rerender_dependencies, - maybe_mocapnet_dependencies + maybe_mocapnet_dependencies, + maybe_styletts2_dependencies, }) } } From ea3a6d3c0b6ce7a9914b656bda2e215cbf3508eb Mon Sep 17 00:00:00 2001 From: Kasisnu <5044120+kasisnu@users.noreply.github.com> Date: Fri, 12 Jan 2024 17:58:12 +0530 Subject: [PATCH 2/3] pass appropriate arguments to styletts inference command --- .../tts/styletts2/model_downloaders.rs | 15 +--- .../tts/styletts2/process_job_create_voice.rs | 45 +++------- .../styletts2/process_job_inference_voice.rs | 44 +++------- .../styletts2/styletts2_inference_command.rs | 88 ++++--------------- 4 files changed, 45 insertions(+), 147 deletions(-) diff --git a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/model_downloaders.rs b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/model_downloaders.rs index 18bcc1c931..c50a6dcfff 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/model_downloaders.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/model_downloaders.rs @@ -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![] } } diff --git a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_create_voice.rs b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_create_voice.rs index d64cc4f809..3c890b3cea 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_create_voice.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_create_voice.rs @@ -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 = 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); diff --git a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_inference_voice.rs b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_inference_voice.rs index cb8d476954..b7fcd141de 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_inference_voice.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_inference_voice.rs @@ -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); diff --git a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/styletts2_inference_command.rs b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/styletts2_inference_command.rs index f23fe05c09..bda3b67f5d 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/styletts2_inference_command.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/styletts2_inference_command.rs @@ -39,17 +39,15 @@ 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_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> { 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 ===== From 5e98049c501eafae0caf91ed1d9f6540c6c037c7 Mon Sep 17 00:00:00 2001 From: Kasisnu <5044120+kasisnu@users.noreply.github.com> Date: Mon, 15 Jan 2024 17:44:29 +0530 Subject: [PATCH 3/3] switch from vallex -> styletts2 for zero shot --- .../media_files/media_file_origin_model_type.rs | 11 ++++++++++- .../create/insert_media_file_from_zero_shot_tts.rs | 6 +++--- .../src/job/job_types/tts/process_single_tts_job.rs | 11 +++++++++-- .../tts/styletts2/process_job_inference_voice.rs | 4 +++- .../tts/styletts2/styletts2_inference_command.rs | 6 ++++++ .../tts/vall_e_x/process_job_inference_voice.rs | 2 ++ .../voice_designer/inference/enqueue_tts_request.rs | 2 +- .../endpoints/voice_designer/voices/create_voice.rs | 2 +- 8 files changed, 35 insertions(+), 9 deletions(-) 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 d9fbd7cdd5..1a3a9e58f8 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 @@ -38,6 +38,9 @@ pub enum MediaFileOriginModelType { #[serde(rename = "mocap_net")] MocapNet, + + #[serde(rename = "styletts2")] + StyleTTS2, } // TODO(bt, 2022-12-21): This desperately needs MySQL integration tests! @@ -56,6 +59,7 @@ impl MediaFileOriginModelType { Self::VallEX => "vall_e_x", Self::Rerender => "rerender", Self::MocapNet => "mocap_net", + Self::StyleTTS2 => "styletts2", } } @@ -68,6 +72,7 @@ impl MediaFileOriginModelType { "vall_e_x" => Ok(Self::VallEX), "rerender" => Ok(Self::Rerender), "mocap_net" => Ok(Self::MocapNet), + "styletts2" => Ok(Self::StyleTTS2), _ => Err(format!("invalid value: {:?}", value)), } } @@ -83,6 +88,7 @@ impl MediaFileOriginModelType { Self::VallEX, Self::Rerender, Self::MocapNet, + Self::StyleTTS2, ]) } } @@ -109,6 +115,7 @@ mod tests { assert_eq!(MediaFileOriginModelType::VallEX.to_str(), "vall_e_x"); assert_eq!(MediaFileOriginModelType::Rerender.to_str(), "rerender"); assert_eq!(MediaFileOriginModelType::MocapNet.to_str(), "mocap_net"); + assert_eq!(MediaFileOriginModelType::StyleTTS2.to_str(), "styletts2"); } #[test] @@ -120,13 +127,14 @@ mod tests { assert_eq!(MediaFileOriginModelType::from_str("vall_e_x").unwrap(), MediaFileOriginModelType::VallEX); assert_eq!(MediaFileOriginModelType::from_str("rerender").unwrap(), MediaFileOriginModelType::Rerender); assert_eq!(MediaFileOriginModelType::from_str("mocap_net").unwrap(), MediaFileOriginModelType::MocapNet); + assert_eq!(MediaFileOriginModelType::from_str("styletts2").unwrap(), MediaFileOriginModelType::StyleTTS2); assert!(MediaFileOriginModelType::from_str("foo").is_err()); } #[test] fn all_variants() { let mut variants = MediaFileOriginModelType::all_variants(); - assert_eq!(variants.len(), 7); + assert_eq!(variants.len(), 8); assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::RvcV2)); assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::SadTalker)); assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::SoVitsSvc)); @@ -134,6 +142,7 @@ mod tests { assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::VallEX)); assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::Rerender)); assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::MocapNet)); + assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::StyleTTS2)); assert_eq!(variants.pop_first(), None); } } diff --git a/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_from_zero_shot_tts.rs b/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_from_zero_shot_tts.rs index e250919752..ff73ed35d5 100644 --- a/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_from_zero_shot_tts.rs +++ b/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_from_zero_shot_tts.rs @@ -18,6 +18,7 @@ pub struct InsertArgs<'a> { pub pool: &'a MySqlPool, pub job: &'a AvailableInferenceJob, + pub origin_model_type: MediaFileOriginModelType, pub maybe_mime_type: Option<&'a str>, pub file_size_bytes: u64, pub sha256_checksum: &'a str, @@ -69,7 +70,6 @@ pub async fn insert_media_file_from_zero_shot( // Since fill this out apart of the funtion this seems like the only thing to change. const ORIGIN_CATEGORY : MediaFileOriginCategory = MediaFileOriginCategory::Inference; const ORIGIN_PRODUCT_CATEGORY : MediaFileOriginProductCategory = MediaFileOriginProductCategory::ZeroShotVoice; - const ORIGIN_MODEL_TYPE : MediaFileOriginModelType = MediaFileOriginModelType::VallEX; const MEDIA_TYPE : MediaFileType = MediaFileType::Audio; let record_id = { @@ -111,8 +111,8 @@ pub async fn insert_media_file_from_zero_shot( ORIGIN_CATEGORY.to_str(), ORIGIN_PRODUCT_CATEGORY.to_str(), - ORIGIN_MODEL_TYPE.to_str(), - + args.origin_model_type.to_str(), + MEDIA_TYPE.to_str(), args.maybe_mime_type, args.file_size_bytes, diff --git a/crates/service/job/inference_job/src/job/job_types/tts/process_single_tts_job.rs b/crates/service/job/inference_job/src/job/job_types/tts/process_single_tts_job.rs index d9a6cf1f53..495901d1dc 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/process_single_tts_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/process_single_tts_job.rs @@ -8,7 +8,8 @@ use mysql_queries::queries::tts::tts_models::get_tts_model_for_inference_improve use crate::job::job_loop::job_success_result::JobSuccessResult; use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; -use crate::job::job_types::tts::{tacotron2_v2_early_fakeyou, vits}; +use crate::job::job_types::tts::{styletts2, tacotron2_v2_early_fakeyou, vits}; +use crate::job::job_types::tts::styletts2::process_job::StyleTTS2ProcessJobArgs; use crate::job::job_types::tts::tacotron2_v2_early_fakeyou::process_job::ProcessJobArgs; use crate::job::job_types::tts::vall_e_x::process_job::VALLEXProcessJobArgs; use crate::job::job_types::tts::vits::process_job::VitsProcessJobArgs; @@ -24,7 +25,7 @@ pub async fn process_single_tts_job( // TODO: Move common checks for slurs, etc. here. match job.maybe_model_type { - Some(InferenceModelType::VallEX) => { + Some(InferenceModelType::VallEX) | Some(InferenceModelType::StyleTTS2) => { // Zero-shot TTS does not need a fine-tuned model token. dispatch_zero_shot_model( job_dependencies, @@ -64,6 +65,12 @@ async fn dispatch_zero_shot_model( job, }).await } + Some(InferenceModelType::StyleTTS2) => { + styletts2::process_job::process_job(StyleTTS2ProcessJobArgs { + job_dependencies, + job, + }).await + } _ => { Err(ProcessSingleJobError::Other(anyhow!("not a zero-shot model"))) } diff --git a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_inference_voice.rs b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_inference_voice.rs index b7fcd141de..b602d04fbf 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_inference_voice.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_inference_voice.rs @@ -11,6 +11,7 @@ use buckets::public::media_files::bucket_file_path::MediaFileBucketPath; use cloud_storage::bucket_client::BucketClient; use cloud_storage::bucket_path_unifier::BucketPathUnifier; use enums::by_table::generic_inference_jobs::inference_result_type::InferenceResultType; +use enums::by_table::media_files::media_file_origin_model_type::MediaFileOriginModelType; use filesys::file_size::file_size; 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; @@ -135,7 +136,7 @@ pub async fn process_inference_voice( let command_exit_status = model_dependencies.inference_command.execute_inference( InferenceArgs { - input_embedding_file_path: &workdir, + input_embedding_file_path: &downloaded_weights_path, 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, @@ -208,6 +209,7 @@ pub async fn process_inference_voice( let (media_file_token, id) = insert_media_file_from_zero_shot(InsertArgs { pool: &args.job_dependencies.db.mysql_pool, job: &job, + origin_model_type: MediaFileOriginModelType::StyleTTS2, maybe_mime_type: Some(&MIME_TYPE), file_size_bytes, sha256_checksum: &file_checksum, diff --git a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/styletts2_inference_command.rs b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/styletts2_inference_command.rs index bda3b67f5d..814cf97939 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/styletts2_inference_command.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/styletts2_inference_command.rs @@ -423,6 +423,12 @@ impl StyleTTS2CreateEmbeddingCommand { command.push_str(" --voice "); command.push_str(&args.audio_file); + command.push_str(" --vcsteps "); + command.push_str(format!("{}", 20).as_str()); + + command.push_str(" --input "); + command.push_str("/dev/null"); + // ===== End Python Args ===== if let Some(docker_options) = self.maybe_docker_options.as_ref() { diff --git a/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/process_job_inference_voice.rs b/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/process_job_inference_voice.rs index 72272f4fab..7646fd7652 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/process_job_inference_voice.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/process_job_inference_voice.rs @@ -11,6 +11,7 @@ use buckets::public::media_files::bucket_file_path::MediaFileBucketPath; use cloud_storage::bucket_client::BucketClient; use cloud_storage::bucket_path_unifier::BucketPathUnifier; use enums::by_table::generic_inference_jobs::inference_result_type::InferenceResultType; +use enums::by_table::media_files::media_file_origin_model_type::MediaFileOriginModelType; use filesys::file_size::file_size; 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; @@ -224,6 +225,7 @@ pub async fn process_inference_voice( let (media_file_token, id) = insert_media_file_from_zero_shot(InsertArgs { pool: &args.job_dependencies.db.mysql_pool, job: &job, + origin_model_type: MediaFileOriginModelType::VallEX, maybe_mime_type: Some(&MIME_TYPE), file_size_bytes, sha256_checksum: &file_checksum, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/inference/enqueue_tts_request.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/inference/enqueue_tts_request.rs index 7e09f40782..5a8e7bb3f0 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/inference/enqueue_tts_request.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/inference/enqueue_tts_request.rs @@ -189,7 +189,7 @@ pub async fn enqueue_tts_request( let query_result = insert_generic_inference_job(InsertGenericInferenceArgs { uuid_idempotency_token: &request.uuid_idempotency_token, inference_category: InferenceCategory::TextToSpeech, - maybe_model_type: Some(InferenceModelType::VallEX), // NB: Model is static during inference + maybe_model_type: Some(InferenceModelType::StyleTTS2), // NB: Model is static during inference maybe_model_token: None, // NB: Model is static during inference maybe_input_source_token: None, maybe_input_source_token_type: None, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voices/create_voice.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voices/create_voice.rs index 6654843007..f8dd1274b9 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voices/create_voice.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voices/create_voice.rs @@ -177,7 +177,7 @@ pub async fn create_voice_handler( let query_result = insert_generic_inference_job(InsertGenericInferenceArgs { uuid_idempotency_token: &request.uuid_idempotency_token, inference_category: InferenceCategory::TextToSpeech, - maybe_model_type: Some(InferenceModelType::VallEX), // NB: Model is static during inference + maybe_model_type: Some(InferenceModelType::StyleTTS2), // NB: Model is static during inference maybe_model_token: None, // NB: Model is static during inference maybe_input_source_token: None, maybe_input_source_token_type: None,