download hubert for inference-job

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;

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(())
+  }
+}

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<JobSuccessResu
       .new_generic_inference(job.inference_job_token.as_str())
       .map_err(|e| ProcessSingleJobError::Other(anyhow!(e)))?;
 
-  // ==================== CONFIRM OR DOWNLOAD RVC (v2) DEPENDENCIES ==================== //
+  // ==================== DOWNLOAD HUBERT ==================== //
 
-  // TODO: Figure out where to get hubert from.
+  args.job_dependencies.pretrained_models.rvc_v2_hubert.download_if_not_on_filesystem(
+    &args.job_dependencies.private_bucket_client,
+    &args.job_dependencies.fs.scoped_temp_dir_creator_for_downloads)
+      .await
+      .map_err(|e| {
+        error!("could not download hubert: {:?}", e);
+        ProcessSingleJobError::from_anyhow_error(e)
+      })?;
 
   // ==================== CONFIRM OR DOWNLOAD RVC (v2) MODEL ==================== //
 
@@ -178,8 +185,6 @@ pub async fn process_job(args: RvcV2ProcessJobArgs<'_>) -> Result<JobSuccessResu
   info!("Running RVC (v2) VC inference...");
 
   info!("Expected output audio filename: {:?}", &output_audio_fs_path);
-  //info!("Expected output metadata filename: {:?}", &output_metadata_fs_path);
-  //info!("Expected output spectrogram filename: {:?}", &output_spectrogram_fs_path);
 
   // TODO: Limit output length for premium.
 
@@ -188,32 +193,6 @@ pub async fn process_job(args: RvcV2ProcessJobArgs<'_>) -> Result<JobSuccessResu
       .map(|args| args.args.as_ref())
       .flatten();
 
-  // TODO: Other inference args
-
-//  // If not specified by the user, turn off auto prediction. It sounds awful.
-//  let auto_predict_f0 = maybe_args
-//      .map(|args| match args {
-//        PolymorphicInferenceArgs::Tts { .. } => 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<JobSuccessResu
       .execute_inference(InferenceArgs {
         model_path: &rvc_v2_model_fs_path,
         maybe_model_index_path: maybe_rvc_v2_model_index_fs_path,
+        hubert_path: &args.job_dependencies.pretrained_models.rvc_v2_hubert.filesystem_path,
         input_path: &input_wav_path,
         output_path: &output_audio_fs_path,
       });

crates/service/job/inference_job/src/job/job_types/vc/rvc_v2/rvc_v2_inference_command.rs

@@ -64,6 +64,9 @@ pub struct InferenceArgs<P: AsRef<Path>, Q: AsRef<Path>> {
   /// --model_index_path: model index path
   pub maybe_model_index_path: Option<Q>,
 
+  /// --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));
 

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,

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,