partially handle softvc model downloads

crates/lib/storage_buckets_common/src/bucket_path_unifier.rs

@@ -16,6 +16,8 @@ pub struct BucketPathUnifier {
   pub w2l_inference_output_root: PathBuf,
   pub w2l_model_root: PathBuf,
   pub w2l_end_bump_root: PathBuf,
+  // VC
+  pub softvc_model_root: PathBuf,
 }
 
 impl BucketPathUnifier {
@@ -33,6 +35,8 @@ impl BucketPathUnifier {
       w2l_inference_output_root: PathBuf::from("/w2l_inference_output"),
       w2l_model_root: PathBuf::from("/w2l_pretrained_models"),
       w2l_end_bump_root: PathBuf::from("/w2l_end_bumps"),
+      // VC
+      softvc_model_root: PathBuf::from("/user_uploaded_softvc_models"),
     }
   }
 
@@ -169,6 +173,19 @@ impl BucketPathUnifier {
       .join(video_filename)
   }
 
+  // ==================== VOICE CONVERSION MODELS ==================== //
+
+  pub fn softvc_model_path(&self, softvc_model_file_hash: &str) -> PathBuf {
+    let hashed_path = Self::hashed_directory_path(softvc_model_file_hash);
+    let model_filename = format!("{}.pt", &softvc_model_file_hash);
+
+    self.softvc_model_root
+        .join(hashed_path)
+        .join(model_filename)
+  }
+
+  // ==================== UTILITY ==================== //
+
   pub fn hashed_directory_path(file_hash: &str) -> String {
     match file_hash.len() {
       0 | 1=> "".to_string(),
@@ -197,6 +214,8 @@ mod tests {
       w2l_inference_output_root: PathBuf::from("/test_path_w2l_output"),
       w2l_model_root: PathBuf::from("/test_path_w2l_pretrained_models"),
       w2l_end_bump_root: PathBuf::from("/test_path_w2l_end_bumps"),
+      // VOICE CONVERSION
+      softvc_model_root: PathBuf::from("/test_path_softvc_models"),
     }
   }
 
@@ -297,6 +316,13 @@ mod tests {
                "/test_path_w2l_output/a/b/c/vocodes_video_TYPEABCDEF.mp4");
   }
 
