diff --git a/crates/lib/mysql_queries/src/queries/voice_conversion/models/insert_voice_conversion_model_from_download_job.rs b/crates/lib/mysql_queries/src/queries/voice_conversion/models/insert_voice_conversion_model_from_download_job.rs index 121febc5dc..c47232f800 100644 --- a/crates/lib/mysql_queries/src/queries/voice_conversion/models/insert_voice_conversion_model_from_download_job.rs +++ b/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> { +pub struct InsertVoiceConversionModelArgs<'a, P: AsRef> { + pub model_type: VoiceConversionModelType, + pub title: &'a str, pub original_download_url: &'a str, @@ -25,7 +28,7 @@ pub struct Args<'a, P: AsRef> { pub async fn insert_voice_conversion_model_from_download_job>( - 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>( 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, diff --git a/crates/service/job/download_job/src/job_types/dispatch_job_to_handler.rs b/crates/service/job/download_job/src/job_types/dispatch_job_to_handler.rs index 48679ddbba..d0e7e8a52a 100644 --- a/crates/service/job/download_job/src/job_types/dispatch_job_to_handler.rs +++ b/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, diff --git a/crates/service/job/download_job/src/job_types/voice_conversion/so_vits_svc/mod.rs b/crates/service/job/download_job/src/job_types/voice_conversion/so_vits_svc/mod.rs index e69de29bb2..c7b0e376fa 100644 --- a/crates/service/job/download_job/src/job_types/voice_conversion/so_vits_svc/mod.rs +++ b/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; \ No newline at end of file diff --git a/crates/service/job/download_job/src/job_types/voice_conversion/so_vits_svc/process_so_vits_svc_model.rs b/crates/service/job/download_job/src/job_types/voice_conversion/so_vits_svc/process_so_vits_svc_model.rs new file mode 100644 index 0000000000..2f0c572349 --- /dev/null +++ b/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 { + + // ==================== 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! + }) +} diff --git a/crates/service/job/download_job/src/job_types/voice_conversion/so_vits_svc/so_vits_svc_model_check_command.rs b/crates/service/job/download_job/src/job_types/voice_conversion/so_vits_svc/so_vits_svc_model_check_command.rs new file mode 100644 index 0000000000..c0f3287a06 --- /dev/null +++ b/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, + + /// eg. `python3` + maybe_override_python_interpreter: Option, + + /// If this is run under Docker (eg. in development), these are the options. + maybe_docker_options: Option, +} + +pub enum Device { + Cuda, + Cpu, +} + +pub struct CheckArgs<'a, P: AsRef> { + /// --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>( + 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, + ) -> 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>( + &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(()) + } +} diff --git a/crates/service/job/download_job/src/job_types/voice_conversion/softvc/process_softvc_model.rs b/crates/service/job/download_job/src/job_types/voice_conversion/softvc/process_softvc_model.rs index 7eb2de3c88..224709ddda 100644 --- a/crates/service/job/download_job/src/job_types/voice_conversion/softvc/process_softvc_model.rs +++ b/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,