diff --git a/crates/service/job/inference_job/src/job/job_types/vc/rvc_v2/mod.rs b/crates/service/job/inference_job/src/job/job_types/vc/rvc_v2/mod.rs index 7433075eca..f9f942e196 100644 --- a/crates/service/job/inference_job/src/job/job_types/vc/rvc_v2/mod.rs +++ b/crates/service/job/inference_job/src/job/job_types/vc/rvc_v2/mod.rs @@ -1,2 +1,3 @@ +pub mod pretrained_hubert_model; pub mod process_job; pub mod rvc_v2_inference_command; \ No newline at end of file diff --git a/crates/service/job/inference_job/src/job/job_types/vc/rvc_v2/pretrained_hubert_model.rs b/crates/service/job/inference_job/src/job/job_types/vc/rvc_v2/pretrained_hubert_model.rs new file mode 100644 index 0000000000..0c21b492c0 --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/vc/rvc_v2/pretrained_hubert_model.rs @@ -0,0 +1,84 @@ +use anyhow::anyhow; +use cloud_storage::bucket_client::BucketClient; +use container_common::filesystem::safe_delete_temp_directory::safe_delete_temp_directory; +use errors::AnyhowResult; +use filesys::create_dir_all_if_missing::create_dir_all_if_missing; +use filesys::file_exists::file_exists; +use log::info; +use std::path::PathBuf; +use tempdir::TempDir; +use crate::util::scoped_temp_dir_creator::ScopedTempDirCreator; + +#[derive(Clone)] +pub struct PretrainedHubertModel { + pub cloud_bucket_path: String, + pub filesystem_path: PathBuf, +} + +impl PretrainedHubertModel { + pub fn from_env() -> Self { + // NB: For now rvc-v2 is all that uses this hubert, but other models + // may (re)use it in the future. + let cloud_bucket_path = easyenv::get_env_string_or_default( + "RVC_V2_PRETRAINED_HUBERT_BUCKET_PATH", + "/hubert_pretrained/rvc_v2_hubert_base.pt"); + + // NB: For now rvc-v2 is all that uses this hubert, but other models + // may (re)use it in the future. + let filesystem_path = easyenv::get_env_pathbuf_or_default( + "RVC_V2_PRETRAINED_HUBERT_FILESYSTEM_PATH", + "/tmp/hubert/rvc_v2_hubert_base.pt"); + + Self { + cloud_bucket_path, + filesystem_path, + } + } + + pub async fn download_if_not_on_filesystem( + &self, + bucket_client: &BucketClient, + scoped_tempdir_creator: &ScopedTempDirCreator, + ) -> AnyhowResult<()> { + if file_exists(&self.filesystem_path) { + return Ok(()); + } + + if let Some(parent_directory_path) = self.filesystem_path.parent() { + create_dir_all_if_missing(parent_directory_path)?; + } + + // NB: Download to temp directory to stop concurrent writes and race conditions from other + // workers writing to a shared volume. + // NB: TempDir exists until it goes out of scope, at which point it should delete from filesystem. + let temp_dir = scoped_tempdir_creator.new_tempdir("hubert_download") + .map_err(|e| anyhow!("problem creating tempdir: {:?}", e))?; + + let temp_path = temp_dir.path().join("download.part"); + + info!("Downloading hubert from bucket path: {:?}", &self.cloud_bucket_path); + + bucket_client.download_file_to_disk(&self.cloud_bucket_path, &temp_path) + .await + .map_err(|e| { + safe_delete_temp_directory(&temp_dir); + anyhow!("couldn't download cloud object to disk: {:?}", e) + })?; + + info!("Downloaded hubert from bucket"); + + info!("Renaming hubert file from {:?} to {:?}!", &temp_path, &self.filesystem_path); + + std::fs::rename(&temp_path, &self.filesystem_path) + .map_err(|e| { + safe_delete_temp_directory(&temp_dir); + anyhow!("couldn't rename disk files: {:?}", e) + })?; + + info!("Finished downloading hubert file to {:?}", &self.filesystem_path); + + safe_delete_temp_directory(&temp_dir); + + Ok(()) + } +} diff --git a/crates/service/job/inference_job/src/job/job_types/vc/rvc_v2/process_job.rs b/crates/service/job/inference_job/src/job/job_types/vc/rvc_v2/process_job.rs index 4f3c84fff3..5055bd586a 100644 --- a/crates/service/job/inference_job/src/job/job_types/vc/rvc_v2/process_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/vc/rvc_v2/process_job.rs @@ -1,4 +1,3 @@ -use std::path::PathBuf; use anyhow::anyhow; use buckets::public::media_uploads::original_file::MediaUploadOriginalFilePath; use buckets::public::voice_conversion_results::original_file::VoiceConversionResultOriginalFilePath; @@ -20,6 +19,7 @@ use mysql_queries::queries::generic_inference::job::list_available_generic_infer use mysql_queries::queries::media_uploads::get_media_upload_for_inference::get_media_upload_for_inference; use mysql_queries::queries::voice_conversion::inference::get_voice_conversion_model_for_inference::VoiceConversionModelForInference; use mysql_queries::queries::voice_conversion::results::insert_voice_conversion_result::{insert_voice_conversion_result, InsertArgs}; +use std::path::PathBuf; use std::time::Instant; use tokens::files::media_upload::MediaUploadToken; use tokens::users::user::UserToken; @@ -40,9 +40,16 @@ pub async fn process_job(args: RvcV2ProcessJobArgs<'_>) -> Result) -> Result) -> Result None, -// PolymorphicInferenceArgs::Vc { auto_predict_f0, .. } => *auto_predict_f0, -// }) -// .flatten() -// .unwrap_or(false); -// -// let maybe_transpose = maybe_args -// .map(|args| match args { -// PolymorphicInferenceArgs::Tts { .. } => None, -// PolymorphicInferenceArgs::Vc { transpose, .. } => *transpose, -// }) -// .flatten(); -// -// let maybe_override_f0_method = maybe_args -// .map(|args| match args { -// PolymorphicInferenceArgs::Tts { .. } => None, -// PolymorphicInferenceArgs::Vc { override_f0_method, .. } => -// *override_f0_method, -// }) -// .flatten(); - // ==================== RUN INFERENCE SCRIPT ==================== // let inference_start_time = Instant::now(); @@ -225,6 +204,7 @@ pub async fn process_job(args: RvcV2ProcessJobArgs<'_>) -> Result, Q: AsRef> { /// --model_index_path: model index path pub maybe_model_index_path: Option, + /// --hubert_model_path: path to the hubert model on the filesystem + pub hubert_path: P, + /// --input_audio_filename: input wav path pub input_path: P, @@ -190,6 +193,9 @@ impl RvcV2InferenceCommand { command.push_str(&path_to_string(model_index_path)); } + command.push_str(" --hubert_model_path "); + command.push_str(&path_to_string(args.hubert_path)); + command.push_str(" --input_audio_filename "); command.push_str(&path_to_string(args.input_path)); diff --git a/crates/service/job/inference_job/src/job_dependencies.rs b/crates/service/job/inference_job/src/job_dependencies.rs index 666ef890a8..ba57506de2 100644 --- a/crates/service/job/inference_job/src/job_dependencies.rs +++ b/crates/service/job/inference_job/src/job_dependencies.rs @@ -20,6 +20,7 @@ use r2d2_redis::r2d2; use sqlx::MySqlPool; use std::path::PathBuf; use concurrency::relaxed_atomic_bool::RelaxedAtomicBool; +use crate::job::job_types::vc::rvc_v2::pretrained_hubert_model::PretrainedHubertModel; use crate::job::job_types::vc::rvc_v2::rvc_v2_inference_command::RvcV2InferenceCommand; pub struct JobDependencies { @@ -87,6 +88,8 @@ pub struct JobDependencies { // Details for each job type (grouped by the job type) pub job_type_details: JobTypeDetails, + pub pretrained_models: PretrainedModels, + pub container: ContainerEnvironment, pub container_db: ContainerEnvironmentArg, // Same info, but for database. @@ -136,6 +139,10 @@ pub struct JobCaches { pub model_cache_counter: TtlKeyCounter, } +pub struct PretrainedModels { + pub rvc_v2_hubert: PretrainedHubertModel, +} + /// Per-job type details pub struct JobTypeDetails { pub tacotron2_old_vocodes: Tacotron2VocodesDetails, diff --git a/crates/service/job/inference_job/src/main.rs b/crates/service/job/inference_job/src/main.rs index 58a3bd28ab..6063951b9b 100644 --- a/crates/service/job/inference_job/src/main.rs +++ b/crates/service/job/inference_job/src/main.rs @@ -38,7 +38,7 @@ use crate::job::job_types::tts::tacotron2_v2_early_fakeyou::tacotron2_inference_ use crate::job::job_types::tts::vits::vits_inference_command::VitsInferenceCommand; use crate::job::job_types::vc::rvc_v2::rvc_v2_inference_command::RvcV2InferenceCommand; use crate::job::job_types::vc::so_vits_svc::so_vits_svc_inference_command::SoVitsSvcInferenceCommand; -use crate::job_dependencies::{FileSystemDetails, JobCaches, JobDependencies, JobTypeDetails, JobWorkerDetails, RvcV2Details, SoVitsSvcDetails, Tacotron2VocodesDetails, VitsDetails}; +use crate::job_dependencies::{FileSystemDetails, JobCaches, JobDependencies, JobTypeDetails, JobWorkerDetails, PretrainedModels, RvcV2Details, SoVitsSvcDetails, Tacotron2VocodesDetails, VitsDetails}; use crate::util::scoped_temp_dir_creator::ScopedTempDirCreator; use filesys::create_dir_all_if_missing::create_dir_all_if_missing; use jobs_common::job_progress_reporter::job_progress_reporter::JobProgressReporterBuilder; @@ -58,6 +58,7 @@ use sqlx::mysql::MySqlPoolOptions; use std::path::PathBuf; use std::time::Duration; use subprocess_common::docker_options::{DockerEnvVar, DockerFilesystemMount, DockerGpu, DockerOptions}; +use crate::job::job_types::vc::rvc_v2::pretrained_hubert_model::PretrainedHubertModel; // Buckets (shared config) const ENV_ACCESS_KEY : &'static str = "ACCESS_KEY"; @@ -308,6 +309,9 @@ async fn main() -> AnyhowResult<()> { inference_command: RvcV2InferenceCommand::from_env()?, }, }, + pretrained_models: PretrainedModels { + rvc_v2_hubert: PretrainedHubertModel::from_env(), + }, container: container_environment.clone(), container_db: ContainerEnvironmentArg { hostname: container_environment.hostname,