SadTalker arguments (wip)

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,
       })),
     };
 

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<LipsyncAnimationImageSource>,
+
+  /// 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<FaceEnhancer>,
+
+  /// 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<u8>,
+}
+
+#[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());
+  }
 }

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<bool>,
+
+  /// SadTalker: parameter to enhance the face using gfpgan.
+  /// If `enhance_face_with` is set, it will take precedence over this arg.
+  enhance_face: Option<bool>,
+
+  /// 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<String>,
 }
 
 /// 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 {