From 93157b50c8b20255fa3102750af70445ce7ed90f Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 6 Apr 2023 18:58:07 -0700 Subject: [PATCH] vits inference command roughly done, need to test --- .../tts/tts_results/insert_tts_result.rs | 7 +- .../job_types/tts/vits/process_vits_model.rs | 2 +- .../tts/vits/vits_model_check_command.rs | 19 +-- .../job/inference_job/inference-job.env | 5 +- .../tacotron2_v2_early_fakeyou/process_job.rs | 3 +- .../src/job/job_types/tts/vits/process_job.rs | 104 ++---------- .../tts/vits/vits_inference_command.rs | 149 +++++------------- .../job/inference_job/src/job_dependencies.rs | 3 +- crates/service/job/inference_job/src/main.rs | 37 ++++- .../src/job_steps/process_single_job.rs | 2 +- 10 files changed, 114 insertions(+), 217 deletions(-) diff --git a/crates/lib/mysql_queries/src/queries/tts/tts_results/insert_tts_result.rs b/crates/lib/mysql_queries/src/queries/tts/tts_results/insert_tts_result.rs index 49cd80a550..1a43939831 100644 --- a/crates/lib/mysql_queries/src/queries/tts/tts_results/insert_tts_result.rs +++ b/crates/lib/mysql_queries/src/queries/tts/tts_results/insert_tts_result.rs @@ -24,7 +24,7 @@ pub async fn insert_tts_result>( pool: &MySqlPool, job: JobType<'_>, text_hash: &str, - pretrained_vocoder_used: VocoderType, + maybe_pretrained_vocoder_used: Option, bucket_audio_results_path: P, bucket_spectrogram_results_path: P, file_size_bytes: u64, @@ -46,6 +46,9 @@ pub async fn insert_tts_result>( .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, diff --git a/crates/service/job/download_job/src/job_types/tts/vits/process_vits_model.rs b/crates/service/job/download_job/src/job_types/tts/vits/process_vits_model.rs index 0b8fa454a5..822b9dfd96 100644 --- a/crates/service/job/download_job/src/job_types/tts/vits/process_vits_model.rs +++ b/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>( diff --git a/crates/service/job/download_job/src/job_types/tts/vits/vits_model_check_command.rs b/crates/service/job/download_job/src/job_types/tts/vits/vits_model_check_command.rs index 4d753afa27..eb65f10450 100644 --- a/crates/service/job/download_job/src/job_types/tts/vits/vits_model_check_command.rs +++ 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>(path: P) -> String { - path.as_ref().display().to_string() -} diff --git a/crates/service/job/inference_job/inference-job.env b/crates/service/job/inference_job/inference-job.env index d38dd0b731..a1689edfe2 100644 --- a/crates/service/job/inference_job/inference-job.env +++ b/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" diff --git a/crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/process_job.rs b/crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/process_job.rs index b9b207a72f..226ac1c373 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/process_job.rs +++ b/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, 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 f4d753ef7c..371f26ec6d 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 @@ -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 = 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, 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 4b4c0a46d6..670a427704 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 @@ -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, } -#[derive(Clone)] -pub enum VocoderForInferenceOption> { - 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> { + /// --checkpoint: input path of the model checkpoint + pub model_checkpoint_path: P, -pub struct InferenceArgs <'a, P: AsRef> { - //# 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

, - - //pub vocoder_type: &'a str, - //pub waveglow_vocoder_checkpoint_path: Option

, - //pub hifigan_vocoder_checkpoint_path: Optin

, - //pub hifigan_superres_vocoder_checkpoint_path: Option

, - - //# 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, - pub maybe_mel_multiply_factor: Option, - - //# 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>( - 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, ) -> 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>( &self, - args: InferenceArgs<'_, P>, + args: VitsInferenceArgs

, ) -> 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); diff --git a/crates/service/job/inference_job/src/job_dependencies.rs b/crates/service/job/inference_job/src/job_dependencies.rs index 4a315141bb..2afd0bd673 100644 --- a/crates/service/job/inference_job/src/job_dependencies.rs +++ b/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, + pub inference_command: VitsInferenceCommand, } diff --git a/crates/service/job/inference_job/src/main.rs b/crates/service/job/inference_job/src/main.rs index 651146a27f..c2df632338 100644 --- a/crates/service/job/inference_job/src/main.rs +++ b/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 { + 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, + )) +} \ No newline at end of file diff --git a/crates/service/job/tts_inference_job/src/job_steps/process_single_job.rs b/crates/service/job/tts_inference_job/src/job_steps/process_single_job.rs index 5b0cdb05bc..75096d2080 100644 --- a/crates/service/job/tts_inference_job/src/job_steps/process_single_job.rs +++ b/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,