inference-job - sadtalker to work with image media files

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<Audi
   })
 }
 
-
 async fn from_media_file(media_file_token: &str, args: &DownloadAudioFileArgs<'_>) -> Result<PathBuf, ProcessSingleJobError> {
   let media_file_token = MediaFileToken::new_from_str(media_file_token);
   let media_file_result = get_media_file(

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<dyn JobProgressReporter>,
-  job: &AvailableInferenceJob,
-  temp_dir_creator: &ScopedTempDirCreator,
-  work_temp_dir: &TempDir,
-  mysql_pool: &MySqlPool,
-) -> Result<ImageFile, ProcessSingleJobError> {
+pub struct DownloadImageFileArgs<'a> {
+  pub image_source: &'a LipsyncAnimationImageSource,
+  pub public_bucket_client: &'a BucketClient,
+  pub job_progress_reporter: &'a mut Box<dyn JobProgressReporter>,
+  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<ImageFile, ProcessSingleJobError> {
 
   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<PathBuf, ProcessSingleJobError> {
+  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<PathBuf, ProcessSingleJobError> {
+  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())
+}

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<JobSuccess
     mysql_pool: &deps.db.mysql_pool
   }).await?;
 
-  info!("Audio file: {:?}", audio_path.filesystem_path);
+  info!("Downloaded audio file: {:?}", audio_path.filesystem_path);
 
-  let image_path = download_image_file(
-    &job_args.image_source,
-    &args.job_dependencies.buckets.public_bucket_client,
-    &mut job_progress_reporter,
+  let image_path = download_image_file(DownloadImageFileArgs {
+    image_source: &job_args.image_source,
+    public_bucket_client: &args.job_dependencies.buckets.public_bucket_client,
+    job_progress_reporter: &mut job_progress_reporter,
     job,
-    &args.job_dependencies.fs.scoped_temp_dir_creator_for_work,
-    &work_temp_dir,
-    &deps.db.mysql_pool
-  ).await?;
+    temp_dir_creator: &args.job_dependencies.fs.scoped_temp_dir_creator_for_work,
+    work_temp_dir: &work_temp_dir,
+    mysql_pool: &deps.db.mysql_pool
+  }).await?;
 
   info!("Downloaded image file: {:?}", image_path.filesystem_path);