From 63143121f61cd5aa3abc1f21f064084fba147aed Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 12 Sep 2024 14:18:51 -0400 Subject: [PATCH] add lipsync details to jobs endpoints --- .../inference_job/common_responses/lipsync.rs | 14 ++++ .../inference_job/common_responses/mod.rs | 1 + .../batch_get_inference_job_status_handler.rs | 13 +++- .../get/get_inference_job_status_handler.rs | 13 +++- .../list/list_session_jobs_handler.rs | 13 +++- .../extractors/extract_lipsync_details.rs | 78 +++++++++++++++++++ .../inference_job/utils/extractors/mod.rs | 1 + 7 files changed, 127 insertions(+), 6 deletions(-) create mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/common_responses/lipsync.rs create mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/extractors/extract_lipsync_details.rs diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/common_responses/lipsync.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/common_responses/lipsync.rs new file mode 100644 index 0000000000..8dbabd9da3 --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/common_responses/lipsync.rs @@ -0,0 +1,14 @@ +use tokens::tokens::media_files::MediaFileToken; +use utoipa::ToSchema; + +/// Details about submitted lipsync jobs (request arguments only) +#[derive(Serialize, ToSchema)] +pub struct JobDetailsLipsyncRequest { + /// Media file token for the source audio. + /// This is probably an audio file, but in the future we might pull audio from video. + pub audio_source_token: MediaFileToken, + + /// Media file token for the source visuals. + /// This is either an image or video. + pub image_or_video_source_token: MediaFileToken, +} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/common_responses/mod.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/common_responses/mod.rs index d114627d88..8514adecdf 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/common_responses/mod.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/common_responses/mod.rs @@ -1 +1,2 @@ +pub mod lipsync; pub mod live_portrait; \ No newline at end of file diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/get/batch_get_inference_job_status_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/get/batch_get_inference_job_status_handler.rs index d063da76fa..fff70dc200 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/get/batch_get_inference_job_status_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/get/batch_get_inference_job_status_handler.rs @@ -23,9 +23,10 @@ use mysql_queries::queries::generic_inference::web::job_status::GenericInference use redis_common::redis_keys::RedisKeys; use tokens::tokens::generic_inference_jobs::InferenceJobToken; use tokens::tokens::media_files::MediaFileToken; - +use crate::http_server::endpoints::inference_job::common_responses::lipsync::JobDetailsLipsyncRequest; use crate::http_server::endpoints::inference_job::common_responses::live_portrait::JobDetailsLivePortraitRequest; use crate::http_server::endpoints::inference_job::utils::estimates::estimate_job_progress::estimate_job_progress; +use crate::http_server::endpoints::inference_job::utils::extractors::extract_lipsync_details::extract_lipsync_details; use crate::http_server::endpoints::inference_job::utils::extractors::extract_live_portrait_details::extract_live_portrait_details; use crate::http_server::endpoints::inference_job::utils::extractors::extract_polymorphic_inference_args::extract_polymorphic_inference_args; use crate::http_server::endpoints::media_files::get::batch_get_media_files_handler::BatchGetMediaFilesQueryParams; @@ -90,6 +91,10 @@ pub struct BatchRequestDetailsResponse { /// OPTIONAL. For Live Portrait jobs, this is additional information on the request. pub maybe_live_portrait_details: Option, + + /// OPTIONAL. For lipsync jobs (face fusion and sad talker), this is additional + /// information on the request. + pub maybe_lipsync_details: Option, } /// Details about the ongoing job status @@ -277,7 +282,11 @@ fn db_record_to_response_payload( maybe_raw_inference_text: record.request_details.maybe_raw_inference_text, maybe_style_name: record.request_details.maybe_style_name, maybe_live_portrait_details: maybe_polymorphic_args - .and_then(|ref args| extract_live_portrait_details(args)), + .as_ref() + .and_then(|args| extract_live_portrait_details(args)), + maybe_lipsync_details: maybe_polymorphic_args + .as_ref() + .and_then(|args| extract_lipsync_details(args)), }, status: BatchStatusDetailsResponse { status: record.status, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/get/get_inference_job_status_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/get/get_inference_job_status_handler.rs index 20761d1ed5..3d3d74ce06 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/get/get_inference_job_status_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/get/get_inference_job_status_handler.rs @@ -20,9 +20,10 @@ use mysql_queries::queries::generic_inference::web::get_inference_job_status::ge use mysql_queries::queries::generic_inference::web::job_status::GenericInferenceJobStatus; use redis_common::redis_keys::RedisKeys; use tokens::tokens::generic_inference_jobs::InferenceJobToken; - +use crate::http_server::endpoints::inference_job::common_responses::lipsync::JobDetailsLipsyncRequest; use crate::http_server::endpoints::inference_job::common_responses::live_portrait::JobDetailsLivePortraitRequest; use crate::http_server::endpoints::inference_job::utils::estimates::estimate_job_progress::estimate_job_progress; +use crate::http_server::endpoints::inference_job::utils::extractors::extract_lipsync_details::extract_lipsync_details; use crate::http_server::endpoints::inference_job::utils::extractors::extract_live_portrait_details::extract_live_portrait_details; use crate::http_server::endpoints::inference_job::utils::extractors::extract_polymorphic_inference_args::extract_polymorphic_inference_args; use crate::http_server::web_utils::filter_model_name::maybe_filter_model_name; @@ -78,6 +79,10 @@ pub struct RequestDetailsResponse { /// OPTIONAL. For Live Portrait jobs, this is additional information on the request. pub maybe_live_portrait_details: Option, + + /// OPTIONAL. For lipsync jobs (face fusion and sad talker), this is additional + /// information on the request. + pub maybe_lipsync_details: Option, } /// Details about the ongoing job status @@ -267,7 +272,11 @@ fn record_to_payload( maybe_raw_inference_text: record.request_details.maybe_raw_inference_text, maybe_style_name: record.request_details.maybe_style_name, maybe_live_portrait_details: maybe_polymorphic_args - .and_then(|ref args| extract_live_portrait_details(args)), + .as_ref() + .and_then(|args| extract_live_portrait_details(args)), + maybe_lipsync_details: maybe_polymorphic_args + .as_ref() + .and_then(|args| extract_lipsync_details(args)), }, status: StatusDetailsResponse { status: record.status, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/list/list_session_jobs_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/list/list_session_jobs_handler.rs index 98a7a1eed7..d6ccc6eaee 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/list/list_session_jobs_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/list/list_session_jobs_handler.rs @@ -24,9 +24,10 @@ use redis_common::redis_keys::RedisKeys; use tokens::tokens::generic_inference_jobs::InferenceJobToken; use tokens::tokens::media_files::MediaFileToken; use tokens::tokens::users::UserToken; - +use crate::http_server::endpoints::inference_job::common_responses::lipsync::JobDetailsLipsyncRequest; use crate::http_server::endpoints::inference_job::common_responses::live_portrait::JobDetailsLivePortraitRequest; use crate::http_server::endpoints::inference_job::utils::estimates::estimate_job_progress::estimate_job_progress; +use crate::http_server::endpoints::inference_job::utils::extractors::extract_lipsync_details::extract_lipsync_details; use crate::http_server::endpoints::inference_job::utils::extractors::extract_live_portrait_details::extract_live_portrait_details; use crate::http_server::endpoints::inference_job::utils::extractors::extract_polymorphic_inference_args::extract_polymorphic_inference_args; use crate::http_server::web_utils::filter_model_name::maybe_filter_model_name; @@ -94,6 +95,10 @@ pub struct ListSessionRequestDetailsResponse { /// OPTIONAL. For Live Portrait jobs, this is additional information on the request. pub maybe_live_portrait_details: Option, + + /// OPTIONAL. For lipsync jobs (face fusion and sad talker), this is additional + /// information on the request. + pub maybe_lipsync_details: Option, } /// Details about the ongoing job status @@ -341,7 +346,11 @@ fn db_record_to_response_payload( maybe_raw_inference_text: record.request_details.maybe_raw_inference_text, maybe_style_name: record.request_details.maybe_style_name, maybe_live_portrait_details: maybe_polymorphic_args - .and_then(|ref args| extract_live_portrait_details(args)), + .as_ref() + .and_then(|args| extract_live_portrait_details(args)), + maybe_lipsync_details: maybe_polymorphic_args + .as_ref() + .and_then(|args| extract_lipsync_details(args)), }, status: ListSessionStatusDetailsResponse { status: record.status, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/extractors/extract_lipsync_details.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/extractors/extract_lipsync_details.rs new file mode 100644 index 0000000000..bd481add66 --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/extractors/extract_lipsync_details.rs @@ -0,0 +1,78 @@ +use crate::http_server::endpoints::inference_job::common_responses::lipsync::JobDetailsLipsyncRequest; +use mysql_queries::payloads::generic_inference_args::generic_inference_args::PolymorphicInferenceArgs; +use mysql_queries::payloads::generic_inference_args::inner_payloads::face_fusion_payload::FaceFusionPayload; +use mysql_queries::payloads::generic_inference_args::inner_payloads::lipsync_payload::{LipsyncAnimationAudioSource, LipsyncAnimationImageSource, LipsyncArgs}; +use tokens::tokens::media_files::MediaFileToken; + +/// Extract args for SadTalker and FaceFusion jobs. +pub fn extract_lipsync_details(args: &PolymorphicInferenceArgs) -> Option { + match args { + // Face Fusion + PolymorphicInferenceArgs::Ff(args) => extract_face_fusion(args), + // Sad Talker + PolymorphicInferenceArgs::La(args) => extract_sad_talker(args), + // Non-matching job types + _ => None, + } +} + +fn extract_face_fusion(args: &FaceFusionPayload) -> Option { + args.audio_media_file_token + .as_ref() + .zip(args.image_or_video_media_file_token.as_ref()) + .map(|(audio, image_or_video)| { + JobDetailsLipsyncRequest { + audio_source_token: audio.clone(), + image_or_video_source_token: image_or_video.clone(), + } + }) +} + +fn extract_sad_talker(args: &LipsyncArgs) -> Option { + args.maybe_audio_source + .as_ref() + .map(|polymorphic| match polymorphic { + // NB: We have enums over *old* record types: TTS results, Media Uploads, etc. + LipsyncAnimationAudioSource::F(token) => Some(MediaFileToken::new_from_str(token)), + _ => None, + }) + .flatten() + .zip(args.maybe_image_source.as_ref() + .map(|polymorphic| match polymorphic { + // NB: We have enums over *old* record types: TTS results, Media Uploads, etc. + LipsyncAnimationImageSource::F(token) => Some(MediaFileToken::new_from_str(token)), + _ => None, + }).flatten()) + .map(|(audio, image_or_video)| { + JobDetailsLipsyncRequest { + audio_source_token: audio, + image_or_video_source_token: image_or_video, + } + }) +} + +#[cfg(test)] +mod tests { + use mysql_queries::payloads::generic_inference_args::generic_inference_args::PolymorphicInferenceArgs; + use mysql_queries::payloads::generic_inference_args::inner_payloads::face_fusion_payload::FaceFusionPayload; + use tokens::tokens::media_files::MediaFileToken; + + use super::extract_lipsync_details; + + #[test] + fn test_valid_args_with_tokens() { + let polymorphic_args = PolymorphicInferenceArgs::Ff(FaceFusionPayload{ + audio_media_file_token: Some(MediaFileToken::new_from_str("audio")), + image_or_video_media_file_token: Some(MediaFileToken::new_from_str("image")), + crop: None, + watermark_type: None, + sleep_millis: None, + }); + + let details = extract_lipsync_details(&polymorphic_args) + .expect("should contain tokens"); + + assert_eq!(details.audio_source_token.as_str(), "audio"); + assert_eq!(details.image_or_video_source_token.as_str(), "image"); + } +} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/extractors/mod.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/extractors/mod.rs index 8300fc5b96..6045d824be 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/extractors/mod.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/extractors/mod.rs @@ -1,3 +1,4 @@ pub mod extract_comfy_workflow_args; +pub mod extract_lipsync_details; pub mod extract_live_portrait_details; pub mod extract_polymorphic_inference_args; \ No newline at end of file