vits dependencies

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;

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<JobSuccessResul
       .new_generic_inference(job.inference_job_token.as_str())
       .map_err(|e| ProcessSingleJobError::Other(anyhow!(e)))?;
 
+  let model_dependencies = args
+      .job_dependencies
+      .job_specific_dependencies
+      .maybe_vits_dependencies
+      .as_ref()
+      .ok_or_else(|| ProcessSingleJobError::JobSystemMisconfiguration(Some("missing VITS dependencies".to_string())))?;
+
   // ==================== CONFIRM OR DOWNLOAD VITS DEPENDENCIES ==================== //
 
   // TODO: Currently VITS downloads models from HuggingFace. This is likely a risk in that they can move.
@@ -123,7 +130,7 @@ pub async fn process_job(args: VitsProcessJobArgs<'_>) -> Result<JobSuccessResul
 
   let inference_start_time = Instant::now();
 
-  let _r = args.job_dependencies.job_type_details.vits.inference_command.execute_inference(VitsInferenceArgs {
+  let _r = model_dependencies.inference_command.execute_inference(VitsInferenceArgs {
     model_checkpoint_path: &vits_traced_synthesizer_fs_path,
     config_path: &config_path,
     device: Device::Cuda,

crates/service/job/inference_job/src/job/job_types/tts/vits/vits_dependencies.rs

@@ -0,0 +1,15 @@
+use errors::AnyhowResult;
+
+use crate::job::job_types::tts::vits::vits_inference_command::VitsInferenceCommand;
+
+pub struct VitsDependencies {
+  pub inference_command: VitsInferenceCommand,
+}
+
+impl VitsDependencies {
+  pub fn setup() -> AnyhowResult<Self> {
+    Ok(Self {
+      inference_command: VitsInferenceCommand::from_env()?,
+    })
+  }
+}

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<P: AsRef<Path>> {
 }
 
 impl VitsInferenceCommand {
+
+  pub fn from_env() -> AnyhowResult<Self> {
+    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<P: AsRef<Path>>(
     vits_root_code_directory: P,
     inference_script_name: P,

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<u8>,
 
-  // 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,
-}

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<SvcDependencies>,
   pub maybe_tacotron2_dependencies: Option<Tacotron2Dependencies>,
   pub maybe_vall_e_x_dependencies: Option<VallExDependencies>,
+  pub maybe_vits_dependencies: Option<VitsDependencies>,
 }
 
 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,
     })
   }
 }

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<VitsInferenceCommand> {
-  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,
-  ))
-}