From 1f7dc268a8663301d2fa0d469da781aae75d6f3e Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Tue, 19 Sep 2023 17:27:34 -0400 Subject: [PATCH] SadTalker arguments (wip) --- .../generic_inference_args.rs | 1 + .../generic_inference_args/lipsync_payload.rs | 50 ++++++++++++++++++- .../animation/enqueue_lipsync_animation.rs | 10 ++++ 3 files changed, 60 insertions(+), 1 deletion(-) diff --git a/crates/lib/mysql_queries/src/payloads/generic_inference_args/generic_inference_args.rs b/crates/lib/mysql_queries/src/payloads/generic_inference_args/generic_inference_args.rs index 99021de4e3..b82e920126 100644 --- a/crates/lib/mysql_queries/src/payloads/generic_inference_args/generic_inference_args.rs +++ b/crates/lib/mysql_queries/src/payloads/generic_inference_args/generic_inference_args.rs @@ -122,6 +122,7 @@ mod tests { args: Some(PolymorphicInferenceArgs::La(LipsyncArgs { maybe_audio_source: Some(LipsyncAnimationAudioSource::U("foo".to_string())), maybe_image_source: Some(LipsyncAnimationImageSource::F("bar".to_string())), + maybe_face_enhancer: None, })), }; diff --git a/crates/lib/mysql_queries/src/payloads/generic_inference_args/lipsync_payload.rs b/crates/lib/mysql_queries/src/payloads/generic_inference_args/lipsync_payload.rs index 8871ee399a..a1eed7e010 100644 --- a/crates/lib/mysql_queries/src/payloads/generic_inference_args/lipsync_payload.rs +++ b/crates/lib/mysql_queries/src/payloads/generic_inference_args/lipsync_payload.rs @@ -9,6 +9,27 @@ pub struct LipsyncArgs { #[serde(rename = "i")] // NB: DO NOT CHANGE. It could break live jobs. Renamed to be fewer bytes. #[serde(skip_serializing_if = "Option::is_none")] pub maybe_image_source: Option, + + /// Which face enhancement algorithm to run. + /// This makes the video larger in dimensions and higher quality. + #[serde(rename = "f")] // NB: DO NOT CHANGE. It could break live jobs. Renamed to be fewer bytes. + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_face_enhancer: Option, + + /// Number between [0, 46) to control pose. + #[serde(rename = "f")] // NB: DO NOT CHANGE. It could break live jobs. Renamed to be fewer bytes. + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_pose_style: Option, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub enum FaceEnhancer { + // NB: DO NOT CHANGE. It could break live jobs. Renamed to be fewer bytes. + /// "gfpgan" + G, + // NB: DO NOT CHANGE. It could break live jobs. Renamed to be fewer bytes. + /// "RestoreFormer" + R, } /// Audio sources can be one of several: @@ -82,13 +103,14 @@ impl LipsyncAnimationImageSource { #[cfg(test)] mod tests { - use crate::payloads::generic_inference_args::lipsync_payload::{LipsyncAnimationAudioSource, LipsyncAnimationImageSource, LipsyncArgs}; + use crate::payloads::generic_inference_args::lipsync_payload::{FaceEnhancer, LipsyncAnimationAudioSource, LipsyncAnimationImageSource, LipsyncArgs}; #[test] fn test_media_file() { let args = LipsyncArgs { maybe_audio_source: Some(LipsyncAnimationAudioSource::media_file_token("audio_media_file")), maybe_image_source: Some(LipsyncAnimationImageSource::media_file_token("image_media_file")), + maybe_face_enhancer: None, }; let json = serde_json::ser::to_string(&args).unwrap(); assert_eq!(json, r#"{"a":{"F":"audio_media_file"},"i":{"F":"image_media_file"}}"#.to_string()); @@ -99,6 +121,7 @@ mod tests { let args = LipsyncArgs { maybe_audio_source: Some(LipsyncAnimationAudioSource::media_upload_token("audio_media_upload")), maybe_image_source: Some(LipsyncAnimationImageSource::media_upload_token("image_media_upload")), + maybe_face_enhancer: None, }; let json = serde_json::ser::to_string(&args).unwrap(); assert_eq!(json, r#"{"a":{"U":"audio_media_upload"},"i":{"U":"image_media_upload"}}"#.to_string()); @@ -109,17 +132,42 @@ mod tests { let args = LipsyncArgs { maybe_audio_source: Some(LipsyncAnimationAudioSource::tts_result_token("audio_tts_result")), maybe_image_source: Some(LipsyncAnimationImageSource::media_upload_token("image_media_upload")), + maybe_face_enhancer: None, }; let json = serde_json::ser::to_string(&args).unwrap(); assert_eq!(json, r#"{"a":{"T":"audio_tts_result"},"i":{"U":"image_media_upload"}}"#.to_string()); } + #[test] fn test_voice_conversion_result() { let args = LipsyncArgs { maybe_audio_source: Some(LipsyncAnimationAudioSource::voice_conversion_result_token("audio_voice_conversion_result")), maybe_image_source: Some(LipsyncAnimationImageSource::media_upload_token("image_media_upload")), + maybe_face_enhancer: None, }; let json = serde_json::ser::to_string(&args).unwrap(); assert_eq!(json, r#"{"a":{"V":"audio_voice_conversion_result"},"i":{"U":"image_media_upload"}}"#.to_string()); } + + #[test] + fn test_face_enhancer_1() { + let args = LipsyncArgs { + maybe_audio_source: None, + maybe_image_source: None, + maybe_face_enhancer: Some(FaceEnhancer::G), + }; + let json = serde_json::ser::to_string(&args).unwrap(); + assert_eq!(json, r#"{"f":"G"}"#.to_string()); + } + + #[test] + fn test_face_enhancer_2() { + let args = LipsyncArgs { + maybe_audio_source: None, + maybe_image_source: None, + maybe_face_enhancer: Some(FaceEnhancer::R), + }; + let json = serde_json::ser::to_string(&args).unwrap(); + assert_eq!(json, r#"{"f":"R"}"#.to_string()); + } } diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/animation/enqueue_lipsync_animation.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/animation/enqueue_lipsync_animation.rs index 2710c02eb3..e415fb9893 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/animation/enqueue_lipsync_animation.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/animation/enqueue_lipsync_animation.rs @@ -47,6 +47,14 @@ pub struct EnqueueLipsyncAnimationRequest { /// SadTalker: parameter to animate the full body. animate_full_body: Option, + + /// SadTalker: parameter to enhance the face using gfpgan. + /// If `enhance_face_with` is set, it will take precedence over this arg. + enhance_face: Option, + + /// SadTalker: parameter to enhance the face using gfpgan or RestoreFormer. + /// If `enhance_face` is set, this `enhance_face_with` argument will take precedence over it. + enhance_face_with: Option, } /// Treated as an enum. Only one of these may be set. @@ -269,6 +277,8 @@ pub async fn enqueue_lipsync_animation_handler( let maybe_args = Some(PolymorphicInferenceArgs::La(LipsyncArgs { maybe_audio_source: Some(audio_source), maybe_image_source: Some(image_source), + maybe_face_enhancer: None, + maybe_pose_style: None, })); let query_result = insert_generic_inference_job(InsertGenericInferenceArgs {