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 =====