From eb0e05f61e862c8ba69f2dfb6fc93e9f310aeefd Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 14 Sep 2023 04:58:28 -0400 Subject: [PATCH] download image and audio --- .../src/job/job_loop/main_loop.rs | 3 + .../job/job_loop/process_single_job_error.rs | 3 + .../lipsync/sad_talker/download_audio_file.rs | 108 ++++++++++++++++ .../lipsync/sad_talker/download_image_file.rs | 81 ++++++++++++ .../job/job_types/lipsync/sad_talker/mod.rs | 2 + .../lipsync/sad_talker/process_job.rs | 120 +++++------------- 6 files changed, 231 insertions(+), 86 deletions(-) create mode 100644 crates/service/job/inference_job/src/job/job_types/lipsync/sad_talker/download_audio_file.rs create mode 100644 crates/service/job/inference_job/src/job/job_types/lipsync/sad_talker/download_image_file.rs diff --git a/crates/service/job/inference_job/src/job/job_loop/main_loop.rs b/crates/service/job/inference_job/src/job/job_loop/main_loop.rs index 4847e6033e..7d12fdd596 100644 --- a/crates/service/job/inference_job/src/job/job_loop/main_loop.rs +++ b/crates/service/job/inference_job/src/job/job_loop/main_loop.rs @@ -130,6 +130,8 @@ async fn process_job_batch(job_dependencies: &JobDependencies, jobs: Vec (true, false, format!("InvalidJob: {:?}", err), Some("invalid job")), + ProcessSingleJobError::NotYetImplemented => + (true, false, "not yet implemented".to_string(), Some("not yet implemented")), // Non-permanent failures ProcessSingleJobError::FilesystemFull => @@ -171,6 +173,7 @@ async fn process_job_batch(job_dependencies: &JobDependencies, jobs: Vec {} ProcessSingleJobError::InvalidJob(_) => {} ProcessSingleJobError::KeepAliveElapsed => {} + ProcessSingleJobError::NotYetImplemented => {} } } } diff --git a/crates/service/job/inference_job/src/job/job_loop/process_single_job_error.rs b/crates/service/job/inference_job/src/job/job_loop/process_single_job_error.rs index 662bc8755b..5cc24ad7a7 100644 --- a/crates/service/job/inference_job/src/job/job_loop/process_single_job_error.rs +++ b/crates/service/job/inference_job/src/job/job_loop/process_single_job_error.rs @@ -15,6 +15,9 @@ pub enum ProcessSingleJobError { /// This is any other kind of error. /// It might be important, we just haven't special cased it yet. Other(anyhow::Error), + /// We hit a feature or a path for this job that has not yet been implemented. + /// Permanently fail the job. + NotYetImplemented, } impl ProcessSingleJobError { diff --git a/crates/service/job/inference_job/src/job/job_types/lipsync/sad_talker/download_audio_file.rs b/crates/service/job/inference_job/src/job/job_types/lipsync/sad_talker/download_audio_file.rs new file mode 100644 index 0000000000..3fb69f9112 --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/lipsync/sad_talker/download_audio_file.rs @@ -0,0 +1,108 @@ +use anyhow::anyhow; +use buckets::public::media_uploads::original_file::MediaUploadOriginalFilePath; +use buckets::public::voice_conversion_results::original_file::VoiceConversionResultOriginalFilePath; +use cloud_storage::bucket_client::BucketClient; +use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; +use crate::util::maybe_download_file_from_bucket::maybe_download_file_from_bucket; +use crate::util::scoped_temp_dir_creator::ScopedTempDirCreator; +use jobs_common::job_progress_reporter::job_progress_reporter::JobProgressReporter; +use log::{error, warn}; +use mysql_queries::payloads::generic_inference_args::lipsync_payload::LipsyncAnimationAudioSource; +use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; +use mysql_queries::queries::media_uploads::get_media_upload_for_inference::get_media_upload_for_inference; +use mysql_queries::queries::voice_conversion::results::get_voice_conversion_result_for_inference::get_voice_conversion_result_for_inference; +use sqlx::MySqlPool; +use std::path::PathBuf; +use tempdir::TempDir; +use tokens::files::media_upload::MediaUploadToken; + +pub struct AudioFile { + pub filesystem_path: PathBuf, +} + +pub async fn download_audio_file( + audio_source: &LipsyncAnimationAudioSource, + public_bucket_client: &BucketClient, + job_progress_reporter: &mut Box, + job: &AvailableInferenceJob, + temp_dir_creator: &ScopedTempDirCreator, + work_temp_dir: &TempDir, + mysql_pool: &MySqlPool, +) -> Result { + + let bucket_object_path; + + match audio_source { + LipsyncAnimationAudioSource::F(media_file_token) => { + // TODO(bt, 2023-09-08): Implement + return Err(ProcessSingleJobError::NotYetImplemented) + } + LipsyncAnimationAudioSource::U(media_upload_token) => { + let media_upload_token = MediaUploadToken::new_from_str(media_upload_token); + let media_upload_result = get_media_upload_for_inference( + &media_upload_token, + mysql_pool + ).await; + + let media_upload_result = match media_upload_result { + Ok(Some(result)) => result, + Ok(None) => { + warn!("could not find media upload: {:?}", media_upload_token); + return Err(ProcessSingleJobError::from_anyhow_error( + anyhow!("could not find media upload: {:?}", media_upload_token))) + } + Err(e) => { + error!("could not query media upload: {:?}", e); + return Err(ProcessSingleJobError::from_anyhow_error(e)) + } + }; + + let media_upload_bucket_path = MediaUploadOriginalFilePath::from_object_hash(&media_upload_result.public_bucket_directory_hash); + bucket_object_path = media_upload_bucket_path.to_full_object_pathbuf(); + } + LipsyncAnimationAudioSource::T(tts_result_token) => { + // TODO(bt, 2023-09-08): Implement + return Err(ProcessSingleJobError::NotYetImplemented) + } + LipsyncAnimationAudioSource::V(voice_conversion_result_token) => { + let voice_conversion_result = get_voice_conversion_result_for_inference( + voice_conversion_result_token, + false, + &mysql_pool, + ).await; + + let voice_conversion_result = match voice_conversion_result { + Ok(Some(result)) => result, + Ok(None) => { + warn!("could not find voice conversion result: {:?}", voice_conversion_result_token); + return Err(ProcessSingleJobError::from_anyhow_error( + anyhow!("could not find voice conversion result: {:?}", voice_conversion_result_token))) + } + Err(e) => { + error!("could not query voice conversion result: {:?}", e); + return Err(ProcessSingleJobError::from_anyhow_error(e)) + } + }; + + let result_bucket_path = VoiceConversionResultOriginalFilePath::from_object_hash(&voice_conversion_result.public_bucket_hash); + bucket_object_path = result_bucket_path.to_full_object_pathbuf(); + } + } + + let downloaded_filesystem_path = work_temp_dir.path().join("audio.bin"); + + maybe_download_file_from_bucket( + "audio file", + &downloaded_filesystem_path, + &bucket_object_path, + public_bucket_client, + job_progress_reporter, + "downloading", + job.id.0, + &temp_dir_creator, + ).await?; + + Ok(AudioFile { + filesystem_path: downloaded_filesystem_path, + }) +} diff --git a/crates/service/job/inference_job/src/job/job_types/lipsync/sad_talker/download_image_file.rs b/crates/service/job/inference_job/src/job/job_types/lipsync/sad_talker/download_image_file.rs new file mode 100644 index 0000000000..4d57032883 --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/lipsync/sad_talker/download_image_file.rs @@ -0,0 +1,81 @@ +use anyhow::anyhow; +use buckets::public::media_uploads::original_file::MediaUploadOriginalFilePath; +use buckets::public::voice_conversion_results::original_file::VoiceConversionResultOriginalFilePath; +use cloud_storage::bucket_client::BucketClient; +use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; +use crate::util::maybe_download_file_from_bucket::maybe_download_file_from_bucket; +use crate::util::scoped_temp_dir_creator::ScopedTempDirCreator; +use jobs_common::job_progress_reporter::job_progress_reporter::JobProgressReporter; +use log::{error, warn}; +use mysql_queries::payloads::generic_inference_args::lipsync_payload::{LipsyncAnimationAudioSource, LipsyncAnimationImageSource}; +use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; +use mysql_queries::queries::media_uploads::get_media_upload_for_inference::get_media_upload_for_inference; +use mysql_queries::queries::voice_conversion::results::get_voice_conversion_result_for_inference::get_voice_conversion_result_for_inference; +use sqlx::MySqlPool; +use std::path::PathBuf; +use tempdir::TempDir; +use tokens::files::media_upload::MediaUploadToken; + +pub struct ImageFile { + pub filesystem_path: PathBuf, +} + +pub async fn download_image_file( + image_source: &LipsyncAnimationImageSource, + public_bucket_client: &BucketClient, + job_progress_reporter: &mut Box, + job: &AvailableInferenceJob, + temp_dir_creator: &ScopedTempDirCreator, + work_temp_dir: &TempDir, + mysql_pool: &MySqlPool, +) -> Result { + + let bucket_object_path; + + match image_source { + LipsyncAnimationImageSource::F(media_file_token) => { + // TODO(bt, 2023-09-08): Implement + return Err(ProcessSingleJobError::NotYetImplemented) + } + LipsyncAnimationImageSource::U(media_upload_token) => { + let media_upload_token = MediaUploadToken::new_from_str(media_upload_token); + let media_upload_result = get_media_upload_for_inference( + &media_upload_token, + mysql_pool + ).await; + + let media_upload_result = match media_upload_result { + Ok(Some(result)) => result, + Ok(None) => { + warn!("could not find media upload: {:?}", media_upload_token); + return Err(ProcessSingleJobError::from_anyhow_error( + anyhow!("could not find media upload: {:?}", media_upload_token))) + } + Err(e) => { + error!("could not query media upload: {:?}", e); + return Err(ProcessSingleJobError::from_anyhow_error(e)) + } + }; + + let media_upload_bucket_path = MediaUploadOriginalFilePath::from_object_hash(&media_upload_result.public_bucket_directory_hash); + bucket_object_path = media_upload_bucket_path.to_full_object_pathbuf(); + } + } + + let downloaded_filesystem_path = work_temp_dir.path().join("image.bin"); + + maybe_download_file_from_bucket( + "image file", + &downloaded_filesystem_path, + &bucket_object_path, + public_bucket_client, + job_progress_reporter, + "downloading", + job.id.0, + &temp_dir_creator, + ).await?; + + Ok(ImageFile { + filesystem_path: downloaded_filesystem_path, + }) +} 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 a21070872a..77cad1c1f3 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 @@ -1,3 +1,5 @@ +pub mod download_audio_file; +pub mod download_image_file; pub mod model_downloaders; pub mod process_job; pub mod validate_job; \ No newline at end of file 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 index 90181b03d9..ac638eb371 100644 --- 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 @@ -1,5 +1,4 @@ use anyhow::anyhow; -use crate::util::model_downloader::ModelDownloader; 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; @@ -7,25 +6,29 @@ use container_common::filesystem::safe_delete_temp_directory::safe_delete_temp_d 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::lipsync::sad_talker::download_audio_file::download_audio_file; +use crate::job::job_types::lipsync::sad_talker::download_image_file::download_image_file; +use crate::job::job_types::lipsync::sad_talker::validate_job::validate_job; 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 crate::util::model_downloader::ModelDownloader; use enums::by_table::generic_inference_jobs::inference_result_type::InferenceResultType; +use errors::AnyhowResult; 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 errors::AnyhowResult; use mysql_queries::payloads::generic_inference_args::generic_inference_args::{GenericInferenceArgs, InferenceCategoryAbbreviated, PolymorphicInferenceArgs}; use mysql_queries::payloads::generic_inference_args::lipsync_payload::{LipsyncAnimationAudioSource, LipsyncAnimationImageSource}; +use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; use mysql_queries::queries::voice_conversion::results::get_voice_conversion_result_for_inference::get_voice_conversion_result_for_inference; +use mysql_queries::queries::voice_conversion::results::insert_voice_conversion_result::{insert_voice_conversion_result, InsertArgs}; +use std::path::PathBuf; +use std::thread; +use std::time::{Duration, Instant}; use tokens::users::user::UserToken; -use crate::job::job_types::lipsync::sad_talker::validate_job::validate_job; pub struct SadTalkerProcessJobArgs<'a> { pub job_dependencies: &'a JobDependencies, @@ -73,88 +76,31 @@ pub async fn process_job(args: SadTalkerProcessJobArgs<'_>) -> Result { - // TODO(bt, 2023-09-08): Implement - return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("not yet implemented"))) - } - LipsyncAnimationAudioSource::U(media_upload_token) => { - // TODO(bt, 2023-09-08): Implement - return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("not yet implemented"))) - } - LipsyncAnimationAudioSource::T(tts_result_token) => { - // TODO(bt, 2023-09-08): Implement - return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("not yet implemented"))) - } - LipsyncAnimationAudioSource::V(voice_conversion_result_token) => { - let voice_conversion_result = get_voice_conversion_result_for_inference( - voice_conversion_result_token, - false, - &args.job_dependencies.mysql_pool, - ).await; + info!("Audio file: {:?}", audio_bucket_path.filesystem_path); - let voice_conversion_result = match voice_conversion_result { - Ok(Some(result)) => result, - Ok(None) => { - warn!("could not find voice conversion result: {:?}", voice_conversion_result_token); - return Err(ProcessSingleJobError::from_anyhow_error( - anyhow!("could not find voice conversion result: {:?}", voice_conversion_result_token))) - } - Err(e) => { - error!("could not query voice conversion result: {:?}", e); - return Err(ProcessSingleJobError::from_anyhow_error(e)) - } - }; + let image_bucket_path = download_image_file( + &job_args.image_source, + &args.job_dependencies.public_bucket_client, + &mut job_progress_reporter, + job, + &args.job_dependencies.fs.scoped_temp_dir_creator_for_work, + &work_temp_dir, + &deps.mysql_pool + ).await?; - audio_bucket_path = VoiceConversionResultOriginalFilePath::from_object_hash(&voice_conversion_result.public_bucket_hash); - } - } - - // ==================== DECIDE WHAT TYPE OF IMAGE TO DOWNLOAD ==================== // - - - match job_args.image_source { - LipsyncAnimationImageSource::F(media_file_token) => {} - LipsyncAnimationImageSource::U(media_upload_token) => {} - } - - - /* - - // ==================== DOWNLOAD IMAGE MEDIA FILE ==================== // - - info!("Download image media file ..."); - - // TODO: If already transcoded, download the transcoded file instead. - // TODO: Turn this into a general utility. - - let image_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 - }; + info!("Image file: {:?}", image_bucket_path.filesystem_path); // ==================== TRANSCODE MEDIA (IF NECESSARY) ==================== // @@ -162,11 +108,13 @@ pub async fn process_job(args: SadTalkerProcessJobArgs<'_>) -> Result