diff --git a/crates/lib/api_clients/cloud_storage/src/bucket_path_unifier.rs b/crates/lib/api_clients/cloud_storage/src/bucket_path_unifier.rs index 91de4a2131..76374d5861 100644 --- a/crates/lib/api_clients/cloud_storage/src/bucket_path_unifier.rs +++ b/crates/lib/api_clients/cloud_storage/src/bucket_path_unifier.rs @@ -213,7 +213,6 @@ impl BucketPathUnifier { .join(speaker_encoding_filename) } - // ==================== VOICE CONVERSION MODELS ==================== // // NB: Entropic hash is not based on file hash and is shared between .pth and .index files. diff --git a/crates/lib/mysql_queries/src/queries/voice_designer/voice_samples/list_dataset_samples_for_dataset_token.rs b/crates/lib/mysql_queries/src/queries/voice_designer/voice_samples/list_dataset_samples_for_dataset_token.rs index 7548bea6c6..7e34a993ad 100644 --- a/crates/lib/mysql_queries/src/queries/voice_designer/voice_samples/list_dataset_samples_for_dataset_token.rs +++ b/crates/lib/mysql_queries/src/queries/voice_designer/voice_samples/list_dataset_samples_for_dataset_token.rs @@ -11,6 +11,8 @@ use tokens::tokens::users::UserToken; use tokens::tokens::zs_voice_dataset_samples::ZsVoiceDatasetSampleToken; use tokens::tokens::zs_voice_datasets::ZsVoiceDatasetToken; + + pub struct DatasetSampleRecordForList { pub sample_token: ZsVoiceDatasetSampleToken, diff --git a/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/process_job.rs b/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/process_job.rs index cb1b91a553..e8fa1b0a73 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/process_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/process_job.rs @@ -6,11 +6,16 @@ use filesys::file_size::file_size; use enums::by_table::generic_inference_jobs::inference_result_type::InferenceResultType; use log::{ error, info, warn }; +use mysql_queries::queries::media_files::get_media_file::MediaFile; use mysql_queries::queries::media_files::insert_media_file_from_zero_shot_tts::InsertArgs; use mysql_queries::queries::media_files::insert_media_file_from_zero_shot_tts::insert_media_file_from_face_zero_shot; use mysql_queries::queries::voice_designer::datasets::get_dataset::get_dataset_by_token; +use mysql_queries::queries::voice_designer::datasets::list_datasets::DatasetRecordForList; +use mysql_queries::queries::voice_designer::voice_samples::list_dataset_samples_for_dataset_token::DatasetSampleRecordForList; + use mysql_queries::queries::voice_designer::voice_samples::get_dataset_sample; use mysql_queries::queries::voice_designer::voice_samples::get_dataset_sample::get_dataset_sample_by_token; +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::get_voice::get_voice_by_token; use std::path::PathBuf; use std::time::Instant; @@ -22,6 +27,7 @@ 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::vall_e_x::validate_job::validate_job; +use crate::job::job_types::tts::vall_e_x::vall_e_x_inference_command::CreateVoiceArgs; use crate::job_dependencies::JobDependencies; use crate::job::job_types::tts::vall_e_x::vall_e_x_inference_command::InferenceArgs; use cloud_storage::bucket_path_unifier::BucketPathUnifier; @@ -35,6 +41,11 @@ pub struct VoiceFile { pub filesystem_path: PathBuf, } +pub struct AudioFile { + pub filesystem_path: PathBuf, +} + + pub async fn download_voice_embedding_from_hash( bucket_hash: &str, name: &str, @@ -44,9 +55,6 @@ pub async fn download_voice_embedding_from_hash( let unifer = BucketPathUnifier::default_paths(); let object_path = unifer.zero_shot_tts_speaker_encoding(bucket_hash, 0); - // todo replace with path ... - //let path: String = String::from("/tmp/downloads/zero_shot_tts/"); - // let filesystem_path = format!("{}{}_weights.npz",path, name); let mut path = path.clone(); let file_name = format!("{}", name); @@ -57,10 +65,35 @@ pub async fn download_voice_embedding_from_hash( let voice_file = VoiceFile { filesystem_path: PathBuf::from(&path.clone()), }; + Ok(voice_file) } +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); + + let result = private_bucket_client.download_file_to_disk(object_path, &path).await; + + let audio_file = AudioFile { + filesystem_path: PathBuf::from(&path.clone()), + }; + + Ok(audio_file) +} + + // This will download everything get into the root host OS then ... will invoke inference using the pathes from the files invoked pub struct VALLEXProcessJobArgs<'a> { pub job_dependencies: &'a JobDependencies, @@ -82,6 +115,74 @@ pub async fn process_create_voice( let voice_dataset_token = tokens::tokens::zs_voice_datasets::ZsVoiceDatasetToken(dataset_token); + // download data set files into 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(); + + let dataset = list_dataset_samples_for_dataset_token(&voice_dataset_token, false, &mysql_pool).await; + // unwrap this..... TODO + let downloaded_dataset:Vec = Vec::new(); + + for (index,record) in dataset.iter().enumerate() { + + let audio_media_file = MediaFileBucketPath::from_object_hash( + &record.public_bucket_directory_hash, + record.maybe_public_bucket_prefix, + record.maybe_public_bucket_extension); + + let file_name_wav = format!("{}.wav", index); + let file_path = PathBuf::new(); + file_path.push(workdir); + file_path.push(file_path); + + // TODO we might want to catch the error and not include the pathes into download dataset + deps.public_bucket_client.download_file_to_disk(audio_media_file.to_full_object_pathbuf(), file_path); + + downloaded_dataset.push(file_path); + } + + // Need to download the models + info!("Download models (if not present)..."); + + for downloader in deps.job_type_details.vall_e_x.downloaders.all_downloaders() { + let result = downloader.download_if_not_on_filesystem( + &args.job_dependencies.private_bucket_client, + &args.job_dependencies.fs.scoped_temp_dir_creator_for_downloads + ).await; + + if let Err(e) = result { + error!("could not download: {:?}", e); + return Err(ProcessSingleJobError::from_anyhow_error(e)); + } + } + + + 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"); + + // Run Inference + let command_exit_status = args.job_dependencies.job_type_details.vall_e_x.create_embedding_command.execute_inference( + CreateVoiceArgs { + output_embedding_path: &workdir, + output_embedding_name: file_name, + audio_files: String::from(""), // " " + stderr_output_file: String::from("zero_shot_create_voice.txt"), + } + ); + + let inference_duration = Instant::now().duration_since(inference_start_time); + + Err(ProcessSingleJobError::Other(anyhow!("Error"))) } @@ -177,7 +278,7 @@ pub async fn process_inference_voice( input_embedding_name: file_name, input_text: String::from(text), // text output_file_name: output_file_name.clone(), // output file name in the output folder - stderr_output_file: String::from("zero_shot.txt"), + stderr_output_file: String::from("zero_shot_inference.txt"), } ); let inference_duration = Instant::now().duration_since(inference_start_time); @@ -274,9 +375,25 @@ pub async fn process_job( } } } + +fn join_paths(paths: Vec)-> String { + paths.into_iter() + .map(|p| format!("\"{}\"",p.display())) + .collect::>() + .join(" ") +} mod tests { + + use super::*; #[test] fn test_length_zero() { - println!("Hello Test") + 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/vall_e_x/vall_e_x_inference_command.rs b/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/vall_e_x_inference_command.rs index 94231cc72e..ffb8c5550d 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/vall_e_x_inference_command.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/vall_e_x_inference_command.rs @@ -2,41 +2,37 @@ use std::collections::HashSet; use std::env; use std::ffi::OsString; use std::fs::File; -use std::path::{Path, PathBuf}; +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 subprocess::{ Popen, PopenConfig, Redirection }; use container_common::anyhow_result::AnyhowResult; use filesys::path_to_string::path_to_string; -use subprocess_common::docker_options::{DockerFilesystemMount, DockerGpu, DockerOptions}; +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), + /// 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", - ]; +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::>() + env_var_names + .iter() + .map(|value| value.to_string()) + .collect::>() }); // The inference command assumes everything has been setup to run inference @@ -46,28 +42,34 @@ pub struct InferenceArgs> { 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 - /// --result_file: path to final file output pub output_file_name: String, // output file name in the output folder - pub stderr_output_file: String - } + pub stderr_output_file: String, +} + +pub struct CreateVoiceArgs> { + pub output_embedding_path: P, + pub output_embedding_name: String, + pub audio_files: String, + pub stderr_output_file: String +} #[derive(Clone)] -pub struct VallEXInferenceCommand { - /// 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 +pub struct VallEXInferenceCommand { + /// 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 VallEXInferenceCommand { @@ -79,123 +81,129 @@ impl VallEXInferenceCommand { maybe_docker_options: Option, maybe_execution_timeout: Option, alternate_checkpoint_dir: Option - ) -> Self { + ) -> 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, + 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 { + pub fn from_env() -> AnyhowResult { let root_code_directory = easyenv::get_env_pathbuf_required( - "VALL_E_X_INFERENCE_ROOT_DIRECTORY")?; - + "VALL_E_X_INFERENCE_ROOT_DIRECTORY" + )?; + let maybe_inference_command = easyenv::get_env_string_optional( - "VALL_E_X_INFERENCE_COMMAND"); - - // EVENTUALLY should remove the check + "VALL_E_X_INFERENCE_COMMAND" + ); + + // EVENTUALLY should remove the check // Optional, eg. `./infer.py`. Typically we'll use the command form instead. let maybe_inference_executable = easyenv::get_env_pathbuf_optional( - "VALL_E_X_INFERENCE_EXECUTABLE"); - + "VALL_E_X_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")), - }, + 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( - "VALL_E_X_INFERENCE_MAYBE_VENV_COMMAND"); - + "VALL_E_X_INFERENCE_MAYBE_VENV_COMMAND" + ); + let maybe_default_config_path = easyenv::get_env_pathbuf_optional( - "VALL_E_X_INFERENCE_MAYBE_DEFAULT_CONFIG_PATH"); - - let maybe_execution_timeout = - easyenv::get_env_duration_seconds_optional("TIMEOUT_SECONDS"); - + "VALL_E_X_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( - "VALL_E_X_INFERENCE_MAYBE_DOCKER_IMAGE") + let maybe_docker_options = easyenv + ::get_env_string_optional("VALL_E_X_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), - } + 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( - "VALL_E_X_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, - }) - } + "VALL_E_X_ALTERNATE_CHECKPOINT_PATH" + ); - pub fn execute_inference>( - &self, - args: InferenceArgs

