wip so-vits-svc

crates/lib/mysql_queries/src/queries/voice_conversion/models/insert_voice_conversion_model_from_download_job.rs

@@ -4,9 +4,12 @@ use errors::AnyhowResult;
 use enums::common::visibility::Visibility;
 use sqlx::MySqlPool;
 use std::path::Path;
+use enums::by_table::voice_conversion_models::voice_conversion_model_type::VoiceConversionModelType;
 use tokens::voice_conversion::model::VoiceConversionModelToken;
 
-pub struct Args<'a, P: AsRef<Path>> {
+pub struct InsertVoiceConversionModelArgs<'a, P: AsRef<Path>> {
+  pub model_type: VoiceConversionModelType,
+
   pub title: &'a str,
 
   pub original_download_url: &'a str,
@@ -25,7 +28,7 @@ pub struct Args<'a, P: AsRef<Path>> {
 
 
 pub async fn insert_voice_conversion_model_from_download_job<P: AsRef<Path>>(
-  args: Args<'_, P>,
+  args: InsertVoiceConversionModelArgs<'_, P>,
 ) -> AnyhowResult<(u64, VoiceConversionModelToken)> {
 
   let model_token = VoiceConversionModelToken::generate();
@@ -41,7 +44,7 @@ pub async fn insert_voice_conversion_model_from_download_job<P: AsRef<Path>>(
 INSERT INTO voice_conversion_models
 SET
   token = ?,
-  model_type = "rocket_vc",
+  model_type = ?,
   title = ?,
   description_markdown = '',
   description_rendered_html = '',
@@ -54,6 +57,7 @@ SET
   file_size_bytes = ?
         "#,
       &model_token,
+      args.model_type.to_str(),
       args.title,
       args.creator_user_token,
       args.creator_ip_address,

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

@@ -51,6 +51,9 @@ pub async fn dispatch_job_to_handler<'a, 'b: 'a>(args: DispatchJobToHandlerArgs<
       entity_token = results.entity_token.clone();
       entity_type = results.entity_type.clone();
     }
+    GenericDownloadType::HifiGanSoVitsSvc => {
+      // TODO
+    }
     GenericDownloadType::RocketVc => {
       let results = process_softvc_model(
         args.job_runner_state,
@@ -62,6 +65,17 @@ pub async fn dispatch_job_to_handler<'a, 'b: 'a>(args: DispatchJobToHandlerArgs<
       entity_token = results.entity_token.clone();
       entity_type = results.entity_type.clone();
     }
+    GenericDownloadType::SoVitsSvc => {
+      let results = process_so_vits_svc_model(
+        args.job_runner_state,
+        args.job,
+        args.temp_dir,
+        args.download_filename,
+        args.redis_logger,
+      ).await?;
+      entity_token = results.entity_token.clone();
+      entity_type = results.entity_type.clone();
+    }
     GenericDownloadType::Tacotron2 => {
       let results = process_tacotron_model(
         args.job_runner_state,

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

@@ -0,0 +1,2 @@
+pub mod process_so_vits_svc_model;
+pub mod so_vits_svc_model_check_command;

crates/service/job/download_job/src/job_types/voice_conversion/so_vits_svc/process_so_vits_svc_model.rs

@@ -0,0 +1,162 @@
+use anyhow::anyhow;
+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::JobState;
+use crate::job_loop::job_results::JobResults;
+use crate::job_types::tts::vits::vits_model_check_command::{CheckArgs, Device};
+use enums::by_table::tts_models::tts_model_type::TtsModelType;
+use enums::common::visibility::Visibility;
+use errors::AnyhowResult;
+use filesys::filename_concat::filename_concat;
+use hashing::sha256::sha256_hash_file::sha256_hash_file;
+use jobs_common::redis_job_status_logger::RedisJobStatusLogger;
+use log::{error, info, warn};
+use mysql_queries::queries::generic_download::job::list_available_generic_download_jobs::AvailableDownloadJob;
+use mysql_queries::queries::tts::tts_models::insert_tts_model_from_download_job::insert_tts_model_from_download_job;
+use mysql_queries::queries::tts::tts_models::insert_tts_model_from_download_job;
+use std::fs::File;
+use std::io::Read;
+use std::path::{Path, PathBuf};
+use std::time::Duration;
+use tempdir::TempDir;
+use enums::by_table::voice_conversion_models::voice_conversion_model_type::VoiceConversionModelType;
+use mysql_queries::queries::voice_conversion::models::insert_voice_conversion_model_from_download_job::{insert_voice_conversion_model_from_download_job, InsertVoiceConversionModelArgs};
+
+/// Returns the token of the entity.
+pub async fn process_so_vits_svc_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 model is valid...");
+
+  redis_logger.log_status("checking so-vits-svc model")?;
+
+  let original_model_file_path = PathBuf::from(download_filename.clone());
+
+  // NB: We'll be creating the traced model in the "check" step and uploading it to GDrive along with the original.
+  let traced_model_file_path = PathBuf::from(filename_concat(&original_model_file_path, "_traced"));
+
+  let config_path = PathBuf::from("configs/ljs_li44_tmbert_nmp_s1_arpa.json"); // TODO: This could be variable.
+
+  let model_check_result = job_state.sidecar_configs.so_vits_svc_model_check_command.execute_check(CheckArgs {
+    traced_model_output_path: &traced_model_file_path,
+    model_checkpoint_path: &original_model_file_path,
+    config_path: &config_path,
+    device: Device::Cuda,
+    test_string: "this is a test of model inference",
+  });
+
+  if let Err(e) = model_check_result {
+    safe_delete_temp_file(&original_model_file_path);
+    safe_delete_temp_file(&traced_model_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 output traced model file exists...");
+
+  check_file_exists(&traced_model_file_path)?;
+
+  // TODO: get model sizes
+
+  // 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 ORIGINAL MODEL FILE ==================== //
+
+  info!("Uploading VITS TTS model to GCS...");
+
+  let private_bucket_hash = sha256_hash_file(&download_filename)?;
+
+  info!("File hash: {}", private_bucket_hash);
+
+  let model_bucket_path = job_state.bucket_path_unifier.tts_synthesizer_path(&private_bucket_hash);
+
+  info!("Destination bucket path: {:?}", &model_bucket_path);
+
+  redis_logger.log_status("uploading VITS TTS model")?;
+
+  if let Err(err) = job_state.bucket_client.upload_filename(&model_bucket_path, &original_model_file_path).await {
+    error!("Problem uploading original model: {:?}", err);
+    error!(" - Model file: {:?}", &original_model_file_path);
+    error!(" - Traced model file: {:?}", &traced_model_file_path);
+    safe_delete_temp_file(&original_model_file_path);
+    safe_delete_temp_file(&traced_model_file_path);
+    safe_delete_temp_directory(&temp_dir);
+    return Err(err);
+  }
+
+  // ==================== UPLOAD TRACED MODEL FILE ==================== //
+
+  info!("Uploading VITS TTS (traced) model to GCS...");
+
+  let traced_model_bucket_path = job_state.bucket_path_unifier.tts_traced_synthesizer_path(&private_bucket_hash);
+
+  info!("Destination bucket path: {:?}", &traced_model_bucket_path);
+
+  redis_logger.log_status("uploading VITS TTS (traced) model")?;
+
+  if let Err(err) = job_state.bucket_client.upload_filename(&traced_model_bucket_path, &traced_model_file_path).await {
+    error!("Problem uploading traced model: {:?}", err);
+    error!(" - Model file: {:?}", &original_model_file_path);
+    error!(" - Traced model file: {:?}", &traced_model_file_path);
+    safe_delete_temp_file(&original_model_file_path);
+    safe_delete_temp_file(&traced_model_file_path);
+    safe_delete_temp_directory(&temp_dir);
+    return Err(err);
+  }
+
+  // ==================== DELETE DOWNLOADED FILE ==================== //
+
+  // NB: We should be using a tempdir, but to make absolutely certain we don't overflow the disk...
+  info!("Done uploading; deleting temporary files and paths...");
+  safe_delete_temp_file(&original_model_file_path);
+  safe_delete_temp_file(&traced_model_file_path);
+  safe_delete_temp_directory(&temp_dir);
+
+  // ==================== SAVE RECORDS ==================== //
+
+  info!("Saving TTS model record...");
+
+  let (_id, model_token) = insert_voice_conversion_model_from_download_job(InsertVoiceConversionModelArgs {
+    model_type: VoiceConversionModelType::SoVitsSvc,
+    title: &job.title,
+    original_download_url: &job.download_url,
+    original_filename: &download_filename,
+    file_size_bytes: 0, // TODO: Get file sizes!
+    creator_user_token: &job.creator_user_token,
+    creator_ip_address: &job.creator_ip_address,
+    creator_set_visibility: Visibility::Public, // TODO: All models default to public at start
+    private_bucket_hash: &private_bucket_hash,
+    private_bucket_object_name: "", // TODO: This should go away.
+    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.to_string()),
+    entity_type: Some(VoiceConversionModelType::SoVitsSvc.to_string()), // NB: This may be different from `GenericDownloadType` in the future!
+  })
+}

crates/service/job/download_job/src/job_types/voice_conversion/so_vits_svc/so_vits_svc_model_check_command.rs

@@ -0,0 +1,140 @@
+use container_common::anyhow_result::AnyhowResult;
+use filesys::path_to_string::path_to_string;
+use log::info;
+use std::fs::OpenOptions;
+use std::path::{Path, PathBuf};
+use subprocess::{Popen, PopenConfig, Redirection};
+use subprocess_common::docker_options::DockerOptions;
+
+/// This command is used to check tacotron for being a real model
+#[derive(Clone)]
+pub struct SoVitsSvcModelCheckCommand {
+  /// Where the so-vits-svc code lives
+  so_vits_svc_root_code_directory: PathBuf,
+
+  /// The name of the check/process script, eg. `export_ts.py`
+  check_script_name: PathBuf,
+
+  /// eg. `source python/bin/activate`
+  maybe_virtual_env_activation_command: Option<String>,
+
+  /// eg. `python3`
+  maybe_override_python_interpreter: Option<String>,
+
+  /// If this is run under Docker (eg. in development), these are the options.
+  maybe_docker_options: Option<DockerOptions>,
+}
+
+pub enum Device {
+  Cuda,
+  Cpu,
+}
+
+pub struct CheckArgs<'a, P: AsRef<Path>> {
+  /// --out-path: output path of the traced model
+  pub traced_model_output_path: P,
+
+  /// --checkpoint: input path of the model checkpoint
+  pub model_checkpoint_path: P,
+
+  /// --config: path of the hparams json file
+  pub config_path: P,
+
+  /// --device: cpu or cuda
+  pub device: Device,
+
+  /// --test-string: string to infer when tracing (I think)
+  pub test_string: &'a str,
+}
+
+impl SoVitsSvcModelCheckCommand {
+  pub fn new<P: AsRef<Path>>(
+    so_vits_svc_root_code_directory: P,
+    check_script_name: P,
+    maybe_override_python_interpreter: Option<&str>,
+    maybe_virtual_env_activation_command: Option<&str>,
+    maybe_docker_options: Option<DockerOptions>,
+  ) -> Self {
+    Self {
+      so_vits_svc_root_code_directory: so_vits_svc_root_code_directory.as_ref().to_path_buf(),
+      check_script_name: check_script_name.as_ref().to_path_buf(),
+      maybe_virtual_env_activation_command: maybe_virtual_env_activation_command.map(|s| s.to_string()),
+      maybe_override_python_interpreter: maybe_override_python_interpreter.map(|s| s.to_string()),
+      maybe_docker_options,
+    }
+  }
+
+  pub fn execute_check<P: AsRef<Path>>(
+    &self,
+    args: CheckArgs<'_, P>,
+  ) -> AnyhowResult<()> {
+
+    let mut command = String::new();
+    command.push_str(&format!("cd {}", path_to_string(&self.so_vits_svc_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(" ");
+    }
+
+    let python_binary = self.maybe_override_python_interpreter
+        .as_deref()
+        .unwrap_or("python");
+
+    command.push_str(" && ");
+    command.push_str(python_binary);
+    command.push_str(" ");
+    command.push_str(&path_to_string(&self.check_script_name));
+
+    // ===== Begin Python Args =====
+
+    command.push_str(" --out-path ");
+    command.push_str(&path_to_string(args.traced_model_output_path));
+
+    command.push_str(" --checkpoint ");
+    command.push_str(&path_to_string(args.model_checkpoint_path));
+
+    command.push_str(" --config ");
+    command.push_str(&path_to_string(args.config_path));
+
+    let device = match args.device {
+      Device::Cuda => "cuda",
+      Device::Cpu => "cpu",
+    };
+
+    command.push_str(" --device ");
+    command.push_str(&path_to_string(device));
+
+    command.push_str(" --test-string ");
+    command.push_str("'");
+    command.push_str(&path_to_string(args.test_string));
+    command.push_str("'");
+
+    // ===== End Python Args =====
+
+    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!("Subprocess PID: {:?}", p.pid());
+
+    let exit_status = p.wait()?;
+
+    info!("Subprocess exit status: {:?}", exit_status);
+
+    Ok(())
+  }
+}

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

@@ -10,7 +10,7 @@ use hashing::sha256::sha256_hash_file::sha256_hash_file;
 use jobs_common::redis_job_status_logger::RedisJobStatusLogger;
 use log::{info, warn};
 use mysql_queries::queries::generic_download::job::list_available_generic_download_jobs::AvailableDownloadJob;
-use mysql_queries::queries::voice_conversion::models::insert_voice_conversion_model_from_download_job::{Args, insert_voice_conversion_model_from_download_job};
+use mysql_queries::queries::voice_conversion::models::insert_voice_conversion_model_from_download_job::{InsertVoiceConversionModelArgs, insert_voice_conversion_model_from_download_job};
 use std::fs::File;
 use std::io::Read;
 use std::path::PathBuf;
@@ -94,7 +94,7 @@ pub async fn process_softvc_model<'a, 'b>(
 
   info!("Saving Soft VC record...");
 
-  let (_id, model_token) = insert_voice_conversion_model_from_download_job(Args {
+  let (_id, model_token) = insert_voice_conversion_model_from_download_job(InsertVoiceConversionModelArgs {
     title: &job.title,
     original_download_url: &job.download_url,
     original_filename: &download_filename,