From 3d4a870190d367e9c7402ba17f80d95cf7dda07a Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Sun, 26 Nov 2023 18:08:14 -0500 Subject: [PATCH] streamline inference-job tt2 --- .../download_static_dependencies.rs | 97 +++++++++++++++++++ .../tts/tacotron2_v2_early_fakeyou/mod.rs | 1 + .../tacotron2_v2_early_fakeyou/process_job.rs | 83 +++------------- 3 files changed, 114 insertions(+), 67 deletions(-) create mode 100644 crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/download_static_dependencies.rs diff --git a/crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/download_static_dependencies.rs b/crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/download_static_dependencies.rs new file mode 100644 index 0000000000..bd46e1a8a2 --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/download_static_dependencies.rs @@ -0,0 +1,97 @@ +use std::path::PathBuf; + +use jobs_common::job_progress_reporter::job_progress_reporter::JobProgressReporter; +use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; + +use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; +use crate::job::job_types::tts::tacotron2_v2_early_fakeyou::tacotron2_dependencies::Tacotron2Dependencies; +use crate::job_dependencies::JobDependencies; +use crate::util::maybe_download_file_from_bucket::maybe_download_file_from_bucket; + +pub struct StaticDependencies { + pub waveglow_vocoder_model_fs_path : PathBuf, + pub pretrained_hifigan_vocoder_model_fs_path: PathBuf, + pub hifigan_superres_vocoder_model_fs_path : PathBuf, +} + + +pub async fn download_static_dependencies<'a>( + job_dependencies: &'a JobDependencies, + job: &'a AvailableInferenceJob, + model_dependencies: &'a Tacotron2Dependencies, + mut job_progress_reporter: &mut Box, +) -> Result { + + + // ==================== CONFIRM OR DOWNLOAD WAVEGLOW VOCODER MODEL ==================== // + + let waveglow_vocoder_model_fs_path = { + // TODO(bt,2023-11-21): Port this to the common downloader code. + let waveglow_vocoder_model_filename = model_dependencies.waveglow_vocoder_model_filename.clone(); + let waveglow_vocoder_model_fs_path = job_dependencies.fs.semi_persistent_cache.tts_pretrained_vocoder_model_path(&waveglow_vocoder_model_filename); + let waveglow_vocoder_model_object_path = job_dependencies.buckets.bucket_path_unifier.tts_pretrained_vocoders_path(&waveglow_vocoder_model_filename); + + maybe_download_file_from_bucket( + "waveglow vocoder model", + &waveglow_vocoder_model_fs_path, + &waveglow_vocoder_model_object_path, + &job_dependencies.buckets.private_bucket_client, + &mut job_progress_reporter, + "downloading vocoder (1 of 3)", + job.id.0, + &job_dependencies.fs.scoped_temp_dir_creator_for_downloads, + ).await?; + + waveglow_vocoder_model_fs_path + }; + + // ==================== CONFIRM OR DOWNLOAD HIFIGAN (NORMAL) VOCODER MODEL ==================== // + + let pretrained_hifigan_vocoder_model_fs_path = { + // TODO(bt,2023-11-21): Port this to the common downloader code. + let hifigan_vocoder_model_filename = model_dependencies.hifigan_vocoder_model_filename.clone(); + let hifigan_vocoder_model_fs_path = job_dependencies.fs.semi_persistent_cache.tts_pretrained_vocoder_model_path(&hifigan_vocoder_model_filename); + let hifigan_vocoder_model_object_path = job_dependencies.buckets.bucket_path_unifier.tts_pretrained_vocoders_path(&hifigan_vocoder_model_filename); + + maybe_download_file_from_bucket( + "hifigan vocoder model", + &hifigan_vocoder_model_fs_path, + &hifigan_vocoder_model_object_path, + &job_dependencies.buckets.private_bucket_client, + &mut job_progress_reporter, + "downloading vocoder (2 of 3)", + job.id.0, + &job_dependencies.fs.scoped_temp_dir_creator_for_downloads, + ).await?; + + hifigan_vocoder_model_fs_path + }; + + // ==================== CONFIRM OR DOWNLOAD HIFIGAN (SUPERRES) VOCODER MODEL ==================== // + + let hifigan_superres_vocoder_model_fs_path = { + // TODO(bt,2023-11-21): Port this to the common downloader code. + let hifigan_superres_vocoder_model_filename = model_dependencies.hifigan_superres_vocoder_model_filename.clone(); + let hifigan_superres_vocoder_model_fs_path = job_dependencies.fs.semi_persistent_cache.tts_pretrained_vocoder_model_path(&hifigan_superres_vocoder_model_filename); + let hifigan_superres_vocoder_model_object_path = job_dependencies.buckets.bucket_path_unifier.tts_pretrained_vocoders_path(&hifigan_superres_vocoder_model_filename); + + maybe_download_file_from_bucket( + "hifigan superres vocoder model", + &hifigan_superres_vocoder_model_fs_path, + &hifigan_superres_vocoder_model_object_path, + &job_dependencies.buckets.private_bucket_client, + &mut job_progress_reporter, + "downloading vocoder (3 of 3)", + job.id.0, + &job_dependencies.fs.scoped_temp_dir_creator_for_downloads, + ).await?; + + hifigan_superres_vocoder_model_fs_path + }; + + Ok(StaticDependencies { + waveglow_vocoder_model_fs_path, + pretrained_hifigan_vocoder_model_fs_path, + hifigan_superres_vocoder_model_fs_path, + }) +} diff --git a/crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/mod.rs b/crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/mod.rs index 3ba466c600..e96b66620b 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/mod.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/mod.rs @@ -1,3 +1,4 @@ +pub mod download_static_dependencies; pub mod process_job; pub mod read_metadata_file; pub mod seconds_to_decoder_steps; 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 710ae3953a..451a223700 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 @@ -23,6 +23,7 @@ use tts_common::text_pipelines::text_pipeline_type::TextPipelineType; use crate::job::job_loop::job_success_result::{JobSuccessResult, ResultEntity}; use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; +use crate::job::job_types::tts::tacotron2_v2_early_fakeyou::download_static_dependencies::download_static_dependencies; 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, VocoderForInferenceOption}; use crate::job_dependencies::JobDependencies; @@ -58,71 +59,19 @@ pub async fn process_job(args: ProcessJobArgs<'_>) -> Result) -> Result { VocoderForInferenceOption::Waveglow { - waveglow_vocoder_checkpoint_path: &waveglow_vocoder_model_fs_path + waveglow_vocoder_checkpoint_path: &static_deps.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, + hifigan_vocoder_checkpoint_path: &static_deps.pretrained_hifigan_vocoder_model_fs_path, + hifigan_superres_vocoder_checkpoint_path: &static_deps.hifigan_superres_vocoder_model_fs_path, } } }; @@ -229,7 +178,7 @@ pub async fn process_job(args: ProcessJobArgs<'_>) -> Result