From 783ffef82c4574dff7e939e1a2147bf4a36faaf6 Mon Sep 17 00:00:00 2001
From: Brandon Thomas <echelon@gmail.com>
Date: Wed, 17 Jul 2024 08:04:24 -0400
Subject: [PATCH] permanently fail on face detection errors

---
 .../lipsync/sad_talker/model_downloaders.rs   |  2 +-
 .../categorize_live_portrait_error.rs         | 52 +++++++++++++++++++
 .../workflow/live_portrait/command_args.rs    |  1 -
 .../job_types/workflow/live_portrait/mod.rs   |  1 +
 .../process_live_portrait_job.rs              | 36 +++++++++++--
 5 files changed, 85 insertions(+), 7 deletions(-)
 create mode 100644 crates/service/job/inference_job/src/job/job_types/workflow/live_portrait/categorize_live_portrait_error.rs

diff --git a/crates/service/job/inference_job/src/job/job_types/lipsync/sad_talker/model_downloaders.rs b/crates/service/job/inference_job/src/job/job_types/lipsync/sad_talker/model_downloaders.rs
index d512b8bee4..3d419f717d 100644
--- a/crates/service/job/inference_job/src/job/job_types/lipsync/sad_talker/model_downloaders.rs
+++ b/crates/service/job/inference_job/src/job/job_types/lipsync/sad_talker/model_downloaders.rs
@@ -117,7 +117,7 @@ mod test {
   use std::path::Path;
 
   use crate::job::job_types::lipsync::sad_talker::model_downloaders::{SadTalkerModelFirstMapping, SadTalkerModelSecondMapping};
-  use crate::util::model_downloader::ModelDownloader;
+  use crate::util::downloaders::model_downloader::ModelDownloader;
 
   #[test]
   fn test_downloader_model_first_mapping() {
diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/live_portrait/categorize_live_portrait_error.rs b/crates/service/job/inference_job/src/job/job_types/workflow/live_portrait/categorize_live_portrait_error.rs
new file mode 100644
index 0000000000..bc3c69924b
--- /dev/null
+++ b/crates/service/job/inference_job/src/job/job_types/workflow/live_portrait/categorize_live_portrait_error.rs
@@ -0,0 +1,52 @@
+use crate::job::job_loop::process_single_job_error::ProcessSingleJobError;
+
+pub fn categorize_live_portrait_error(stderr_contents: &str) -> Option<ProcessSingleJobError> {
+  if stderr_contents.contains("No face detected in the source image!") {
+    return Some(ProcessSingleJobError::FaceDetectionFailure);
+  }
+  None
+}
+
+#[cfg(test)]
+mod tests {
+  use anyhow::bail;
+
+  use errors::AnyhowResult;
+
+  use crate::job::job_loop::process_single_job_error::ProcessSingleJobError;
+  use crate::job::job_types::workflow::live_portrait::categorize_live_portrait_error::categorize_live_portrait_error;
+
+  fn assert_face_tracking_error(maybe_error: Option<ProcessSingleJobError>) -> AnyhowResult<()> {
+    // NB: We have to do this since ProcessSingleJobError can't derive `PartialEq`
+    match maybe_error {
+      Some(ProcessSingleJobError::FaceDetectionFailure) => Ok(()),
+      None => bail!("error is absent"),
+      Some(_) => bail!("error is wrong type"),
+    }
+  }
+
+  #[test]
+  fn stack_trace_0() {
+    // Failure video:
+    // https://staging.fakeyou.com/media/m_xwphfcm7y8eh4dabva24sd5hwqm6z2
+    let trace = r#"
+    FaceAnalysisDIY warmup time: 1.020s
+Detecting and cropping..:   0%|          | 0/170 [00:00<?, ?it/s]
+!!! Exception during processing!!! No face detected in the source image!
+Traceback (most recent call last):
+  File "/app/ComfyUI/execution.py", line 151, in recursive_execute
+    output_data, output_ui = get_output_data(obj, input_data_all)
+  File "/app/ComfyUI/execution.py", line 81, in get_output_data
+    return_values = map_node_over_list(obj, input_data_all, obj.FUNCTION, allow_interrupt=True)
+  File "/app/ComfyUI/execution.py", line 74, in map_node_over_list
+    results.append(getattr(obj, func)(**slice_dict(input_data_all, i)))
+  File "/app/ComfyUI/custom_nodes/ComfyUI-LivePortraitKJ/nodes.py", line 408, in process
+    crop_info = self.cropper.crop_single_image(source_image_np[i], dsize, scale, vy_ratio, vx_ratio, face_index, rotate)
+  File "/app/ComfyUI/custom_nodes/ComfyUI-LivePortraitKJ/liveportrait/utils/cropper.py", line 55, in crop_single_image
+    raise Exception("No face detected in the source image!")
+Exception: No face detected in the source image!
+    "#;
+
+    assert_face_tracking_error(categorize_live_portrait_error(trace)).unwrap();
+  }
+}
diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/live_portrait/command_args.rs b/crates/service/job/inference_job/src/job/job_types/workflow/live_portrait/command_args.rs
index 9aa819271e..389b2218bb 100644
--- a/crates/service/job/inference_job/src/job/job_types/workflow/live_portrait/command_args.rs
+++ b/crates/service/job/inference_job/src/job/job_types/workflow/live_portrait/command_args.rs
@@ -1,7 +1,6 @@
 use std::path::Path;
 
 use filesys::path_to_string::path_to_string;
-use primitives::bool_to_str::bool_to_str;
 use subprocess_common::command_runner::command_args::CommandArgs;
 
 #[derive(Debug)]
diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/live_portrait/mod.rs b/crates/service/job/inference_job/src/job/job_types/workflow/live_portrait/mod.rs
index 2e9efa9e40..568104d14a 100644
--- a/crates/service/job/inference_job/src/job/job_types/workflow/live_portrait/mod.rs
+++ b/crates/service/job/inference_job/src/job/job_types/workflow/live_portrait/mod.rs
@@ -1,3 +1,4 @@
+pub mod categorize_live_portrait_error;
 pub mod command_args;
 pub mod extract_live_portrait_payload_from_job;
 pub mod process_live_portrait_job;
\ No newline at end of file
diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/live_portrait/process_live_portrait_job.rs b/crates/service/job/inference_job/src/job/job_types/workflow/live_portrait/process_live_portrait_job.rs
index f62329713c..8a0116b054 100644
--- a/crates/service/job/inference_job/src/job/job_types/workflow/live_portrait/process_live_portrait_job.rs
+++ b/crates/service/job/inference_job/src/job/job_types/workflow/live_portrait/process_live_portrait_job.rs
@@ -43,7 +43,9 @@ use videos::ffprobe_get_dimensions::ffprobe_get_dimensions;
 
 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::workflow::comfy_ui_inference_command::{InferenceArgs, InferenceDetails};
+use crate::job::job_types::workflow::live_portrait::categorize_live_portrait_error::categorize_live_portrait_error;
 use crate::job::job_types::workflow::live_portrait::command_args::LivePortraitCommandArgs;
 use crate::job::job_types::workflow::live_portrait::extract_live_portrait_payload_from_job::extract_live_portrait_payload_from_job;
 use crate::job::job_types::workflow::video_style_transfer::extract_vst_workflow_payload_from_job::extract_vst_workflow_payload_from_job;
@@ -193,16 +195,19 @@ pub async fn process_live_portrait_job(
     info!("Comfy output video dimensions: {}x{}", dimensions.width, dimensions.height);
   }
 
+  // ==================== CHECK STATUS ======================== //
+
+  // if !command_exit_status.is_success() {
+  //   error!("Inference failed: {:?}", command_exit_status);
+  // }
+
   // ==================== CHECK OUTPUT FILE ======================== //
 
   if let Err(err) = check_file_exists(&output_file_path) {
     error!("Output file does not  exist: {:?}", err);
+    error!("Inference failed with exit status: {:?}", command_exit_status);
 
-    error!("Inference failed: {:?}", command_exit_status);
-
-    if let Ok(contents) = read_to_string(&stderr_output_file) {
-      warn!("Captured stderr output: {}", contents);
-    }
+    print_and_detect_stderr_issues(&stderr_output_file)?;
 
     safe_delete_temp_file(&stderr_output_file);
     safe_delete_temp_file(&stdout_output_file);
@@ -335,3 +340,24 @@ pub async fn process_live_portrait_job(
     inference_duration,
   })
 }
+
+fn print_and_detect_stderr_issues(stderr_output_file: &Path) -> Result<(), ProcessSingleJobError> {
+  let contents = match read_to_string(stderr_output_file) {
+    Ok(contents) => {
+      warn!("Captured stderr output: {}", contents);
+      contents
+    },
+    Err(err) => {
+      error!("Error reading stderr output: {:?}", err);
+      return Ok(());
+    }
+  };
+
+  match categorize_live_portrait_error(&contents) {
+    Some(ProcessSingleJobError::FaceDetectionFailure) => {
+      warn!("Face not detected in source");
+      Err(ProcessSingleJobError::FaceDetectionFailure)
+    }
+    _ => Ok(())
+  }
+}