From baefb9f7d670643c300af57db4470cd4fb59eb12 Mon Sep 17 00:00:00 2001
From: Brandon Thomas <echelon@gmail.com>
Date: Thu, 28 Sep 2023 10:11:10 -0400
Subject: [PATCH] Do not respawn SadTalker jobs when face detection fails

Currently SadTalker jobs will retry on failure, which is great for
worker nodes with transient issues. Unfortunately failure to detect a
face is a common user input error occuring in just around half of all
image uploads. This will permanently tombstone these jobs.

The new failure class can be reused for future animation systems that
rely on face detection.

There is still some minor work to do to surface this error to the
frontend / user in a friendly way.
---
 .../src/job/job_loop/main_loop.rs             |  3 +++
 .../job/job_loop/process_single_job_error.rs  |  5 +++++
 .../lipsync/sad_talker/process_job.rs         | 20 +++++++++++++++++--
 .../sad_talker_inference_command.rs           | 12 +++++++++--
 4 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/crates/service/job/inference_job/src/job/job_loop/main_loop.rs b/crates/service/job/inference_job/src/job/job_loop/main_loop.rs
index d0271657c4..f99d376c74 100644
--- a/crates/service/job/inference_job/src/job/job_loop/main_loop.rs
+++ b/crates/service/job/inference_job/src/job/job_loop/main_loop.rs
@@ -135,6 +135,8 @@ async fn process_job_batch(job_dependencies: &JobDependencies, jobs: Vec<Availab
                 (true, false, format!("InvalidJob: {:?}", err), Some("invalid job")),
               ProcessSingleJobError::NotYetImplemented =>
                 (true, false, "not yet implemented".to_string(), Some("not yet implemented")),
+              ProcessSingleJobError::FaceDetectionFailure =>
+                (true, false, "face not detected".to_string(), Some("face not detected")),
 
               // Non-permanent failures
               ProcessSingleJobError::FilesystemFull =>
@@ -177,6 +179,7 @@ async fn process_job_batch(job_dependencies: &JobDependencies, jobs: Vec<Availab
           ProcessSingleJobError::InvalidJob(_) => {}
           ProcessSingleJobError::KeepAliveElapsed => {}
           ProcessSingleJobError::NotYetImplemented => {}
+          ProcessSingleJobError::FaceDetectionFailure => {}
         }
       }
     }
diff --git a/crates/service/job/inference_job/src/job/job_loop/process_single_job_error.rs b/crates/service/job/inference_job/src/job/job_loop/process_single_job_error.rs
index 5cc24ad7a7..c6509c1dde 100644
--- a/crates/service/job/inference_job/src/job/job_loop/process_single_job_error.rs
+++ b/crates/service/job/inference_job/src/job/job_loop/process_single_job_error.rs
@@ -3,6 +3,11 @@ use anyhow::anyhow;
 /// Error from processing a single job
 #[derive(Debug)]
 pub enum ProcessSingleJobError {
+  /// For inference jobs that are required to find faces, fail the job
+  /// immediately if no face is found. We can also surface the error to
+  /// the user in a friendly way.
+  /// Example job types: SadTalker, Wav2Lip.
+  FaceDetectionFailure,
   /// The filesystem is out of space and we need to free it up.
   FilesystemFull,
   /// The job is invalid (bad state, etc.)
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 4c112ef3a9..c88a8b2279 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
@@ -1,7 +1,8 @@
+use std::fs::read_to_string;
 use std::time::Instant;
 
 use anyhow::anyhow;
-use log::{error, info};
+use log::{error, info, warn};
 
 use buckets::public::media_files::original_file::MediaFileBucketPath;
 use container_common::filesystem::check_file_exists::check_file_exists;
@@ -132,6 +133,8 @@ pub async fn process_job(args: SadTalkerProcessJobArgs<'_>) -> Result<JobSuccess
 
   let workdir = work_temp_dir.path().to_path_buf();
 
+  let stderr_output_file = work_temp_dir.path().join("stderr.txt");
+
   let inference_start_time = Instant::now();
 
   let command_exit_status = args.job_dependencies
@@ -143,6 +146,7 @@ pub async fn process_job(args: SadTalkerProcessJobArgs<'_>) -> Result<JobSuccess
         input_image: &image_path.filesystem_path,
         work_dir: &workdir,
         output_file: &output_video_fs_path,
+        stderr_output_file: &stderr_output_file,
       });
 
   let inference_duration = Instant::now().duration_since(inference_start_time);
@@ -151,11 +155,23 @@ pub async fn process_job(args: SadTalkerProcessJobArgs<'_>) -> Result<JobSuccess
 
   if !command_exit_status.is_success() {
     error!("Inference failed: {:?}", command_exit_status);
+
+    let mut error = ProcessSingleJobError::Other(anyhow!("CommandExitStatus: {:?}", command_exit_status));
+
+    if let Ok(contents) = read_to_string(&stderr_output_file) {
+      if contents.contains("can not detect the landmark from source image") {
+        warn!("Face not detected in source image");
+        error = ProcessSingleJobError::FaceDetectionFailure;
+      }
+    }
+
     safe_delete_temp_file(&audio_path.filesystem_path);
     safe_delete_temp_file(&image_path.filesystem_path);
     safe_delete_temp_file(&output_video_fs_path);
+    safe_delete_temp_file(&stderr_output_file);
     safe_delete_temp_directory(&work_temp_dir);
-    return Err(ProcessSingleJobError::Other(anyhow!("CommandExitStatus: {:?}", command_exit_status)));
+
+    return Err(error);
   }
 
   // ==================== CHECK NON-WATERMARKED RESULT ==================== //
diff --git a/crates/service/job/inference_job/src/job/job_types/lipsync/sad_talker/sad_talker_inference_command.rs b/crates/service/job/inference_job/src/job/job_types/lipsync/sad_talker/sad_talker_inference_command.rs
index 2e4868e323..d8101c9394 100644
--- a/crates/service/job/inference_job/src/job/job_types/lipsync/sad_talker/sad_talker_inference_command.rs
+++ b/crates/service/job/inference_job/src/job/job_types/lipsync/sad_talker/sad_talker_inference_command.rs
@@ -1,13 +1,14 @@
 use std::collections::HashSet;
 use std::env;
 use std::ffi::OsString;
+use std::fs::{File, read_to_string};
 use std::path::{Path, PathBuf};
 use std::time::Duration;
 
 use anyhow::anyhow;
-use log::info;
+use log::{error, info};
 use once_cell::sync::Lazy;
-use subprocess::{Popen, PopenConfig};
+use subprocess::{Exec, Popen, PopenConfig, Redirection};
 
 use container_common::anyhow_result::AnyhowResult;
 use filesys::path_to_string::path_to_string;
@@ -79,6 +80,8 @@ pub struct InferenceArgs<P: AsRef<Path>> {
   /// --result_file: path to final file output
   pub output_file: P,
 
+  pub stderr_output_file: P,
+
   // TODO: Other SadTalker args
 }
 
@@ -243,6 +246,11 @@ impl SadTalkerInferenceCommand {
 
     let mut config = PopenConfig::default();
 
+    info!("stderr will be written to file: {:?}", args.stderr_output_file.as_ref());
+
+    let stderr_file = File::create(&args.stderr_output_file)?;
+    config.stderr = Redirection::File(stderr_file);
+
     if !env_vars.is_empty() {
       config.env = Some(env_vars);
     }