From de8eb05db68a39636e0ee40cd8ffdf2848d64490 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Fri, 8 Dec 2023 04:45:28 -0500 Subject: [PATCH] inference-job - sadtalker to work with image media files --- .../lipsync/sad_talker/download_audio_file.rs | 3 +- .../lipsync/sad_talker/download_image_file.rs | 113 ++++++++++++------ .../lipsync/sad_talker/process_job.rs | 20 ++-- 3 files changed, 86 insertions(+), 50 deletions(-) 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 index 5e6c8c7d47..4e6f8973e7 100644 --- 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 @@ -4,8 +4,8 @@ use anyhow::anyhow; use log::{error, warn}; use sqlx::MySqlPool; use tempdir::TempDir; -use buckets::public::media_files::original_file::MediaFileBucketPath; +use buckets::public::media_files::original_file::MediaFileBucketPath; use buckets::public::media_uploads::original_file::MediaUploadOriginalFilePath; use buckets::public::voice_conversion_results::original_file::VoiceConversionResultOriginalFilePath; use cloud_storage::bucket_client::BucketClient; @@ -75,7 +75,6 @@ pub async fn download_audio_file(args: DownloadAudioFileArgs<'_>) -> Result) -> Result { let media_file_token = MediaFileToken::new_from_str(media_file_token); let media_file_result = get_media_file( 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 index 6ac3b4d3e9..3d401f0a7a 100644 --- 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 @@ -5,12 +5,15 @@ use log::{error, warn}; use sqlx::MySqlPool; use tempdir::TempDir; +use buckets::public::media_files::original_file::MediaFileBucketPath; use buckets::public::media_uploads::original_file::MediaUploadOriginalFilePath; use cloud_storage::bucket_client::BucketClient; use jobs_common::job_progress_reporter::job_progress_reporter::JobProgressReporter; use mysql_queries::payloads::generic_inference_args::lipsync_payload::LipsyncAnimationImageSource; use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; +use mysql_queries::queries::media_files::get_media_file::get_media_file; use mysql_queries::queries::media_uploads::get_media_upload_for_inference::get_media_upload_for_inference; +use tokens::tokens::media_files::MediaFileToken; use tokens::tokens::media_uploads::MediaUploadToken; use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; @@ -21,59 +24,40 @@ 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 { +pub struct DownloadImageFileArgs<'a> { + pub image_source: &'a LipsyncAnimationImageSource, + pub public_bucket_client: &'a BucketClient, + pub job_progress_reporter: &'a mut Box, + pub job: &'a AvailableInferenceJob, + pub temp_dir_creator: &'a ScopedTempDirCreator, + pub work_temp_dir: &'a TempDir, + pub mysql_pool: &'a MySqlPool, +} + +pub async fn download_image_file(args: DownloadImageFileArgs<'_>) -> Result { let bucket_object_path; - match image_source { + match args.image_source { LipsyncAnimationImageSource::F(media_file_token) => { - // TODO(bt, 2023-09-08): Implement - return Err(ProcessSingleJobError::NotYetImplemented) + bucket_object_path = from_media_file(media_file_token, &args).await?; } 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(); + bucket_object_path = from_media_upload(media_upload_token, &args).await?; } } - let downloaded_filesystem_path = work_temp_dir.path().join("image.bin"); + let downloaded_filesystem_path = args.work_temp_dir.path().join("image.bin"); maybe_download_file_from_bucket(MaybeDownloadArgs { name_or_description_of_file: "image file", final_filesystem_file_path: &downloaded_filesystem_path, bucket_object_path: &bucket_object_path, - bucket_client: public_bucket_client, - job_progress_reporter, + bucket_client: args.public_bucket_client, + job_progress_reporter: args.job_progress_reporter, job_progress_update_description: "downloading", - job_id: job.id.0, - scoped_tempdir_creator: &temp_dir_creator, + job_id: args.job.id.0, + scoped_tempdir_creator: &args.temp_dir_creator, maybe_existing_file_minimum_size_required: None, }).await?; @@ -81,3 +65,56 @@ pub async fn download_image_file( filesystem_path: downloaded_filesystem_path, }) } + +async fn from_media_file(media_file_token: &str, args: &DownloadImageFileArgs<'_>) -> Result { + let media_file_token = MediaFileToken::new_from_str(media_file_token); + let media_file_result = get_media_file( + &media_file_token, + true, + args.mysql_pool + ).await; + + let media_file_result = match media_file_result { + Ok(Some(result)) => result, + Ok(None) => { + warn!("could not find media file : {:?}", media_file_token); + return Err(ProcessSingleJobError::from_anyhow_error( + anyhow!("could not find media file : {:?}", media_file_token))) + } + Err(e) => { + error!("could not query media file : {:?}", e); + return Err(ProcessSingleJobError::from_anyhow_error(e)) + } + }; + + let media_file_bucket_path = MediaFileBucketPath::from_object_hash( + &media_file_result.public_bucket_directory_hash, + media_file_result.maybe_public_bucket_prefix.as_deref(), + media_file_result.maybe_public_bucket_extension.as_deref()); + + Ok(media_file_bucket_path.to_full_object_pathbuf()) +} + +async fn from_media_upload(media_upload_token: &str, args: &DownloadImageFileArgs<'_>) -> Result { + let media_upload_token = MediaUploadToken::new_from_str(media_upload_token); + let media_upload_result = get_media_upload_for_inference( + &media_upload_token, + args.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); + Ok(media_upload_bucket_path.to_full_object_pathbuf()) +} 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 0d75a844bb..cb6d563156 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 @@ -20,7 +20,7 @@ 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::categorize_error::categorize_error; use crate::job::job_types::lipsync::sad_talker::download_audio_file::{download_audio_file, DownloadAudioFileArgs}; -use crate::job::job_types::lipsync::sad_talker::download_image_file::download_image_file; +use crate::job::job_types::lipsync::sad_talker::download_image_file::{download_image_file, DownloadImageFileArgs}; use crate::job::job_types::lipsync::sad_talker::resize_image::resize_image; use crate::job::job_types::lipsync::sad_talker::sad_talker_inference_command::InferenceArgs; use crate::job::job_types::lipsync::sad_talker::validate_job::validate_job; @@ -107,17 +107,17 @@ pub async fn process_job(args: SadTalkerProcessJobArgs<'_>) -> Result