, - ) -> CommandExitStatus { + 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) }, + Ok(exit_status) => exit_status, + Err(error) => + CommandExitStatus::FailureWithReason { reason: format!("error: {:?}", error) }, } - } - + } + fn do_execute_inference>( &self, - args: InferenceArgs

, - ) -> AnyhowResult { - + 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(" && "); + 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(" "); - } + 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(&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)); @@ -209,185 +217,333 @@ impl VallEXInferenceCommand { 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( + &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( + &path_to_string(Path::new("/tmp/downloads/zero_shot_tts/vall-e-x_1.0/vocos-encodec/")) + ); command.push_str(" --vallex-path "); - command.push_str(&path_to_string(Path::new("/tmp/downloads/zero_shot_tts/vall-e-x_1.0/vallex-checkpoint.pt"))); + command.push_str( + &path_to_string( + Path::new("/tmp/downloads/zero_shot_tts/vall-e-x_1.0/vallex-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); + command = docker_options.to_command_string(&command); } - + info!("Command: {:?}", command); - - let command_parts = [ - "bash", - "-c", - &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), - )); + 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); - + 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); + 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); + 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 VallEXCreateEmbeddingCommand { - /// 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 +pub struct VallEXCreateEmbeddingCommand { + /// 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 VallEXCreateEmbeddingCommand { - - 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, +impl VallEXCreateEmbeddingCommand { + 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( - "VALL_E_X_INFERENCE_ROOT_DIRECTORY")?; - - let maybe_inference_command = easyenv::get_env_string_optional( - "VALL_E_X_INFERENCE_COMMAND"); - + // 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( + "VALL_E_X_INFERENCE_ROOT_DIRECTORY" + )?; + + let maybe_inference_command = easyenv::get_env_string_optional( + "VALL_E_X_INFERENCE_COMMAND" + ); + // EVENTUALLY should remove the check - // Optional, eg. `./infer.py`. Typically we'll use the command form instead. - let maybe_inference_executable = easyenv::get_env_pathbuf_optional( - "VALL_E_X_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( - "VALL_E_X_INFERENCE_MAYBE_VENV_COMMAND"); - - let maybe_default_config_path = easyenv::get_env_pathbuf_optional( - "VALL_E_X_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( - "VALL_E_X_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), + // Optional, eg. `./infer.py`. Typically we'll use the command form instead. + let maybe_inference_executable = easyenv::get_env_pathbuf_optional( + "VALL_E_X_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( + "VALL_E_X_INFERENCE_MAYBE_VENV_COMMAND" + ); + + let maybe_default_config_path = easyenv::get_env_pathbuf_optional( + "VALL_E_X_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("VALL_E_X_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( + "VALL_E_X_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: CreateVoiceArgs

) -> 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: CreateVoiceArgs

+ ) -> 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(" "); } - }); - - // Override for --checkpoint_dir at inference time - let alternate_checkpoint_dir = easyenv::get_env_pathbuf_optional( - "VALL_E_X_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, - }) + 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(" --vallex-path "); + command.push_str( + &path_to_string( + Path::new("/tmp/downloads/zero_shot_tts/vall-e-x_1.0/vallex-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); + + 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)) + } + } + } + } } } -