model downloader trait wip

Cargo.lock

@@ -2831,6 +2831,7 @@ dependencies = [
  "actix-web 4.0.1",
  "actix_helpers",
  "anyhow",
+ "async-trait",
  "bootstrap",
  "buckets",
  "chrono",

crates/service/job/inference_job/Cargo.toml

@@ -77,3 +77,9 @@ actix-web = "4.0.0-beta.10"
 serde = "1.0.125"
 serde_derive = "1.0.125"
 serde_json = "1.0.64"
+
+# Async traits
+# NB: note: `async` trait functions are not currently supported
+# The initial round of stabilizations for the async/await language feature in Rust 1.39 did not
+# include support for async fn in traits.
+async-trait = "0.1.58"

crates/service/job/inference_job/src/job/job_types/lipsync/process_single_lipsync_job.rs

@@ -1,31 +1,20 @@
 use anyhow::anyhow;
 use crate::job::job_loop::job_success_result::JobSuccessResult;
 use crate::job::job_loop::process_single_job_error::ProcessSingleJobError;
-use crate::job::job_types::vc::rvc_v2::process_job::RvcV2ProcessJobArgs;
-use crate::job::job_types::vc::so_vits_svc::process_job::SoVitsSvcProcessJobArgs;
-use crate::job::job_types::vc::{rvc_v2, so_vits_svc};
+use crate::job::job_types::lipsync::sad_talker::process_job::SadTalkerProcessJobArgs;
+use crate::job::job_types::lipsync::sad_talker;
 use crate::job_dependencies::JobDependencies;
-use enums::by_table::voice_conversion_models::voice_conversion_model_type::VoiceConversionModelType;
-use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob;
-use mysql_queries::queries::voice_conversion::inference::get_voice_conversion_model_for_inference::get_voice_conversion_model_for_inference;
-use std::time::Duration;
-use log::{error, info};
 use enums::by_table::generic_inference_jobs::inference_model_type::InferenceModelType;
-use mysql_queries::queries::media_uploads::get_media_upload_for_inference::get_media_upload_for_inference;
-use tokens::files::media_upload::MediaUploadToken;
+use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob;
 
 pub async fn process_single_lipsync_job(job_dependencies: &JobDependencies, job: &AvailableInferenceJob) -> Result<JobSuccessResult, ProcessSingleJobError> {
 
   let job_success_result = match job.maybe_model_type {
     Some(InferenceModelType::SadTalker) => {
-      //rvc_v2::process_job::process_job(RvcV2ProcessJobArgs {
-      //  job_dependencies,
-      //  job,
-      //  vc_model: &vc_model,
-      //  media_upload_token: &media_upload_token,
-      //  media_upload: &media_upload,
-      //}).await?
-      return Err(ProcessSingleJobError::Other(anyhow!("not yet implemented")))
+      sad_talker::process_job::process_job(SadTalkerProcessJobArgs {
+        job_dependencies,
+        job,
+      }).await?
     }
     Some(model_type) => return Err(ProcessSingleJobError::Other(anyhow!("wrong model type: {:?}", model_type))),
     None => return Err(ProcessSingleJobError::Other(anyhow!("no model type in record"))),
@@ -33,4 +22,3 @@ pub async fn process_single_lipsync_job(job_dependencies: &JobDependencies, job:
 
   Ok(job_success_result)
 }
-

crates/service/job/inference_job/src/job/job_types/lipsync/sad_talker/mod.rs

@@ -0,0 +1,2 @@
+pub mod process_job;
+pub mod model_downloaders;

crates/service/job/inference_job/src/job/job_types/lipsync/sad_talker/process_job.rs

@@ -0,0 +1,362 @@
+use anyhow::anyhow;
+use buckets::public::media_uploads::original_file::MediaUploadOriginalFilePath;
+use buckets::public::voice_conversion_results::original_file::VoiceConversionResultOriginalFilePath;
+use container_common::filesystem::check_file_exists::check_file_exists;
+use container_common::filesystem::safe_delete_temp_directory::safe_delete_temp_directory;
+use container_common::filesystem::safe_delete_temp_file::safe_delete_temp_file;
+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::vc::rvc_v2::rvc_v2_inference_command::InferenceArgs;
+use crate::job_dependencies::JobDependencies;
+use crate::util::maybe_download_file_from_bucket::maybe_download_file_from_bucket;
+use enums::by_table::generic_inference_jobs::inference_result_type::InferenceResultType;
+use filesys::create_dir_all_if_missing::create_dir_all_if_missing;
+use filesys::file_size::file_size;
+use log::{error, info, warn};
+use media::decode_basic_audio_info::decode_basic_audio_file_info;
+use mimetypes::mimetype_for_file::get_mimetype_for_file;
+use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob;
+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::users::user::UserToken;
+
+pub struct SadTalkerProcessJobArgs<'a> {
+  pub job_dependencies: &'a JobDependencies,
+  pub job: &'a AvailableInferenceJob,
+}
+
+pub async fn process_job(args: SadTalkerProcessJobArgs<'_>) -> Result<JobSuccessResult, ProcessSingleJobError> {
+  let job = args.job;
+
+  /*
+  let mut job_progress_reporter = args.job_dependencies
+      .job_progress_reporter
+      .new_generic_inference(job.inference_job_token.as_str())
+      .map_err(|e| ProcessSingleJobError::Other(anyhow!(e)))?;
+
+  // ==================== DOWNLOAD HUBERT ==================== //
+
+  info!("Download RVC hubert model (if not present)...");
+
+  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 ==================== //
+
+  //info!("Download RVC model (if not present)...");
+
+  let rvc_v2_model_fs_path = {
+    let filename = format!("{}.pt", vc_model.token.as_str());
+    let fs_path = args.job_dependencies.fs.semi_persistent_cache.voice_conversion_model_path(&filename);
+
+    create_dir_all_if_missing(args.job_dependencies.fs.semi_persistent_cache.voice_conversion_model_directory())
+        .map_err(|e| {
+          error!("could not create model storage directory: {:?}", e);
+          ProcessSingleJobError::from_io_error(e)
+        })?;
+
+    let model_object_path = args.job_dependencies.bucket_path_unifier.rvc_v2_model_path(&vc_model.private_bucket_hash);
+
+    maybe_download_file_from_bucket(
+      "rvc (v2) model",
+      &fs_path,
+      &model_object_path,
+      &args.job_dependencies.private_bucket_client,
+      &mut job_progress_reporter,
+      "downloading rvc (v2) model",
+      job.id.0,
+      &args.job_dependencies.fs.scoped_temp_dir_creator_for_downloads,
+    ).await?;
+
+    fs_path
+  };
+
+  // ==================== CONFIRM OR DOWNLOAD RVC (v2) MODEL INDICES ==================== //
+
+  // Index files are optional
+  let mut maybe_rvc_v2_model_index_fs_path: Option<PathBuf> = None;
+
+  if vc_model.has_index_file {
+    info!("Download RVC index (if not present)...");
+
+    maybe_rvc_v2_model_index_fs_path = {
+      let filename = format!("{}.index", vc_model.token.as_str());
+      let fs_path = args.job_dependencies.fs.semi_persistent_cache.voice_conversion_model_path(&filename);
+
+      create_dir_all_if_missing(args.job_dependencies.fs.semi_persistent_cache.voice_conversion_model_directory())
+          .map_err(|e| {
+            error!("could not create model storage directory: {:?}", e);
+            ProcessSingleJobError::from_io_error(e)
+          })?;
+
+      let model_index_object_path = args.job_dependencies.bucket_path_unifier.rvc_v2_model_index_path(&vc_model.private_bucket_hash);
+
+      maybe_download_file_from_bucket(
+        "rvc (v2) model index",
+        &fs_path,
+        &model_index_object_path,
+        &args.job_dependencies.private_bucket_client,
+        &mut job_progress_reporter,
+        "downloading rvc (v2) model index",
+        job.id.0,
+        &args.job_dependencies.fs.scoped_temp_dir_creator_for_downloads,
+      ).await?;
+
+      Some(fs_path)
+    };
+  }
+
+  // ==================== TEMP DIR ==================== //
+
+  let work_temp_dir = format!("temp_rvc_v2_inference_{}", job.id.0);
+
+  // NB: TempDir exists until it goes out of scope, at which point it should delete from filesystem.
+  let work_temp_dir = args.job_dependencies
+      .fs
+      .scoped_temp_dir_creator_for_work
+      .new_tempdir(&work_temp_dir)
+      .map_err(|e| ProcessSingleJobError::from_io_error(e))?;
+
+  // ==================== DOWNLOAD MEDIA FILE ==================== //
+
+  info!("Download media for RVC voice conversion...");
+
+  // TODO: If already transcoded, download the transcoded file instead.
+  // TODO: Turn this into a general utility.
+
+  let original_media_upload_fs_path = {
+    let original_media_upload_fs_path = work_temp_dir.path().join("original.bin");
+
+    let media_upload_bucket_path =
+        MediaUploadOriginalFilePath::from_object_hash(&args.media_upload.public_bucket_directory_hash);
+
+    let bucket_object_path = media_upload_bucket_path.to_full_object_pathbuf();
+
+    info!("Downloading media from bucket path: {:?}", &bucket_object_path);
+
+    maybe_download_file_from_bucket(
+      "media upload (original file)",
+      &original_media_upload_fs_path,
+      &bucket_object_path,
+      &args.job_dependencies.public_bucket_client,
+      &mut job_progress_reporter,
+      "downloading",
+      job.id.0,
+      &args.job_dependencies.fs.scoped_temp_dir_creator_for_work,
+    ).await?;
+
+    original_media_upload_fs_path
+  };
+
+  // ==================== TRANSCODE MEDIA (IF NECESSARY) ==================== //
+
+  // TODO
+
+  // ==================== SETUP FOR INFERENCE ==================== //
+
+  info!("Ready for RVC (v2) inference...");
+
+  job_progress_reporter.log_status("running inference")
+      .map_err(|e| ProcessSingleJobError::Other(e))?;
+
+  //let config_path = PathBuf::from("/models/voice_conversion/so-vits-svc/example_config.json"); // TODO: This could be variable.
+  let input_wav_path = original_media_upload_fs_path;
+
+  let output_audio_fs_path = work_temp_dir.path().join("output.wav");
+  //let output_metadata_fs_path = temp_dir.path().join("metadata.json");
+  //let output_spectrogram_fs_path = temp_dir.path().join("spectrogram.json");
+
+  info!("Running RVC (v2) VC inference...");
+
+  info!("Expected output audio filename: {:?}", &output_audio_fs_path);
+
+  // TODO: Limit output length for premium.
+
+  let maybe_args = job.maybe_inference_args
+      .as_ref()
+      .map(|args| args.args.as_ref())
+      .flatten();
+
+  // ==================== RUN INFERENCE SCRIPT ==================== //
+
+  let inference_start_time = Instant::now();
+
+  let command_exit_status = args.job_dependencies
+      .job_type_details
+      .rvc_v2
+      .inference_command
+      .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,
+      });
+
+  let inference_duration = Instant::now().duration_since(inference_start_time);
+
+  info!("Inference took duration to complete: {:?}", &inference_duration);
+
+  if !command_exit_status.is_success() {
+    error!("Inference failed: {:?}", command_exit_status);
+    safe_delete_temp_file(&input_wav_path);
+    safe_delete_temp_file(&output_audio_fs_path);
+    safe_delete_temp_directory(&work_temp_dir);
+    return Err(ProcessSingleJobError::Other(anyhow!("CommandExitStatus: {:?}", command_exit_status)));
+  }
+
+  // ==================== 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_metadata_fs_path).map_err(|e| ProcessSingleJobError::Other(e))?;
+  //check_file_exists(&output_spectrogram_fs_path).map_err(|e| ProcessSingleJobError::Other(e))?;
+
+  info!("Interrogating result file properties...");
+
+  let file_size_bytes = file_size(&output_audio_fs_path)
+      .map_err(|err| ProcessSingleJobError::Other(err))?;
+
+  let maybe_mimetype = get_mimetype_for_file(&output_audio_fs_path)
+      .map_err(|err| ProcessSingleJobError::from_io_error(err))?
+      .map(|mime| mime.to_string());
+
+  let audio_info = decode_basic_audio_file_info(&output_audio_fs_path, maybe_mimetype.as_deref(), None)
+      .map_err(|err| ProcessSingleJobError::Other(err))?;
+
+  if audio_info.required_full_decode {
+    warn!("Required a full decode of the output data to get duration! That's inefficient!");
+  }
+
+  // TODO: Make a new python image that generates spectrograms from any audio file.
+
+  let file_metadata = FileMetadata {
+    duration_millis: audio_info.duration_millis,
+    mimetype: maybe_mimetype,
+    file_size_bytes,
+  };
+
+  //safe_delete_temp_file(&output_metadata_fs_path);
+
+  // ==================== UPLOAD AUDIO TO BUCKET ==================== //
+
+  job_progress_reporter.log_status("uploading result")
+      .map_err(|e| ProcessSingleJobError::Other(e))?;
+
+  let result_bucket_location = VoiceConversionResultOriginalFilePath::generate_new();
+
+  let result_bucket_object_pathbuf = result_bucket_location.to_full_object_pathbuf();
+
+  info!("Audio destination bucket path: {:?}", &result_bucket_object_pathbuf);
+
+  info!("Uploading audio...");
+
+  args.job_dependencies.public_bucket_client.upload_filename_with_content_type(
+    &result_bucket_object_pathbuf,
+    &output_audio_fs_path,
+    "audio/wav")
+      .await
+      .map_err(|e| ProcessSingleJobError::Other(e))?;
+
+  safe_delete_temp_file(&output_audio_fs_path);
+
+//  // ==================== UPLOAD SPECTROGRAM TO BUCKETS ==================== //
+//
+//  let spectrogram_result_object_path = args.job_dependencies.bucket_path_unifier.tts_inference_spectrogram_output_path(
+//    &job.uuid_idempotency_token); // TODO: Don't use this!
+//
+//  info!("Spectrogram destination bucket path: {:?}", &spectrogram_result_object_path);
+//
+//  info!("Uploading spectrogram...");
+//
+//  args.job_dependencies.public_bucket_client.upload_filename_with_content_type(
+//    &spectrogram_result_object_path,
+//    &output_spectrogram_fs_path,
+//    "application/json")
+//      .await
+//      .map_err(|e| ProcessSingleJobError::Other(e))?;
+//
+//  safe_delete_temp_file(&output_spectrogram_fs_path);
+
+  // ==================== DELETE DOWNLOADED FILE ==================== //
+
+  // NB: We should be using a tempdir, but to make absolutely certain we don't overflow the disk...
+  safe_delete_temp_directory(&work_temp_dir);
+
+  // ==================== SAVE RECORDS ==================== //
+
+  info!("Saving vc inference record...");
+
+  let (inference_result_token, id) = insert_voice_conversion_result(InsertArgs {
+    pool: &args.job_dependencies.mysql_pool,
+    job: &job,
+    public_bucket_hash: result_bucket_location.get_object_hash(),
+    file_size_bytes: file_metadata.file_size_bytes,
+    duration_millis: file_metadata.duration_millis.unwrap_or(0),
+    is_on_prem: args.job_dependencies.container.is_on_prem,
+    worker_hostname: &args.job_dependencies.container.hostname,
+    worker_cluster: &args.job_dependencies.container.cluster_name,
+    is_debug_worker: args.job_dependencies.worker_details.is_debug_worker
+  })
+      .await
+      .map_err(|e| ProcessSingleJobError::Other(e))?;
+
+  info!("VC Done.");
+
+  // TODO: Update upstream to be strongly typed
+  let maybe_user_token = job.maybe_creator_user_token.as_deref()
+      .map(|token| UserToken::new_from_str(token));
+
+  args.job_dependencies.firehose_publisher.vc_inference_finished(
+    maybe_user_token.as_ref(),
+    &job.inference_job_token,
+    inference_result_token.as_str())
+      .await
+      .map_err(|e| {
+        error!("error publishing event: {:?}", e);
+        ProcessSingleJobError::Other(anyhow!("error publishing event"))
+      })?;
+
+  job_progress_reporter.log_status("done")
+      .map_err(|e| ProcessSingleJobError::Other(e))?;
+
+  info!("Job {:?} complete success! Downloaded, ran inference, and uploaded. Saved model record: {}, Result Token: {}",
+        job.id, id, &inference_result_token);
+
+  Ok(JobSuccessResult {
+    maybe_result_entity: Some(ResultEntity {
+      entity_type: InferenceResultType::VoiceConversion,
+      entity_token: inference_result_token.to_string(),
+    }),
+    inference_duration,
+  })
+
+   */
+
+
+  // TODO - REMOVE -
+  let inference_start_time = Instant::now();
+  let inference_duration = Instant::now().duration_since(inference_start_time);
+  Ok(JobSuccessResult {
+    maybe_result_entity: Some(ResultEntity {
+      entity_type: InferenceResultType::VoiceConversion,
+      entity_token: "todo".to_string(),
+    }),
+    inference_duration,
+  })
+}
+
+#[derive(Deserialize, Default)]
+struct FileMetadata {
+  pub duration_millis: Option<u64>,
+  pub mimetype: Option<String>,
+  pub file_size_bytes: u64,
+}

crates/service/job/inference_job/src/util/from_env.rs

@@ -0,0 +1 @@
+

crates/service/job/inference_job/src/util/mod.rs

@@ -1,4 +1,6 @@
 pub mod audiowmark;
-pub mod scoped_execution;
 pub mod maybe_download_file_from_bucket;
+pub mod model_downloader;
+pub mod scoped_execution;
 pub mod scoped_temp_dir_creator;
+mod from_env;

crates/service/job/inference_job/src/util/model_downloader.rs

@@ -0,0 +1,126 @@
+use anyhow::anyhow;
+use cloud_storage::bucket_client::BucketClient;
+use container_common::filesystem::safe_delete_temp_directory::safe_delete_temp_directory;
+use crate::util::scoped_temp_dir_creator::ScopedTempDirCreator;
+use errors::AnyhowResult;
+use filesys::create_dir_all_if_missing::create_dir_all_if_missing;
+use filesys::file_exists::file_exists;
+use log::{error, info};
+use std::path::{Path, PathBuf};
+use async_trait::async_trait;
+use filesys::rename_across_devices::rename_across_devices;
+
+// /// Implementing this implies we can return an instance constructed from environment variables.
+// /// This isn't to be manually used, but rather with a macro.
+// trait FromEnv {
+//   fn from_env() -> Box<dyn Self>;
+// }
+
+
+#[async_trait]
+trait ModelDownloader {
+  //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",
+  //    // NB: For now this path is on the shared SMB drive in GKE K8S
+  //    "/tmp/downloads/hubert/rvc_v2_hubert_base.pt");
+
+  //  Self {
+  //    cloud_bucket_path,
+  //    filesystem_path,
+  //  }
+  //}
+
+  fn get_model_name(&self) -> &str;
+
+  fn get_cloud_bucket_path(&self) -> &str;
+
+  fn get_filesystem_path(&self) -> &Path;
+
+  async fn download_if_not_on_filesystem(
+    &self,
+    bucket_client: &BucketClient,
+    scoped_tempdir_creator: &ScopedTempDirCreator,
+  ) -> AnyhowResult<()> {
+    let filesystem_path = self.get_filesystem_path();
+
+    if file_exists(filesystem_path) {
+      return Ok(());
+    }
+
+    if let Some(parent_directory_path) = 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("model_download")
+        .map_err(|e| anyhow!("problem creating tempdir: {:?}", e))?;
+
+    let temp_path = temp_dir.path().join("download.part");
+
+    let model_name = self.get_model_name();
+    let cloud_bucket_path = self.get_cloud_bucket_path();
+
+    info!("Downloading {} from bucket path: {:?}", model_name, cloud_bucket_path);
+
+    bucket_client.download_file_to_disk(cloud_bucket_path, &temp_path)
+        .await
+        .map_err(|e| {
+          error!("could not download {} to disk: {:?}", model_name, e);
+          safe_delete_temp_directory(&temp_dir);
+          anyhow!("couldn't download {} cloud object to disk: {:?}", model_name, e)
+        })?;
+
+    info!("Downloaded {} from bucket", model_name);
+
+    info!("Renaming {} file from {:?} to {:?}!", model_name, &temp_path, filesystem_path);
+
+    rename_across_devices(&temp_path, filesystem_path)
+        .map_err(|e| {
+          error!("could not rename on disk: {:?}", e);
+          safe_delete_temp_directory(&temp_dir);
+          anyhow!("couldn't rename disk files: {:?}", e)
+        })?;
+
+    info!("Finished downloading {} file to {:?}", model_name, filesystem_path);
+
+    safe_delete_temp_directory(&temp_dir);
+
+    Ok(())
+  }
+}
+
+
+macro_rules! impl_model_downloader {
+  ($struct_name:ident) => {
+    #[derive(Debug, Clone)]
+    struct $struct_name {
+      pub model_name: String,
+      pub cloud_bucket_path: String,
+      pub filesystem_path: PathBuf,
+    }
+
+    #[async_trait]
+    impl ModelDownloader for $struct_name {
+      fn get_model_name(&self) -> &str {
+        &self.model_name
+      }
+      fn get_cloud_bucket_path(&self) -> &str {
+        &self.cloud_bucket_path
+      }
+      fn get_filesystem_path(&self) -> &Path {
+        &self.filesystem_path
+      }
+    }
+  }
+}