vits inference command roughly done, need to test

crates/lib/mysql_queries/src/queries/tts/tts_results/insert_tts_result.rs

@@ -24,7 +24,7 @@ pub async fn insert_tts_result<P: AsRef<Path>>(
   pool: &MySqlPool,
   job: JobType<'_>,
   text_hash: &str,
-  pretrained_vocoder_used: VocoderType,
+  maybe_pretrained_vocoder_used: Option<VocoderType>,
   bucket_audio_results_path: P,
   bucket_spectrogram_results_path: P,
   file_size_bytes: u64,
@@ -46,6 +46,9 @@ pub async fn insert_tts_result<P: AsRef<Path>>(
       .display()
       .to_string();
 
+  let maybe_pretrained_vocoder_used = maybe_pretrained_vocoder_used
+      .map(|v| v.to_str());
+
   let raw_inference_text;
   let maybe_creator_user_token;
   let tts_model_token;
@@ -165,7 +168,7 @@ SET
         "#,
       inference_result_token,
       tts_model_token,
-      pretrained_vocoder_used.to_str(),
+      maybe_pretrained_vocoder_used,
       raw_inference_text,
       text_hash,
       normalized_inference_text,

crates/service/job/download_job/src/job_types/tts/vits/process_vits_model.rs

@@ -8,6 +8,7 @@ use crate::job_types::tts::vits::vits_model_check_command::{CheckArgs, Device};
 use enums::by_table::tts_models::tts_model_type::TtsModelType;
 use enums::common::visibility::Visibility;
 use errors::AnyhowResult;
+use filesys::filename_concat::filename_concat;
 use hashing::sha256::sha256_hash_file::sha256_hash_file;
 use jobs_common::redis_job_status_logger::RedisJobStatusLogger;
 use log::{error, info, warn};
@@ -19,7 +20,6 @@ use std::io::Read;
 use std::path::{Path, PathBuf};
 use std::time::Duration;
 use tempdir::TempDir;
-use filesys::filename_concat::filename_concat;
 
 /// Returns the token of the entity.
 pub async fn process_vits_model<'a, 'b>(

crates/service/job/download_job/src/job_types/tts/vits/vits_model_check_command.rs

@@ -1,4 +1,5 @@
 use container_common::anyhow_result::AnyhowResult;
+use filesys::path_to_string::path_to_string;
 use log::info;
 use std::fs::OpenOptions;
 use std::path::{Path, PathBuf};
@@ -69,7 +70,7 @@ impl VitsModelCheckCommand {
   ) -> AnyhowResult<()> {
 
     let mut command = String::new();
-    command.push_str(&format!("cd {}", path_to_str(&self.vits_root_code_directory)));
+    command.push_str(&format!("cd {}", path_to_string(&self.vits_root_code_directory)));
 
     if let Some(venv_command) = self.maybe_virtual_env_activation_command.as_deref() {
       command.push_str(" && ");
@@ -84,18 +85,18 @@ impl VitsModelCheckCommand {
     command.push_str(" && ");
     command.push_str(python_binary);
     command.push_str(" ");
-    command.push_str(&path_to_str(&self.check_script_name));
+    command.push_str(&path_to_string(&self.check_script_name));
 
     // ===== Begin Python Args =====
 
     command.push_str(" --out-path ");
-    command.push_str(&path_to_str(args.traced_model_output_path));
+    command.push_str(&path_to_string(args.traced_model_output_path));
 
     command.push_str(" --checkpoint ");
-    command.push_str(&path_to_str(args.model_checkpoint_path));
+    command.push_str(&path_to_string(args.model_checkpoint_path));
 
     command.push_str(" --config ");
-    command.push_str(&path_to_str(args.config_path));
+    command.push_str(&path_to_string(args.config_path));
 
     let device = match args.device {
       Device::Cuda => "cuda",
@@ -103,11 +104,11 @@ impl VitsModelCheckCommand {
     };
 
     command.push_str(" --device ");
-    command.push_str(&path_to_str(device));
+    command.push_str(&path_to_string(device));
 
     command.push_str(" --test-string ");
     command.push_str("'");
-    command.push_str(&path_to_str(args.test_string));
+    command.push_str(&path_to_string(args.test_string));
     command.push_str("'");
 
     // ===== End Python Args =====
@@ -137,7 +138,3 @@ impl VitsModelCheckCommand {
     Ok(())
   }
 }
-
-fn path_to_str<P: AsRef<Path>>(path: P) -> String {
-  path.as_ref().display().to_string()
-}

crates/service/job/inference_job/inference-job.env

@@ -1,3 +1,6 @@
 # inference-job default .env configurations (these can be overridden with k8s)
 
-VITS_DOCKER_IMAGE_SHA="1e2357c3c154"
+VITS_INFERENCE_ROOT_DIRECTORY="/vits"
+VITS_INFERENCE_SCRIPT="infer_ts_job.py"
+VITS_INFERENCE_MAYBE_PYTHON_INTERPRETER="python3"
+VITS_INFERENCE_MAYBE_DOCKER_IMAGE_SHA="1e2357c3c154"

crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/process_job.rs

@@ -248,7 +248,6 @@ pub async fn process_job(args: ProcessJobArgs<'_>) -> Result<(), ProcessSingleJo
 
   // ==================== RUN INFERENCE SCRIPT ==================== //
 
-  //let tacotron_code_root_directory = "/home/bt/dev/storyteller/storyteller-ml/tts/tacotron2_v1_early_fakeyou";
   let tacotron_code_root_directory = "/models/tts";
   //let python = Some("python3");
   let python = None;
@@ -352,7 +351,7 @@ pub async fn process_job(args: ProcessJobArgs<'_>) -> Result<(), ProcessSingleJo
     &args.job_dependencies.mysql_pool,
     JobType::GenericInferenceJob(&job),
     &text_hash,
-    pretrained_vocoder,
+    Some(pretrained_vocoder),
     &audio_result_object_path,
     &spectrogram_result_object_path,
     file_metadata.file_size_bytes,

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

@@ -5,9 +5,11 @@ use container_common::filesystem::safe_delete_temp_file::safe_delete_temp_file;
 use crate::job::job_loop::process_single_job_error::ProcessSingleJobError;
 use crate::job::job_types::tts::tacotron2_v2_early_fakeyou::seconds_to_decoder_steps::seconds_to_decoder_steps;
 use crate::job::job_types::tts::tacotron2_v2_early_fakeyou::tacotron2_inference_command::{InferenceArgs, MelMultiplyFactor, Tacotron2InferenceCommand, VocoderForInferenceOption};
+use crate::job::job_types::tts::vits::vits_inference_command::{Device, VitsInferenceArgs, VitsInferenceCommand};
 use crate::job_dependencies::JobDependencies;
 use crate::util::maybe_download_file_from_bucket::maybe_download_file_from_bucket;
 use errors::AnyhowResult;
+use filesys::filename_concat::filename_concat;
 use hashing::sha256::sha256_hash_string::sha256_hash_string;
 use log::{error, info};
 use mysql_queries::column_types::vocoder_type::VocoderType;
@@ -108,112 +110,35 @@ pub async fn process_job(args: VitsProcessJobArgs<'_>) -> Result<(), ProcessSing
   info!("Running TTS inference...");
 
   info!("Expected output audio filename: {:?}", &output_audio_fs_path);
-  //info!("Expected output spectrogram filename: {:?}", &output_spectrogram_fs_path);
   info!("Expected output metadata filename: {:?}", &output_metadata_fs_path);
+  //info!("Expected output spectrogram filename: {:?}", &output_spectrogram_fs_path);
 
-  let mut pretrained_vocoder = VocoderType::HifiGanSuperResolution;
-  if let Some(default_vocoder) = tts_model.maybe_default_pretrained_vocoder.as_deref() {
-    pretrained_vocoder = VocoderType::from_str(default_vocoder)
-        .map_err(|e| ProcessSingleJobError::Other(e))?;
-  }
-
-  info!("With pretrained vocoder: {:?}", pretrained_vocoder);
-
-  // TODO: Clean up the vocoder selection logic to make this crystal clear.
-  let mut vocoder_option = match pretrained_vocoder {
-    // We most likely will *not* use WaveGlow.
-    VocoderType::WaveGlow => {
-      VocoderForInferenceOption::Waveglow {
-        waveglow_vocoder_checkpoint_path: &waveglow_vocoder_model_fs_path
-      }
-    }
-    VocoderType::HifiGanSuperResolution => {
-      VocoderForInferenceOption::HifiganSuperres {
-        hifigan_vocoder_checkpoint_path: &pretrained_hifigan_vocoder_model_fs_path,
-        hifigan_superres_vocoder_checkpoint_path: &hifigan_superres_vocoder_model_fs_path,
-      }
-    }
-  };
-
-  if let Some(ref custom_vocoder_path) = custom_vocoder_fs_path {
-      info!("using custom user-trained HiFi-GAN vocoder: {:?}", custom_vocoder_fs_path);
-      vocoder_option = VocoderForInferenceOption::HifiganSuperres {
-        hifigan_vocoder_checkpoint_path: custom_vocoder_path,
-        hifigan_superres_vocoder_checkpoint_path: &hifigan_superres_vocoder_model_fs_path,
-      };
-  };
-
-  let text_pipeline_type_or_guess = tts_model.text_pipeline_type
-      .as_deref()
-      // NB: If there's an error deserializing, turn it to None.
-      .and_then(|pipeline_type| TextPipelineType::from_str(pipeline_type).ok())
-      .unwrap_or_else(|| guess_text_pipeline_heuristic(Some(tts_model.created_at)));
-
-  info!("With text pipeline type `{:?} ` (or guess: {:?})",
-    &tts_model.text_pipeline_type,
-    &text_pipeline_type_or_guess);
-
+  // TODO: Limit output length for premium.
   // NB: Tacotron operates on decoder steps. 1000 steps is the default and correlates to
   //  roughly 12 seconds max. Here we map seconds to decoder steps.
-  let max_decoder_steps = seconds_to_decoder_steps(job.max_duration_seconds);
+  //let max_decoder_steps = seconds_to_decoder_steps(job.max_duration_seconds);
 
-  // NB: Docker options are for development, not production.
-  let mut maybe_docker_options = None;
-
-  if let Some(docker_sha) = args.job_dependencies.job_type_details.tacotron2_old_vocodes.maybe_docker_image_sha.as_deref() {
-    maybe_docker_options = Some(DockerOptions {
-      image_name: docker_sha.to_string(),
-      maybe_bind_mount: Some(DockerFilesystemMount {
-        local_filesystem: "/tmp".to_string(),
-        container_filesystem: "/tmp".to_string(),
-      }),
-      maybe_gpu: Some(DockerGpu::All),
-    })
-  }
 
   // ==================== RUN INFERENCE SCRIPT ==================== //
 
-  //let tacotron_code_root_directory = "/home/bt/dev/storyteller/storyteller-ml/tts/tacotron2_v1_early_fakeyou";
-  let tacotron_code_root_directory = "/models/tts";
-  //let python = Some("python3");
-  let python = None;
+  let config_path = PathBuf::from("configs/ljs_li44_tmbert_nmp_s1_arpa.json"); // TODO: This could be variable.
 
-  let inference_command = Tacotron2InferenceCommand::new(
-    tacotron_code_root_directory,
-    python,
-    Some("source python/bin/activate"),
-    "vocodes_inference_updated.py",
-    maybe_docker_options,
-  );
-
-  let mut maybe_mel_multiply_factor = None;
-
-  if let Some(factor) = tts_model.maybe_custom_mel_multiply_factor {
-    maybe_mel_multiply_factor = Some(MelMultiplyFactor::CustomMultiplyFactor(factor));
-  } else if tts_model.use_default_mel_multiply_factor {
-    maybe_mel_multiply_factor = Some(MelMultiplyFactor::DefaultMultiplyFactor);
-  }
-
-  let _r = inference_command.execute_inference(InferenceArgs {
-    synthesizer_checkpoint_path: &tts_synthesizer_fs_path,
-    text_pipeline_type: text_pipeline_type_or_guess.to_str(),
-    vocoder: vocoder_option,
-    maybe_mel_multiply_factor,
-    max_decoder_steps,
+  let _r = args.job_dependencies.job_type_details.vits.inference_command.execute_inference(VitsInferenceArgs {
+    model_checkpoint_path: &vits_traced_synthesizer_fs_path,
+    config_path: &config_path,
+    device: Device::Cuda,
     input_text_filename: &text_input_fs_path,
     output_audio_filename: &output_audio_fs_path,
-    output_spectrogram_filename: &output_spectrogram_fs_path,
     output_metadata_filename: &output_metadata_fs_path,
   });
 
-
   // ==================== CHECK ALL FILES EXIST AND GET METADATA ==================== //
 
   info!("Checking that output files exist...");
 
   check_file_exists(&output_audio_fs_path).map_err(|e| ProcessSingleJobError::Other(e))?;
-  check_file_exists(&output_spectrogram_fs_path).map_err(|e| ProcessSingleJobError::Other(e))?;
   check_file_exists(&output_metadata_fs_path).map_err(|e| ProcessSingleJobError::Other(e))?;
+  //check_file_exists(&output_spectrogram_fs_path).map_err(|e| ProcessSingleJobError::Other(e))?;
 
   let file_metadata = read_metadata_file(&output_metadata_fs_path)
       .map_err(|e| ProcessSingleJobError::Other(e))?;
@@ -274,15 +199,16 @@ pub async fn process_job(args: VitsProcessJobArgs<'_>) -> Result<(), ProcessSing
   info!("Saving tts inference record...");
 
   // NB: The stupid DB field for spectrograms is not nullable, so we'll just set empty string.
-  const FAKE_SPECTROGRAM_PATH : PathBuf = PathBuf::from("");
+  let fake_spectrogram_path = PathBuf::from("");
+  const NO_PRETRAINED_VOCODER : Option<VocoderType> = None;
 
   let (id, inference_result_token) = insert_tts_result(
     &args.job_dependencies.mysql_pool,
     JobType::GenericInferenceJob(&job),
     &text_hash,
-    pretrained_vocoder,
+    NO_PRETRAINED_VOCODER,
     &audio_result_object_path,
-    &FAKE_SPECTROGRAM_PATH, // NB: No spectogram!
+    &fake_spectrogram_path, // NB: No spectogram!
     file_metadata.file_size_bytes,
     file_metadata.duration_millis.unwrap_or(0),
     args.job_dependencies.worker_details.is_on_prem,

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

@@ -7,11 +7,11 @@ use subprocess_common::docker_options::DockerOptions;
 
 /// This command is used to run tacotron2 (v1 "early fakeyou") inference
 #[derive(Clone)]
-pub struct Tacotron2InferenceCommand {
-  /// Where the TT2 code lives
-  tacotron_code_root_directory: PathBuf,
+pub struct VitsInferenceCommand {
+  /// Where the VITS code lives
+  vits_root_code_directory: PathBuf,
 
-  /// The name of the inference script, eg. `vocodes_inference_updated.py`
+  /// The name of the check/process script, eg. `infer_ts.py`
   inference_script_name: PathBuf,
 
   /// eg. `source python/bin/activate`
@@ -24,92 +24,55 @@ pub struct Tacotron2InferenceCommand {
   maybe_docker_options: Option<DockerOptions>,
 }
 
-#[derive(Clone)]
-pub enum VocoderForInferenceOption<P: AsRef<Path>> {
-  Waveglow {
-    waveglow_vocoder_checkpoint_path: P,
-  },
-  HifiganSuperres {
-    hifigan_vocoder_checkpoint_path: P,
-    hifigan_superres_vocoder_checkpoint_path: P,
-  }
+pub enum Device {
+  Cuda,
+  Cpu,
 }
 
-pub enum MelMultiplyFactor {
-  // NB: Default is typically "1.4"
-  DefaultMultiplyFactor,
-  // Custom values tend to range from 1.1 to 1.5
-  CustomMultiplyFactor(f64),
-}
+pub struct VitsInferenceArgs<P: AsRef<Path>> {
+  /// --checkpoint: input path of the model checkpoint
+  pub model_checkpoint_path: P,
 
-pub struct InferenceArgs <'a, P: AsRef<Path>> {
-  //# Model parameters
-  //parser.add_argument('--synthesizer_checkpoint_path', type=str, help='path the TTS synthesizer model', required=True)
-  //parser.add_argument('--text_pipeline_type', type=str, help='', required=True)
-  //parser.add_argument('--vocoder_type', type=str, help='', required=True)
-  //parser.add_argument('--waveglow_vocoder_checkpoint_path', type=str, help='path the TTS vocoder model')
-  //parser.add_argument('--hifigan_vocoder_checkpoint_path', type=str, help='path the TTS vocoder model')
-  //parser.add_argument('--hifigan_superres_vocoder_checkpoint_path', type=str, help='path the TTS vocoder model')
+  /// --config: path of the hparams json file
+  pub config_path: P,
 
-  pub synthesizer_checkpoint_path: P,
-  pub text_pipeline_type: &'a str, // TODO: Enum
+  /// --device: cpu or cuda
+  pub device: Device,
 
-  pub vocoder: VocoderForInferenceOption<P>,
-
-  //pub vocoder_type: &'a str,
-  //pub waveglow_vocoder_checkpoint_path: Option<P>,
-  //pub hifigan_vocoder_checkpoint_path: Optin<P>,
-  //pub hifigan_superres_vocoder_checkpoint_path: Option<P>,
-
-  //# Optional mel scaling before vocoding
-  //parser.add_argument('--use_default_mel_multiply_factor', type=bool, help='', action='store_true')
-  //parser.add_argument('--maybe_custom_mel_multiply_factor', type=int, help='')
-  //pub use_default_mel_multiply_factor: bool,
-  //pub maybe_custom_mel_multiply_factor: Option<f32>,
-  pub maybe_mel_multiply_factor: Option<MelMultiplyFactor>,
-
-  //# Premium features
-  //parser.add_argument('--max_decoder_steps', type=int, help='')
-  pub max_decoder_steps: u32,
-
-  //# User input
-  //parser.add_argument('--input_text_filename', type=str, help='path the file containing text to run', required=True)
+  /// --input-text-filename: input file containing the inference text
   pub input_text_filename: P,
 
-  //# Output files
-  //parser.add_argument('--output_audio_filename', type=str, help='where to save result audio', required=True)
-  //parser.add_argument('--output_spectrogram_filename', type=str, help='where to save result spectrogram', required=True)
-  //parser.add_argument('--output_metadata_filename', type=str, help='where to save extra metadata', required=True)
-
+  /// --output-audio-filename: resulting inference output audio file
   pub output_audio_filename: P,
-  pub output_spectrogram_filename: P,
+
+  /// --output-metadata-filename: resulting inference output metadata file
   pub output_metadata_filename: P,
 }
 
-impl Tacotron2InferenceCommand {
+impl VitsInferenceCommand {
   pub fn new<P: AsRef<Path>>(
-    tacotron_code_root_directory: P,
+    vits_root_code_directory: P,
+    inference_script_name: P,
     maybe_override_python_interpreter: Option<&str>,
     maybe_virtual_env_activation_command: Option<&str>,
-    inference_script_name: P,
     maybe_docker_options: Option<DockerOptions>,
   ) -> Self {
     Self {
-      tacotron_code_root_directory: tacotron_code_root_directory.as_ref().to_path_buf(),
-      maybe_override_python_interpreter: maybe_override_python_interpreter.map(|s| s.to_string()),
-      maybe_virtual_env_activation_command: maybe_virtual_env_activation_command.map(|s| s.to_string()),
+      vits_root_code_directory: vits_root_code_directory.as_ref().to_path_buf(),
       inference_script_name: inference_script_name.as_ref().to_path_buf(),
+      maybe_virtual_env_activation_command: maybe_virtual_env_activation_command.map(|s| s.to_string()),
+      maybe_override_python_interpreter: maybe_override_python_interpreter.map(|s| s.to_string()),
       maybe_docker_options,
     }
   }
 
   pub fn execute_inference<P: AsRef<Path>>(
     &self,
-    args: InferenceArgs<'_, P>,
+    args: VitsInferenceArgs<P>,
   ) -> AnyhowResult<()> {
 
     let mut command = String::new();
-    command.push_str(&format!("cd {}", path_to_string(&self.tacotron_code_root_directory)));
+    command.push_str(&format!("cd {}", path_to_string(&self.vits_root_code_directory)));
 
     if let Some(venv_command) = self.maybe_virtual_env_activation_command.as_deref() {
       command.push_str(" && ");
@@ -126,62 +89,32 @@ impl Tacotron2InferenceCommand {
     command.push_str(" ");
     command.push_str(&path_to_string(&self.inference_script_name));
 
-    // ===== Begin Python Inference Args =====
+    // ===== Begin Python Args =====
 
-    command.push_str(" --synthesizer_checkpoint_path ");
-    command.push_str(&path_to_string(args.synthesizer_checkpoint_path));
+    command.push_str(" --checkpoint ");
+    command.push_str(&path_to_string(args.model_checkpoint_path));
 
-    command.push_str(" --text_pipeline_type ");
-    command.push_str(args.text_pipeline_type);
+    command.push_str(" --config ");
+    command.push_str(&path_to_string(args.config_path));
 
-    match args.vocoder {
-      VocoderForInferenceOption::Waveglow { waveglow_vocoder_checkpoint_path } => {
-        command.push_str(" --vocoder_type ");
-        command.push_str("waveglow");
+    let device = match args.device {
+      Device::Cuda => "cuda",
+      Device::Cpu => "cpu",
+    };
 
-        command.push_str(" --waveglow_vocoder_checkpoint_path ");
-        command.push_str(&path_to_string(waveglow_vocoder_checkpoint_path));
-      }
-      VocoderForInferenceOption::HifiganSuperres {
-        hifigan_vocoder_checkpoint_path,
-        hifigan_superres_vocoder_checkpoint_path
-      } => {
-        command.push_str(" --vocoder_type ");
-        command.push_str("hifigan-superres");
+    command.push_str(" --device ");
+    command.push_str(&path_to_string(device));
 
-        command.push_str(" --hifigan_vocoder_checkpoint_path ");
-        command.push_str(&path_to_string(hifigan_vocoder_checkpoint_path));
-
-        command.push_str(" --hifigan_superres_vocoder_checkpoint_path ");
-        command.push_str(&path_to_string(hifigan_superres_vocoder_checkpoint_path));
-      }
-    }
-
-    match args.maybe_mel_multiply_factor {
-      None => {}
-      Some(MelMultiplyFactor::DefaultMultiplyFactor) => {
-        command.push_str(" --maybe_custom_mel_multiply_factor ");
-        //command.push_str("True");
-      }
-      Some(MelMultiplyFactor::CustomMultiplyFactor(factor)) => {
-        command.push_str(" --custom_mel_multiply_factor ");
-        command.push_str(&factor.to_string());
-      }
-    }
-
-    command.push_str(" --input_text_filename ");
+    command.push_str(" --input-text-filename ");
     command.push_str(&path_to_string(args.input_text_filename));
 
-    command.push_str(" --output_audio_filename ");
+    command.push_str(" --output-audio-filename ");
     command.push_str(&path_to_string(args.output_audio_filename));
 
-    command.push_str(" --output_spectrogram_filename ");
-    command.push_str(&path_to_string(args.output_spectrogram_filename));
-
-    command.push_str(" --output_metadata_filename ");
+    command.push_str(" --output-metadata-filename ");
     command.push_str(&path_to_string(args.output_metadata_filename));
 
-    // ===== End Python Inference Args =====
+    // ===== End Python Args =====
 
     if let Some(docker_options) = self.maybe_docker_options.as_ref() {
       command = docker_options.to_command_string(&command);

crates/service/job/inference_job/src/job_dependencies.rs

@@ -13,6 +13,7 @@ use r2d2_redis::RedisConnectionManager;
 use r2d2_redis::r2d2;
 use sqlx::MySqlPool;
 use std::path::PathBuf;
+use crate::job::job_types::tts::vits::vits_inference_command::VitsInferenceCommand;
 
 pub struct JobDependencies {
   pub download_temp_directory: PathBuf,
@@ -135,5 +136,5 @@ pub struct Tacotron2VocodesDetails {
 }
 
 pub struct VitsDetails {
-  pub maybe_docker_image_sha: Option<String>,
+  pub inference_command: VitsInferenceCommand,
 }

crates/service/job/inference_job/src/main.rs

@@ -45,6 +45,8 @@ use r2d2_redis::r2d2;
 use sqlx::mysql::MySqlPoolOptions;
 use std::path::PathBuf;
 use std::time::Duration;
+use subprocess_common::docker_options::{DockerFilesystemMount, DockerGpu, DockerOptions};
+use crate::job::job_types::tts::vits::vits_inference_command::VitsInferenceCommand;
 
 // Buckets (shared config)
 const ENV_ACCESS_KEY : &'static str = "ACCESS_KEY";
@@ -278,7 +280,7 @@ async fn main() -> AnyhowResult<()> {
         hifigan_superres_vocoder_model_filename,
       },
       vits: VitsDetails {
-        maybe_docker_image_sha: easyenv::get_env_string_optional("VITS_DOCKER_IMAGE_SHA"),
+        inference_command: vits_inference_command()?,
       }
     },
   };
@@ -287,3 +289,36 @@ 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_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_gpu: Some(DockerGpu::All),
+        }
+      });
+
+  Ok(VitsInferenceCommand::new(
+    root_directory,
+    inference_script,
+    maybe_python_interpreter.as_deref(),
+    maybe_venv_command.as_deref(),
+    maybe_docker_options,
+  ))
+}

crates/service/job/tts_inference_job/src/job_steps/process_single_job.rs

@@ -406,7 +406,7 @@ pub async fn process_single_job(
     &job_args.mysql_pool,
     JobType::TtsJob(&job),
     &text_hash,
-    pretrained_vocoder,
+    Some(pretrained_vocoder),
     &audio_result_object_path,
     &spectrogram_result_object_path,
     file_metadata.file_size_bytes,