diff --git a/Cargo.lock b/Cargo.lock index 07c5b96ee3..d7cdf2624d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2831,6 +2831,7 @@ dependencies = [ "actix-web 4.0.1", "actix_helpers", "anyhow", + "async-trait", "bootstrap", "buckets", "chrono", diff --git a/crates/service/job/inference_job/Cargo.toml b/crates/service/job/inference_job/Cargo.toml index d1976b27f4..a9787a1bb1 100644 --- a/crates/service/job/inference_job/Cargo.toml +++ b/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" diff --git a/crates/service/job/inference_job/src/job/job_types/lipsync/process_single_lipsync_job.rs b/crates/service/job/inference_job/src/job/job_types/lipsync/process_single_lipsync_job.rs index 3bfb76d5dd..1069942ac5 100644 --- a/crates/service/job/inference_job/src/job/job_types/lipsync/process_single_lipsync_job.rs +++ b/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 { 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) } - diff --git a/crates/service/job/inference_job/src/job/job_types/lipsync/sad_talker/mod.rs b/crates/service/job/inference_job/src/job/job_types/lipsync/sad_talker/mod.rs index e69de29bb2..dd6905c913 100644 --- a/crates/service/job/inference_job/src/job/job_types/lipsync/sad_talker/mod.rs +++ b/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; \ No newline at end of file diff --git a/crates/service/job/inference_job/src/job/job_types/lipsync/sad_talker/model_downloaders.rs b/crates/service/job/inference_job/src/job/job_types/lipsync/sad_talker/model_downloaders.rs new file mode 100644 index 0000000000..e69de29bb2 diff --git a/crates/service/job/inference_job/src/job/job_types/lipsync/sad_talker/process_job.rs b/crates/service/job/inference_job/src/job/job_types/lipsync/sad_talker/process_job.rs new file mode 100644 index 0000000000..87b595a764 --- /dev/null +++ b/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 { + 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 = 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, + pub mimetype: Option, + pub file_size_bytes: u64, +} diff --git a/crates/service/job/inference_job/src/util/from_env.rs b/crates/service/job/inference_job/src/util/from_env.rs new file mode 100644 index 0000000000..8b13789179 --- /dev/null +++ b/crates/service/job/inference_job/src/util/from_env.rs @@ -0,0 +1 @@ + diff --git a/crates/service/job/inference_job/src/util/mod.rs b/crates/service/job/inference_job/src/util/mod.rs index 63f98be1ca..26b016bb6e 100644 --- a/crates/service/job/inference_job/src/util/mod.rs +++ b/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; diff --git a/crates/service/job/inference_job/src/util/model_downloader.rs b/crates/service/job/inference_job/src/util/model_downloader.rs new file mode 100644 index 0000000000..5a5a3a631e --- /dev/null +++ b/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; +// } + + +#[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 + } + } + } +}