From 5e2f85964c9348ab43891a7fe808d057366b010e Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Tue, 21 Nov 2023 22:18:32 -0500 Subject: [PATCH] vits dependencies --- .../src/job/job_types/tts/vits/mod.rs | 1 + .../src/job/job_types/tts/vits/process_job.rs | 9 ++- .../job_types/tts/vits/vits_dependencies.rs | 15 +++++ .../tts/vits/vits_inference_command.rs | 60 ++++++++++++++++- .../job/inference_job/src/job_dependencies.rs | 13 ---- .../src/job_specific_dependencies.rs | 9 +++ crates/service/job/inference_job/src/main.rs | 65 +------------------ 7 files changed, 93 insertions(+), 79 deletions(-) create mode 100644 crates/service/job/inference_job/src/job/job_types/tts/vits/vits_dependencies.rs diff --git a/crates/service/job/inference_job/src/job/job_types/tts/vits/mod.rs b/crates/service/job/inference_job/src/job/job_types/tts/vits/mod.rs index 871f8562c0..ff8287b121 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/vits/mod.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/vits/mod.rs @@ -1,3 +1,4 @@ pub mod process_job; pub mod read_metadata_file; +pub mod vits_dependencies; pub mod vits_inference_command; diff --git a/crates/service/job/inference_job/src/job/job_types/tts/vits/process_job.rs b/crates/service/job/inference_job/src/job/job_types/tts/vits/process_job.rs index 607add18db..a09092464b 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/vits/process_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/vits/process_job.rs @@ -46,6 +46,13 @@ pub async fn process_job(args: VitsProcessJobArgs<'_>) -> Result) -> Result AnyhowResult { + Ok(Self { + inference_command: VitsInferenceCommand::from_env()?, + }) + } +} diff --git a/crates/service/job/inference_job/src/job/job_types/tts/vits/vits_inference_command.rs b/crates/service/job/inference_job/src/job/job_types/tts/vits/vits_inference_command.rs index 4e42412d59..758ae79d49 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/vits/vits_inference_command.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/vits/vits_inference_command.rs @@ -5,7 +5,7 @@ use log::info; use subprocess::{Popen, PopenConfig}; use filesys::path_to_string::path_to_string; -use subprocess_common::docker_options::DockerOptions; +use subprocess_common::docker_options::{DockerEnvVar, DockerFilesystemMount, DockerGpu, DockerOptions}; use crate::AnyhowResult; @@ -60,6 +60,64 @@ pub struct VitsInferenceArgs> { } impl VitsInferenceCommand { + + pub fn from_env() -> AnyhowResult { + let root_directory = easyenv::get_env_string_required( + "VITS_INFERENCE_ROOT_DIRECTORY")?; + + let inference_script = easyenv::get_env_string_or_default( + "VITS_INFERENCE_SCRIPT", + "infer_ts_job.py"); + + let maybe_venv_command = easyenv::get_env_string_optional( + "VITS_INFERENCE_MAYBE_VENV_COMMAND"); + + let maybe_python_interpreter = easyenv::get_env_string_optional( + "VITS_INFERENCE_MAYBE_PYTHON_INTERPRETER"); + + let maybe_huggingface_dataset_cache = easyenv::get_env_string_optional( + "HF_DATASETS_CACHE"); + + let maybe_nltk_data_cache = easyenv::get_env_string_optional( + "NLTK_DATA"); + + let mut docker_env_vars = Vec::new(); + + if let Some(cache_dir) = maybe_huggingface_dataset_cache.as_deref() { + docker_env_vars.push(DockerEnvVar::new("HF_DATASETS_CACHE", cache_dir)); + docker_env_vars.push(DockerEnvVar::new("HF_HOME", cache_dir)); + } + + if let Some(cache_dir) = maybe_nltk_data_cache.as_deref() { + docker_env_vars.push(DockerEnvVar::new("NLTK_DATA", cache_dir)); + docker_env_vars.push(DockerEnvVar::new("NLTK_DATA_PATH", cache_dir)); + } + + let maybe_docker_env_vars = + if docker_env_vars.is_empty() { None } else { Some(docker_env_vars) }; + + let maybe_docker_options = easyenv::get_env_string_optional( + "VITS_INFERENCE_MAYBE_DOCKER_IMAGE_SHA") + .map(|image_name| { + DockerOptions { + image_name, + maybe_bind_mount: Some(DockerFilesystemMount::tmp_to_tmp()), + maybe_environment_variables: maybe_docker_env_vars, + maybe_gpu: Some(DockerGpu::All), + } + }); + + Ok(Self::new( + root_directory, + inference_script, + maybe_python_interpreter.as_deref(), + maybe_venv_command.as_deref(), + maybe_huggingface_dataset_cache.as_deref(), + maybe_nltk_data_cache.as_deref(), + maybe_docker_options, + )) + } + pub fn new>( vits_root_code_directory: P, inference_script_name: P, diff --git a/crates/service/job/inference_job/src/job_dependencies.rs b/crates/service/job/inference_job/src/job_dependencies.rs index 1c3ca5b46d..f88c164697 100644 --- a/crates/service/job/inference_job/src/job_dependencies.rs +++ b/crates/service/job/inference_job/src/job_dependencies.rs @@ -24,7 +24,6 @@ use mysql_queries::queries::tts::tts_models::get_tts_model_for_inference_improve use mysql_queries::queries::voice_conversion::inference::get_voice_conversion_model_for_inference::VoiceConversionModelForInference; use newrelic_telemetry::Client as NewRelicClient; -use crate::job::job_types::tts::vits::vits_inference_command::VitsInferenceCommand; use crate::job_specific_dependencies::JobSpecificDependencies; use crate::util::scoped_execution::ScopedExecution; use crate::util::scoped_temp_dir_creator::ScopedTempDirCreator; @@ -100,9 +99,6 @@ pub struct JobDependencies { // This finds jobs of equal or greater priority. pub maybe_minimum_priority: Option, - // Details for each job type (grouped by the job type) - pub job_type_details: JobTypeDetails, - pub container: ContainerEnvironment, pub container_db: ContainerEnvironmentArg, // Same info, but for database. @@ -151,12 +147,3 @@ pub struct JobCaches { /// If the counter elapses a delta, proceed with calculation. pub model_cache_counter: TtlKeyCounter, } - -/// Per-job type details -pub struct JobTypeDetails { - pub vits: VitsDetails, -} - -pub struct VitsDetails { - pub inference_command: VitsInferenceCommand, -} 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 92b58120ee..4df0060661 100644 --- a/crates/service/job/inference_job/src/job_specific_dependencies.rs +++ b/crates/service/job/inference_job/src/job_specific_dependencies.rs @@ -6,6 +6,7 @@ use errors::AnyhowResult; use crate::job::job_types::lipsync::sad_talker::sad_talker_dependencies::SadTalkerDependencies; 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; use crate::job::job_types::vc::rvc_v2::rvc_v2_dependencies::RvcV2Dependencies; use crate::job::job_types::vc::so_vits_svc::svc_dependencies::SvcDependencies; use crate::util::scoped_execution::ScopedExecution; @@ -16,6 +17,7 @@ pub struct JobSpecificDependencies { pub maybe_svc_dependencies: Option, pub maybe_tacotron2_dependencies: Option, pub maybe_vall_e_x_dependencies: Option, + pub maybe_vits_dependencies: Option, } impl JobSpecificDependencies { @@ -26,6 +28,7 @@ impl JobSpecificDependencies { let mut maybe_svc_dependencies = None; let mut maybe_tacotron2_dependencies = None; let mut maybe_vall_e_x_dependencies = None; + let mut maybe_vits_dependencies = None; if scoped_execution.can_run_job(InferenceModelType::RvcV2) { info!("Setting RVCv2 dependencies..."); @@ -52,12 +55,18 @@ impl JobSpecificDependencies { maybe_vall_e_x_dependencies = Some(VallExDependencies::setup()?); } + if scoped_execution.can_run_job(InferenceModelType::Vits) { + info!("Setting Vits dependencies..."); + maybe_vits_dependencies = Some(VitsDependencies::setup()?); + } + Ok(JobSpecificDependencies { maybe_rvc_v2_dependencies, maybe_sad_talker_dependencies, maybe_svc_dependencies, maybe_tacotron2_dependencies, maybe_vall_e_x_dependencies, + maybe_vits_dependencies, }) } } diff --git a/crates/service/job/inference_job/src/main.rs b/crates/service/job/inference_job/src/main.rs index 8e70071d3d..a6b281b690 100644 --- a/crates/service/job/inference_job/src/main.rs +++ b/crates/service/job/inference_job/src/main.rs @@ -45,13 +45,11 @@ use memory_caching::ttl_key_counter::TtlKeyCounter; use mysql_queries::common_inputs::container_environment_arg::ContainerEnvironmentArg; use mysql_queries::mediators::firehose_publisher::FirehosePublisher; use newrelic_telemetry::ClientBuilder; -use subprocess_common::docker_options::{DockerEnvVar, DockerFilesystemMount, DockerGpu, DockerOptions}; use crate::http_server::run_http_server::CreateServerArgs; use crate::http_server::run_http_server::launch_http_server; use crate::job::job_loop::main_loop::main_loop; -use crate::job::job_types::tts::vits::vits_inference_command::VitsInferenceCommand; -use crate::job_dependencies::{FileSystemDetails, JobCaches, JobDependencies, JobTypeDetails, JobWorkerDetails, VitsDetails}; +use crate::job_dependencies::{FileSystemDetails, JobCaches, JobDependencies, JobWorkerDetails}; use crate::job_specific_dependencies::JobSpecificDependencies; use crate::util::scoped_execution::ScopedExecution; use crate::util::scoped_temp_dir_creator::ScopedTempDirCreator; @@ -291,11 +289,6 @@ async fn main() -> AnyhowResult<()> { sidecar_max_synthesizer_models, low_priority_starvation_prevention_every_nth, maybe_minimum_priority, - job_type_details: JobTypeDetails { - vits: VitsDetails { - inference_command: vits_inference_command()?, - }, - }, container: container_environment.clone(), container_db: ContainerEnvironmentArg { hostname: container_environment.hostname, @@ -320,59 +313,3 @@ async fn main() -> AnyhowResult<()> { Ok(()) } -fn vits_inference_command() -> AnyhowResult { - let root_directory = easyenv::get_env_string_required( - "VITS_INFERENCE_ROOT_DIRECTORY")?; - - let inference_script = easyenv::get_env_string_or_default( - "VITS_INFERENCE_SCRIPT", - "infer_ts_job.py"); - - let maybe_venv_command = easyenv::get_env_string_optional( - "VITS_INFERENCE_MAYBE_VENV_COMMAND"); - - let maybe_python_interpreter = easyenv::get_env_string_optional( - "VITS_INFERENCE_MAYBE_PYTHON_INTERPRETER"); - - let maybe_huggingface_dataset_cache = easyenv::get_env_string_optional( - "HF_DATASETS_CACHE"); - - let maybe_nltk_data_cache = easyenv::get_env_string_optional( - "NLTK_DATA"); - - let mut docker_env_vars = Vec::new(); - - if let Some(cache_dir) = maybe_huggingface_dataset_cache.as_deref() { - docker_env_vars.push(DockerEnvVar::new("HF_DATASETS_CACHE", cache_dir)); - docker_env_vars.push(DockerEnvVar::new("HF_HOME", cache_dir)); - } - - if let Some(cache_dir) = maybe_nltk_data_cache.as_deref() { - docker_env_vars.push(DockerEnvVar::new("NLTK_DATA", cache_dir)); - docker_env_vars.push(DockerEnvVar::new("NLTK_DATA_PATH", cache_dir)); - } - - let maybe_docker_env_vars = - if docker_env_vars.is_empty() { None } else { Some(docker_env_vars) }; - - let maybe_docker_options = easyenv::get_env_string_optional( - "VITS_INFERENCE_MAYBE_DOCKER_IMAGE_SHA") - .map(|image_name| { - DockerOptions { - image_name, - maybe_bind_mount: Some(DockerFilesystemMount::tmp_to_tmp()), - maybe_environment_variables: maybe_docker_env_vars, - maybe_gpu: Some(DockerGpu::All), - } - }); - - Ok(VitsInferenceCommand::new( - root_directory, - inference_script, - maybe_python_interpreter.as_deref(), - maybe_venv_command.as_deref(), - maybe_huggingface_dataset_cache.as_deref(), - maybe_nltk_data_cache.as_deref(), - maybe_docker_options, - )) -}