+  #[test]
+  fn test_softvc_model_path() {
+    let paths = get_instance();
+    assert_eq!(paths.softvc_model_path("foobar").to_str().unwrap(),
+               "/test_path_softvc_models/f/o/o/foobar.pt");
+  }
+
   #[test]
   fn hashed_directory_path_length_zero() {
     assert_eq!(&BucketPathUnifier::hashed_directory_path(""), "");

crates/service/job/download_job/src/job_state.rs

@@ -1,5 +1,6 @@
 use crate::job_types::hifigan::hifigan_model_check_command::HifiGanModelCheckCommand;
 use crate::job_types::hifigan_softvc::hifigan_softvc_model_check_command::HifiGanSoftVcModelCheckCommand;
+use crate::job_types::softvc::softvc_model_check_command::SoftVcModelCheckCommand;
 use crate::job_types::tacotron::tacotron_model_check_command::TacotronModelCheckCommand;
 use database_queries::mediators::badge_granter::BadgeGranter;
 use database_queries::mediators::firehose_publisher::FirehosePublisher;
@@ -42,6 +43,7 @@ pub struct JobState {
 /// Configurations and interfaces to code deployed as sidecars or container mounts.
 pub struct SidecarConfigs {
   pub google_drive_downloader: GoogleDriveDownloadCommand,
+  pub softvc_model_check_command: SoftVcModelCheckCommand,
   pub tacotron_model_check_command: TacotronModelCheckCommand,
   pub hifigan_model_check_command: HifiGanModelCheckCommand,
   pub hifigan_softvc_model_check_command: HifiGanSoftVcModelCheckCommand,

crates/service/job/download_job/src/job_types/hifigan_softvc/hifigan_softvc_model_check_command.rs

@@ -4,10 +4,10 @@ use std::path::Path;
 use subprocess::{Popen, PopenConfig};
 use subprocess_common::docker_options::DockerOptions;
 
-/// This command is used to check hifigan for being a real model
+/// This command is used to check hifigan (softvc) for being a real model
 #[derive(Clone)]
 pub struct HifiGanSoftVcModelCheckCommand {
-  /// Where the HifiGan code lives
+  /// Where the HifiGan (softvc version) code lives
   hifigan_softvc_root_code_directory: String,
   
   /// eg. `source python/bin/activate`
@@ -21,15 +21,15 @@ pub struct HifiGanSoftVcModelCheckCommand {
 
 impl HifiGanSoftVcModelCheckCommand {
   pub fn new(
-    hifigan_root_code_directory: &str,
+    hifigan_softvc_root_code_directory: &str,
     maybe_virtual_env_activation_command: Option<&str>,
-    hifigan_model_check_script_name: &str,
+    hifigan_softvc_model_check_script_name: &str,
     maybe_docker_options: Option<DockerOptions>,
   ) -> Self {
     Self {
-      hifigan_softvc_root_code_directory: hifigan_root_code_directory.to_string(),
+      hifigan_softvc_root_code_directory: hifigan_softvc_root_code_directory.to_string(),
       maybe_virtual_env_activation_command: maybe_virtual_env_activation_command.map(|s| s.to_string()),
-      hifigan_softvc_model_check_script_name: hifigan_model_check_script_name.to_string(),
+      hifigan_softvc_model_check_script_name: hifigan_softvc_model_check_script_name.to_string(),
       maybe_docker_options,
     }
   }

crates/service/job/download_job/src/job_types/mod.rs

@@ -1,4 +1,5 @@
 pub mod hifigan;
 pub mod hifigan_softvc;
+pub mod softvc;
 pub mod tacotron;
 pub mod talknet;

crates/service/job/download_job/src/job_types/softvc/mod.rs

@@ -0,0 +1,2 @@
+pub mod softvc_model_check_command;
+pub mod process_softvc_model;

crates/service/job/download_job/src/job_types/softvc/process_softvc_model.rs

@@ -0,0 +1,140 @@
+use anyhow::anyhow;
+use container_common::anyhow_result::AnyhowResult;
+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 container_common::hashing::hash_file_sha2::hash_file_sha2;
+use crate::JobState;
+use crate::job_steps::job_results::JobResults;
+use database_queries::queries::generic_download::job::list_available_generic_download_jobs::AvailableDownloadJob;
+use database_queries::queries::tts::tts_models::insert_tts_model_from_download_job::insert_tts_model_from_download_job;
+use database_queries::queries::tts::tts_models::insert_tts_model_from_download_job;
+use jobs_common::redis_job_status_logger::RedisJobStatusLogger;
+use log::{info, warn};
+use reusable_types::db::enums::entity_visibility::EntityVisibility;
+use std::fs::File;
+use std::io::Read;
+use std::path::PathBuf;
+use tempdir::TempDir;
+
+/// Returns the token of the entity.
+pub async fn process_softvc_model<'a, 'b>(
+  job_state: &JobState,
+  job: &AvailableDownloadJob,
+  temp_dir: &TempDir,
+  download_filename: &str,
+  redis_logger: &'a mut RedisJobStatusLogger<'b>,
+) -> AnyhowResult<JobResults> {
+
+  // ==================== RUN MODEL CHECK ==================== //
+
+  info!("Checking that softvc model is valid...");
+
+  redis_logger.log_status("checking rocket vc model")?; // NB: "rocket vc" is a codename
+
+  let file_path = PathBuf::from(download_filename.clone());
+
+  let output_metadata_fs_path = temp_dir.path().join("metadata.json");
+
+  let model_check_result = job_state.sidecar_configs.softvc_model_check_command.execute(
+    &file_path,
+    &output_metadata_fs_path,
+  );
+
+  if let Err(e) = model_check_result {
+    safe_delete_temp_file(&file_path);
+    safe_delete_temp_directory(&temp_dir);
+    return Err(anyhow!("model check error: {:?}", e));
+  }
+
+  // ==================== CHECK ALL FILES EXIST AND GET METADATA ==================== //
+
+  info!("Checking that metadata output file exists...");
+
+  check_file_exists(&output_metadata_fs_path)?;
+
+  let file_metadata = match read_metadata_file(&output_metadata_fs_path) {
+    Ok(metadata) => metadata,
+    Err(e) => {
+      safe_delete_temp_file(&file_path);
+      safe_delete_temp_file(&output_metadata_fs_path);
+      safe_delete_temp_directory(&temp_dir);
+      return Err(e);
+    }
+  };
+
+  // ==================== UPLOAD MODEL FILE ==================== //
+
+  info!("Uploading SoftVC model to GCS...");
+
+  let private_bucket_hash = hash_file_sha2(&download_filename)?;
+
+  info!("File hash: {}", private_bucket_hash);
+
+  let model_bucket_path = job_state.bucket_path_unifier.softvc_model_path(&private_bucket_hash);
+
+  info!("Destination bucket path: {:?}", &model_bucket_path);
+
+  redis_logger.log_status("uploading rocket vc  model")?; // NB: "rocket vc" is a codename
+
+  if let Err(e) = job_state.bucket_client.upload_filename(&model_bucket_path, &file_path).await {
+    safe_delete_temp_file(&output_metadata_fs_path);
+    safe_delete_temp_file(&file_path);
+    safe_delete_temp_directory(&temp_dir);
+    return Err(e);
+  }
+
+  // ==================== DELETE DOWNLOADED FILE ==================== //
+
+  // NB: We should be using a tempdir, but to make absolutely certain we don't overflow the disk...
+  safe_delete_temp_file(&output_metadata_fs_path);
+  safe_delete_temp_file(&file_path);
+  safe_delete_temp_directory(&temp_dir);
+
+  // ==================== SAVE RECORDS ==================== //
+
+  info!("Saving Soft VC record...");
+
+  // TODO FIX BELOW QUERIES FOR SOFTVC --
+  // TODO FIX BELOW QUERIES FOR SOFTVC --
+  // TODO FIX BELOW QUERIES FOR SOFTVC --
+  // TODO FIX BELOW QUERIES FOR SOFTVC --
+
+  let (_id, model_token) = insert_tts_model_from_download_job(insert_tts_model_from_download_job::Args {
+    title: &job.title,
+    original_download_url: &job.download_url,
+    original_filename: &download_filename,
+    file_size_bytes: file_metadata.file_size_bytes,
+    creator_user_token: &job.creator_user_token,
+    creator_ip_address: &job.creator_ip_address,
+    creator_set_visibility: EntityVisibility::Public, // TODO: All models default to public at start
+    private_bucket_hash: &private_bucket_hash,
+    private_bucket_object_name: &model_bucket_path,
+    mysql_pool: &job_state.mysql_pool,
+  }).await?;
+
+  job_state.badge_granter.maybe_grant_tts_model_uploads_badge(&job.creator_user_token)
+      .await
+      .map_err(|e| {
+        warn!("error maybe awarding badge: {:?}", e);
+        anyhow!("error maybe awarding badge")
+      })?;
+
+  Ok(JobResults {
+    entity_token: Some(model_token),
+    entity_type: Some("tacotron2".to_string()), // NB: This may be different from `GenericDownloadType` in the future!
+  })
+}
+
+#[derive(Deserialize)]
+struct FileMetadata {
+  pub file_size_bytes: u64,
+}
+
+fn read_metadata_file(filename: &PathBuf) -> AnyhowResult<FileMetadata> {
+  let mut file = File::open(filename)?;
+  let mut buffer = String::new();
+  file.read_to_string(&mut buffer)?;
+  Ok(serde_json::from_str(&buffer)?)
+}
+

crates/service/job/download_job/src/job_types/softvc/softvc_model_check_command.rs

@@ -0,0 +1,87 @@
+use container_common::anyhow_result::AnyhowResult;
+use log::info;
+use std::fs::OpenOptions;
+use std::path::Path;
+use subprocess::{Popen, PopenConfig, Redirection};
+use subprocess_common::docker_options::DockerOptions;
+
+/// This command is used to check softvc for being a real model
+#[derive(Clone)]
+pub struct SoftVcModelCheckCommand {
+  /// Where the Tacotron code lives
+  softvc_root_code_directory: String,
+  
+  /// eg. `source python/bin/activate`
+  maybe_virtual_env_activation_command: Option<String>,
+  
+  softvc_model_check_script_name: String,
+
+  /// If this is run under Docker (eg. in development), these are the options.
+  maybe_docker_options: Option<DockerOptions>,
+}
+
+impl SoftVcModelCheckCommand {
+  pub fn new(
+    softvc_root_code_directory: &str,
+    maybe_virtual_env_activation_command: Option<&str>,
+    softvc_model_check_script_name: &str,
+    maybe_docker_options: Option<DockerOptions>,
+  ) -> Self {
+    Self {
+      softvc_root_code_directory: softvc_root_code_directory.to_string(),
+      maybe_virtual_env_activation_command: maybe_virtual_env_activation_command.map(|s| s.to_string()),
+      softvc_model_check_script_name: softvc_model_check_script_name.to_string(),
+      maybe_docker_options,
+    }
+  }
+
+  pub fn execute<P: AsRef<Path>>(
+    &self,
+    synthesizer_checkpoint_path: P,
+    output_metadata_filename: P,
+  ) -> AnyhowResult<()> {
+    let mut command = String::new();
+
+    command.push_str("echo 'test'");
+    command.push_str(" && ");
+    command.push_str(&format!("cd {}", self.softvc_root_code_directory));
+
+    if let Some(venv_command) = self.maybe_virtual_env_activation_command.as_deref() {
+      command.push_str(" && ");
+      command.push_str(venv_command);
+      command.push_str(" ");
+    }
+
+    command.push_str(" && ");
+    command.push_str("python ");
+    command.push_str(&self.softvc_model_check_script_name);
+    command.push_str(" --synthesizer_checkpoint_path ");
+    command.push_str(&synthesizer_checkpoint_path.as_ref().display().to_string());
+    command.push_str(" --output_metadata_filename ");
+    command.push_str(&output_metadata_filename.as_ref().display().to_string());
+
+    if let Some(docker_options) = self.maybe_docker_options.as_ref() {
+      command = docker_options.to_command_string(&command);
+    }
+
+    info!("Command: {:?}", command);
+
+    let command_parts = [
+      "bash",
+      "-c",
+      &command
+    ];
+
+    let mut p = Popen::create(&command_parts, PopenConfig {
+      ..Default::default()
+    })?;
+
+    info!("Pid : {:?}", p.pid());
+
+    let exit_status = p.wait()?;
+
+    info!("Exit status: {:?}", exit_status);
+
+    Ok(())
+  }
+}