stub for studio gen 2 inference job

crates/service/job/inference_job/src/job/job_loop/process_single_job.rs

@@ -23,6 +23,7 @@ use crate::job::job_types::image_generation::process_single_ig_job::process_sing
 use crate::job::job_types::lipsync::process_single_lipsync_job::process_single_lipsync_job;
 use crate::job::job_types::mocap::process_single_mc_job::process_single_mc_job;
 use crate::job::job_types::render_engine_scene::process_render_engine_scene::process_single_render_engine_scene_job;
+use crate::job::job_types::studio_gen2::process_single_studio_gen2_job::process_single_studio_gen2_job;
 use crate::job::job_types::tts::process_single_tts_job::process_single_tts_job;
 use crate::job::job_types::vc::process_single_vc_job::process_single_vc_job;
 use crate::job::job_types::vc::rvc_v2::dispatch_rvc_v2_job::dispatch_rvc_v2_job;
@@ -273,6 +274,9 @@ async fn new_dispatch(
 ) -> Result<JobSuccessResult, ProcessSingleJobError> {
 
   let job_success_result = match job.job_type {
+    InferenceJobType::StudioGen2 => {
+      process_single_studio_gen2_job(job_dependencies, job).await?
+    },
     InferenceJobType::VideoRender
     | InferenceJobType::LivePortrait
     | InferenceJobType::FaceFusion

crates/service/job/inference_job/src/job/job_types/mod.rs

@@ -1,11 +1,12 @@
+pub mod f5_tts;
 pub mod format_conversion;
+pub mod gpt_sovits;
 pub mod image_generation;
 pub mod lipsync;
 pub mod mocap;
 pub mod render_engine_scene;
+pub mod studio_gen2;
 pub mod tts;
 pub mod vc;
 pub mod videofilter;
 pub mod workflow;
-pub mod gpt_sovits;
-pub mod f5_tts;

crates/service/job/inference_job/src/job/job_types/studio_gen2/mod.rs

@@ -0,0 +1 @@
+pub mod process_single_studio_gen2_job;

crates/service/job/inference_job/src/job/job_types/studio_gen2/process_single_studio_gen2_job.rs

@@ -0,0 +1,654 @@
+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::workflow::comfy_ui_inference_command::{InferenceArgs, InferenceDetails};
+use crate::job::job_types::workflow::face_fusion::process_face_fusion_job::process_face_fusion_job;
+use crate::job::job_types::workflow::live_portrait::process_live_portrait_job::process_live_portrait_job;
+use crate::job::job_types::workflow::video_style_transfer::extract_vst_workflow_payload_from_job::extract_vst_workflow_payload_from_job;
+use crate::job::job_types::workflow::video_style_transfer::process_video_style_transfer_job::process_video_style_transfer_job;
+use crate::job::job_types::workflow::video_style_transfer::steps::check_and_validate_job::check_and_validate_job;
+use crate::job::job_types::workflow::video_style_transfer::steps::download_global_ipa_image::{download_global_ipa_image, DownloadGlobalIpaImageArgs};
+use crate::job::job_types::workflow::video_style_transfer::steps::download_input_videos::{download_input_videos, DownloadInputVideoArgs};
+use crate::job::job_types::workflow::video_style_transfer::steps::post_process_add_watermark::{post_process_add_watermark, PostProcessAddWatermarkArgs};
+use crate::job::job_types::workflow::video_style_transfer::steps::post_process_restore_audio::{post_process_restore_audio, PostProcessRestoreVideoArgs};
+use crate::job::job_types::workflow::video_style_transfer::steps::preprocess_save_audio::{preprocess_save_audio, ProcessSaveAudioArgs};
+use crate::job::job_types::workflow::video_style_transfer::steps::preprocess_trim_and_resample_videos::{preprocess_trim_and_resample_videos, ProcessTrimAndResampleVideoArgs};
+use crate::job::job_types::workflow::video_style_transfer::steps::validate_and_save_results::{validate_and_save_results, SaveResultsArgs};
+use crate::job::job_types::workflow::video_style_transfer::util::comfy_dirs::ComfyDirs;
+use crate::job::job_types::workflow::video_style_transfer::util::process_preview_updates::PreviewProcessor;
+use crate::job::job_types::workflow::video_style_transfer::util::video_pathing::{PrimaryInputVideoAndPaths, SecondaryInputVideoAndPaths, VideoPathing};
+use crate::job::job_types::workflow::video_style_transfer::util::write_workflow_prompt::{write_workflow_prompt, WorkflowPromptArgs};
+use crate::state::job_dependencies::JobDependencies;
+use anyhow::anyhow;
+use bucket_paths::legacy::remote_file_manager_paths::remote_cloud_bucket_details::RemoteCloudBucketDetails;
+use bucket_paths::legacy::typified_paths::public::media_files::bucket_directory::MediaFileBucketDirectory;
+use cloud_storage::remote_file_manager::remote_cloud_file_manager::RemoteCloudFileClient;
+use enums::by_table::generic_inference_jobs::inference_job_type::InferenceJobType;
+use enums::by_table::generic_inference_jobs::inference_result_type::InferenceResultType;
+use filesys::check_file_exists::check_file_exists;
+use filesys::file_deletion::safe_delete_possible_files_and_directories::safe_delete_possible_files_and_directories;
+use filesys::file_deletion::safe_recursively_delete_files::safe_recursively_delete_files;
+use filesys::path_to_string::path_to_string;
+use log::{error, info, warn};
+use mysql_queries::payloads::generic_inference_args::generic_inference_args::PolymorphicInferenceArgs::Cu;
+use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob;
+use mysql_queries::queries::model_weights::get::get_weight::get_weight_by_token_with_transactor;
+use mysql_queries::utils::transactor::Transactor;
+use std::fs::read_to_string;
+use std::path::{Path, PathBuf};
+use std::thread;
+use std::time::{Duration, Instant};
+use tokio::sync::oneshot;
+use tokio::sync::oneshot::channel;
+use videos::ffprobe_get_dimensions::ffprobe_get_dimensions;
+
+// TODO(bt,2025-01-16): This is a stub for Studio Gen2
+// TODO(bt,2025-01-16): This is a stub for Studio Gen2
+// TODO(bt,2025-01-16): This is a stub for Studio Gen2
+// TODO(bt,2025-01-16): This is a stub for Studio Gen2
+
+pub async fn process_single_studio_gen2_job(
+  deps: &JobDependencies,
+  job: &AvailableInferenceJob
+) -> Result<JobSuccessResult, ProcessSingleJobError> {
+
+  // TODO(bt,2025-01-16): This is a stub for Studio Gen2
+  // TODO(bt,2025-01-16): This is a stub for Studio Gen2
+  // TODO(bt,2025-01-16): This is a stub for Studio Gen2
+  // TODO(bt,2025-01-16): This is a stub for Studio Gen2
+
+
+  let mut job_progress_reporter = deps
+      .clients
+      .job_progress_reporter
+      .new_generic_inference(job.inference_job_token.as_str())
+      .map_err(|e| ProcessSingleJobError::Other(anyhow!(e)))?;
+
+  let comfy_deps = deps
+      .job
+      .job_specific_dependencies
+      .maybe_comfy_ui_dependencies
+      .as_ref()
+      .ok_or_else(|| ProcessSingleJobError::JobSystemMisconfiguration(Some("Missing ComfyUI dependencies".to_string())))?;
+
+  // ==================== UNPACK + VALIDATE INFERENCE ARGS ==================== //
+  // check for lack of maybe_json_modifications
+
+  let job_args = check_and_validate_job(job)?;
+
+  // ==================== TEMP DIR ==================== //
+
+  let work_temp_dir = format!("temp_comfy_inference{}", job.id.0);
+
+  // NB: TempDir exists until it goes out of scope, at which point it should delete from filesystem.
+  let work_temp_dir = deps
+      .fs
+      .scoped_temp_dir_creator_for_work
+      .new_tempdir(&work_temp_dir)
+      .map_err(|e| ProcessSingleJobError::from_io_error(e))?;
+
+  // ===================== DOWNLOAD REQUIRED MODELS IF NOT EXIST ===================== //
+
+  // TODO: Replace all other paths with this
+  let comfy_dirs = ComfyDirs::new(&comfy_deps);
+
+  let root_comfy_path = comfy_deps.inference_command.mounts_directory.clone();
+
+  let remote_cloud_file_client = RemoteCloudFileClient::get_remote_cloud_file_client().await;
+  let remote_cloud_file_client = match remote_cloud_file_client {
+    Ok(res) => {
+      res
+    }
+    Err(_) => {
+      return Err(ProcessSingleJobError::from(anyhow!("failed to get remote cloud file client")));
+    }
+  };
+
+  // Download workflow to ComfyRunner
+  let workflow_dir = root_comfy_path.join("prompt");
+  // make folder if not exist
+  if !workflow_dir.exists() {
+    std::fs::create_dir_all(&workflow_dir).map_err(|err| ProcessSingleJobError::IoError(err))?;
+  }
+
+  let mut workflow_path = workflow_dir.join("prompt.json").to_str().unwrap().to_string();
+
+  info!("Grabbing redis connection from pool");
+
+  let redis_pool_dep = deps
+      .db
+      .maybe_keepalive_redis_pool.clone();
+
+  let redis_pool = redis_pool_dep
+      .ok_or_else(|| ProcessSingleJobError::Other(anyhow!("failed to get redis pool")))?;
+  // .map(|redis| redis.get())
+  // .transpose()
+  // .map_err(|err| ProcessSingleJobError::Other(anyhow!("redis pool error: {:?}", err)))?;
+
+  info!("Grabbing mysql connection from pool");
+
+  let mut mysql_connection = deps.db.mysql_pool.acquire()
+      .await
+      .map_err(|e| {
+        warn!("Could not acquire DB pool: {:?}", e);
+        ProcessSingleJobError::Other(anyhow!("Could not acquire DB pool: {:?}", e))
+      })?;
+
+  // download workflow if not none
+  match job_args.workflow_source {
+    Some(workflow_token) => {
+      let retrieved_workflow_record =  get_weight_by_token_with_transactor(
+        workflow_token,
+        false,
+        Transactor::for_connection(&mut mysql_connection)
+      ).await?.ok_or_else(|| ProcessSingleJobError::Other(anyhow!("Workflow not found")))?;
+
+      let bucket_details = RemoteCloudBucketDetails {
+        object_hash: retrieved_workflow_record.public_bucket_hash,
+        prefix: retrieved_workflow_record.maybe_public_bucket_prefix.unwrap(),
+        suffix: retrieved_workflow_record.maybe_public_bucket_extension.unwrap(),
+      };
+      remote_cloud_file_client.download_file(bucket_details, workflow_path.clone()).await?;
+      info!("Downloaded workflow to {:?}", workflow_path);
+    }
+    _ => { }
+  }
+
+  let maybe_args = job.maybe_inference_args
+      .as_ref()
+      .map(|args| args.args.as_ref())
+      .flatten();
+
+  let poly_args = match maybe_args {
+    None => return Err(ProcessSingleJobError::Other(anyhow!("ComfyUi args not found"))),
+    Some(args) => args,
+  };
+
+  let comfy_args = match poly_args {
+    Cu(args) => args,
+    _ => return Err(ProcessSingleJobError::Other(anyhow!("ComfyUi args not found"))),
+  };
+
+  info!("Workflow args: {:?}", comfy_args);
+
+  // ==================== WRITE WORKFLOW PROMPT ==================== //
+
+  workflow_path = write_workflow_prompt(WorkflowPromptArgs {
+    workflow_path: &workflow_path,
+    comfy_args: &comfy_args,
+    model_dependencies: &comfy_deps,
+    maybe_positive_prompt: comfy_args.positive_prompt.as_deref(),
+    maybe_negative_prompt: comfy_args.negative_prompt.as_deref(),
+  })?;
+
+  // ==================== QUERY AND DOWNLOAD FILES ==================== //
+
+  // Download Lora model if specified
+  let mut maybe_lora_path: Option<PathBuf> = None;
+
+  match job_args.maybe_lora_model {
+    Some(lora_model_weight_token) => {
+      let lora_dir = root_comfy_path.join("models").join("loras");
+      // make if not exist
+      if !lora_dir.exists() {
+        std::fs::create_dir_all(&lora_dir)
+            .map_err(|err| ProcessSingleJobError::IoError(err))?;
+      }
+
+      info!("Querying lora model by token: {:?} ...", &lora_model_weight_token);
+
+      let retrieved_lora_record =  get_weight_by_token_with_transactor(
+        lora_model_weight_token,
+        true,
+        Transactor::for_connection(&mut mysql_connection)
+      ).await?.ok_or_else(|| {
+        error!("Lora model not found: {:?}", lora_model_weight_token);
+        ProcessSingleJobError::Other(anyhow!("Lora model not found: {:?}", lora_model_weight_token))
+      })?;
+
+      let bucket_details = RemoteCloudBucketDetails {
+        object_hash: retrieved_lora_record.public_bucket_hash,
+        prefix: retrieved_lora_record.maybe_public_bucket_prefix.unwrap(),
+        suffix: retrieved_lora_record.maybe_public_bucket_extension.unwrap(),
+      };
+
+      let lora_filename = "lora.safetensors";
+      let lora_path = lora_dir.join(lora_filename).to_str().unwrap().to_string();
+      remote_cloud_file_client.download_file(bucket_details, lora_path.clone()).await?;
+      maybe_lora_path = Some(lora_path.parse().unwrap());
+      info!("Downloaded Lora model to {:?}", lora_path);
+      maybe_lora_path = Some(lora_path.parse().unwrap());
+    }
+    None => {}
+  }
+
+  let input_dir = root_comfy_path.join("input");
+
+  if !input_dir.exists() {
+    std::fs::create_dir_all(&input_dir)
+        .map_err(|err| ProcessSingleJobError::IoError(err))?;
+  }
+
+  // Keep track of all the video files we generate
+  info!("Root comfy path: {:?}", &root_comfy_path);
+  info!("Job output path will be (fix this code! the job shouldn't set this path!): {:?}", &job_args.output_path);
+
+  // TODO(bt,2024-04-20): Clean up this mess.
+
+  // ==================== DOWNLOAD GLOBAL IPA IMAGE (IF SET) ==================== //
+
+  let mut global_ipa_image = None;
+
+  if let Some(ipa_media_token) = comfy_args.global_ip_adapter_token.as_ref() {
+    let results = download_global_ipa_image(DownloadGlobalIpaImageArgs {
+      ipa_media_token,
+      comfy_input_directory: &input_dir,
+      remote_cloud_file_client: &remote_cloud_file_client,
+    }, Transactor::for_connection(&mut mysql_connection)).await?;
+
+    info!("Downloaded global IPA image to {:?}", results.ipa_image_path);
+
+    global_ipa_image = Some(results);
+  }
+
+  // ==================== DOWNLOAD VIDEOS ==================== //
+
+  let mut videos = download_input_videos(DownloadInputVideoArgs {
+    job_args: &job_args,
+    comfy_dirs: &comfy_dirs,
+    mysql_connection: &mut mysql_connection,
+    remote_cloud_file_client: &remote_cloud_file_client,
+  }).await?;
+
+  info!("Downloaded video!");
+
+  videos.debug_print_video_paths();
+
+  if let Ok(Some(dimensions)) = ffprobe_get_dimensions(&videos.primary_video.original_download_path) {
+    info!("Download video dimensions: {}x{}", dimensions.width, dimensions.height);
+  }
+
+  // ========================= TRIM AND PREPROCESS VIDEO ======================== //
+
+  let expected_frame_count = preprocess_trim_and_resample_videos(ProcessTrimAndResampleVideoArgs {
+    comfy_args,
+    comfy_deps,
+    comfy_dirs: &comfy_dirs,
+    videos: &mut videos,
+  })?;
+
+
+  // ========================= PREPROCESS AUDIO ======================== //
+
+  let mut lipsync_enabled = comfy_args.lipsync_enabled.unwrap_or(false);
+  if let Err(err) = preprocess_save_audio(ProcessSaveAudioArgs {
+    comfy_deps,
+    videos: &mut videos,
+  }) {
+    error!("Audio extraction failed: {:?}", err);
+
+    lipsync_enabled = false;
+  }
+
+  // ========================= CREATE OUTPUT DIR ======================== //
+
+  // make outputs dir if not exist
+  let output_dir = root_comfy_path.join("output");
+  if !output_dir.exists() {
+    std::fs::create_dir_all(&output_dir)
+        .map_err(|err| ProcessSingleJobError::IoError(err))?;
+  }
+
+  // ==================== RUN COMFY INFERENCE ==================== //
+
+  info!("Preparing for ComfyUI inference...");
+
+  job_progress_reporter.log_status("running inference")
+      .map_err(|e| ProcessSingleJobError::Other(e))?;
+
+  let stderr_output_file = work_temp_dir.path().join("stderr.txt");
+  let stdout_output_file = work_temp_dir.path().join("stdout.txt");
+
+  let positive_prompt_file = work_temp_dir.path().join("positive_prompt.txt");
+  let negative_prompt_file = work_temp_dir.path().join("negative_prompt.txt");
+  let travel_prompt_file = work_temp_dir.path().join("travel_prompt.txt");
+  let preview_frames_dir = work_temp_dir.path().join("preview_frames");
+
+  let maybe_positive_prompt_filename = comfy_args.positive_prompt
+      .as_deref()
+      .map(|prompt| {
+        std::fs::write(&positive_prompt_file, prompt)
+            .map(|_| positive_prompt_file.as_path())
+            .map_err(|e| ProcessSingleJobError::IoError(e))
+      })
+      .transpose()?;
+
+  let maybe_negative_prompt_filename = comfy_args.negative_prompt
+      .as_deref()
+      .map(|prompt| {
+        std::fs::write(&negative_prompt_file, prompt)
+            .map(|_| negative_prompt_file.as_path())
+            .map_err(|e| ProcessSingleJobError::IoError(e))
+      })
+      .transpose()?;
+
+  let maybe_travel_prompt_filename = comfy_args.travel_prompt
+      .as_deref()
+      .map(|prompt| {
+        std::fs::write(&travel_prompt_file, prompt)
+            .map(|_| travel_prompt_file.as_path())
+            .map_err(|e| ProcessSingleJobError::IoError(e))
+      })
+      .transpose()?;
+
+  let inference_details = InferenceDetails::NewPythonArgs {
+    maybe_style: comfy_args.style_name,
+    maybe_positive_prompt_filename,
+    maybe_negative_prompt_filename,
+    maybe_travel_prompt_filename,
+  };
+
+  let inference_start_time = Instant::now();
+
+  let ephemeral_bucket_client = deps.buckets.auto_gc_bucket_client.clone();
+  let previews_media_file_directory = MediaFileBucketDirectory::generate_new();
+
+
+  // let mysql_pool = deps.db.mysql_pool.clone();
+  let (_cancel_tx, mut cancel_rx) = oneshot::channel();
+  // let mut interval = tokio::time::interval(Duration::from_millis(30_000));
+  // let job_token = job.inference_job_token.clone();
+  //
+  // let _cancellation_watcher = tokio::spawn(async move {
+  //     let mut connection = match mysql_pool.acquire().await {
+  //         Ok(connection) => connection,
+  //         Err(e) => {
+  //             error!("Error acquiring mysql connection: {:?}", e);
+  //             return
+  //         }
+  //     };
+  //     let mut fail_count = 0;
+  //     loop {
+  //         interval.tick().await;
+  //         let current_status = get_inference_job_status_from_connection(
+  //             &job_token,
+  //             &mut connection
+  //         ).await;
+  //         let current_status = match current_status {
+  //             Ok(status) => status,
+  //             Err(e) => {
+  //                 fail_count += 1;
+  //                 error!("Error #{fail_count} getting job status: {:?}", e);
+  //                 if fail_count > 3 {
+  //                     return;
+  //                 }
+  //                 tokio::time::sleep(Duration::from_secs(3)).await;
+  //                 continue
+  //             }
+  //         };
+  //         fail_count = 0;
+  //         match current_status {
+  //             Some(status) => {
+  //                 if status.status == JobStatusPlus::CancelledByUser {
+  //                     cancel_tx.send(()).unwrap();
+  //                     break
+  //                 }
+  //             }
+  //             None => {
+  //                 debug!("Job status not found. This has to be an error");
+  //             }
+  //         }
+  //     }
+  // });
+  //
+  let preview_frames_enabled = comfy_args.generate_fast_previews.unwrap_or(false);
+  let (preview_cancellation_tx, rx) = channel();
+  let mut preview_processor = PreviewProcessor::new(
+    job.inference_job_token.clone(),
+    redis_pool,
+    preview_frames_dir.clone(),
+    previews_media_file_directory,
+    expected_frame_count as u32,
+    rx,
+  );
+
+  tokio::spawn(async move {
+    preview_processor.start_processing(&ephemeral_bucket_client).await;;
+  });
+
+  info!("Running ComfyUI inference...");
+
+  let command_exit_status = comfy_deps
+      .inference_command
+      .execute_inference(
+        &mut cancel_rx,
+        InferenceArgs {
+          generate_previews: comfy_args.generate_fast_previews.unwrap_or(false),
+          preview_frames_directory: Some(preview_frames_dir.as_ref()),
+          stderr_output_file: &stderr_output_file,
+          stdout_output_file: &stdout_output_file,
+          inference_details,
+          face_detailer_enabled: comfy_args.use_face_detailer.unwrap_or(false),
+          upscaler_enabled: comfy_args.use_upscaler.unwrap_or(false),
+          lipsync_enabled,
+          disable_lcm: comfy_args.disable_lcm.unwrap_or(false),
+          use_cinematic: comfy_args.use_cinematic.unwrap_or(false),
+          use_cogvideo: comfy_args.use_cogvideo.unwrap_or(false),
+          maybe_strength: comfy_args.strength,
+          frame_skip: comfy_args.frame_skip,
+          global_ipa_image_filename: global_ipa_image
+              .as_ref()
+              .map(|image| path_to_string(&image.ipa_image_path)),
+          global_ipa_strength: None, // TODO: Expose a UI slider
+          depth_video_path: videos.maybe_depth.as_ref()
+              .map(|v| v.maybe_processed_path.as_deref())
+              .flatten(),
+          normal_video_path: videos.maybe_normal.as_ref()
+              .map(|v| v.maybe_processed_path.as_deref())
+              .flatten(),
+          outline_video_path: videos.maybe_outline.as_ref()
+              .map(|v| v.maybe_processed_path.as_deref())
+              .flatten(),
+        }).await;
+
+  let inference_duration = Instant::now().duration_since(inference_start_time);
+
+  info!("Inference command exited with status: {:?}", command_exit_status);
+
+  info!("Inference took duration to complete: {:?}", &inference_duration);
+
+  // check stdout for success and check if file exists
+  if let Ok(contents) = read_to_string(&stdout_output_file) {
+    info!("Captured stdout output: {}", contents);
+  }
+
+  if let Ok(contents) = read_to_string(&stderr_output_file) {
+    info!("Captured stderr output: {}", contents);
+  }
+
+  videos.debug_print_video_paths();
+
+  if let Ok(Some(dimensions)) = ffprobe_get_dimensions(&videos.primary_video.comfy_output_video_path) {
+    info!("Comfy output video dimensions: {}x{}", dimensions.width, dimensions.height);
+  }
+
+  let _ = preview_cancellation_tx.send(());
+
+  // ==================== CHECK OUTPUT FILE ======================== //
+
+  if let Err(err) = check_file_exists(&videos.primary_video.comfy_output_video_path) {
+    error!("Output file does not  exist: {:?}", err);
+
+    error!("Inference failed: {:?}", command_exit_status);
+
+    if let Ok(contents) = read_to_string(&stderr_output_file) {
+      warn!("Captured stderr output: {}", contents);
+    }
+
+    safe_delete_all_input_videos(&videos);
+    safe_recursively_delete_files(&comfy_dirs.comfy_output_dir); // NB: Don't delete directory itself.
+
+    // NB: Forcing generic type to `&Path` with turbofish
+    safe_delete_possible_files_and_directories::<&Path>(&[
+      Some(&negative_prompt_file),
+      Some(&output_dir),
+      Some(&positive_prompt_file),
+      Some(&preview_frames_dir),
+      Some(&stderr_output_file),
+      Some(&stdout_output_file),
+      Some(&travel_prompt_file),
+      Some(work_temp_dir.path()),
+      Some(workflow_path.as_ref()),
+      global_ipa_image.as_ref().map(|ipa| ipa.ipa_image_path.as_path()),
+      maybe_lora_path.as_deref(),
+    ]);
+
+    return Err(ProcessSingleJobError::Other(anyhow!("Output file did not exist: {:?}",
+            &videos.primary_video.comfy_output_video_path)));
+  }
+
+  // ==================== COPY BACK AUDIO ==================== //
+
+  post_process_restore_audio(PostProcessRestoreVideoArgs {
+    comfy_deps,
+    videos: &mut videos,
+  });
+
+  // ==================== OPTIONAL WATERMARK ==================== //
+
+  post_process_add_watermark(PostProcessAddWatermarkArgs {
+    comfy_deps,
+    videos: &mut videos,
+  });
+
+  // ==================== DEBUG ======================== //
+
+  videos.debug_print_video_paths();
+
+  if let Ok(Some(dimensions)) = ffprobe_get_dimensions(&videos.primary_video.get_final_video_to_upload()) {
+    info!("Final video upload dimensions: {}x{}", dimensions.width, dimensions.height);
+  }
+
+  // ==================== VALIDATE AND SAVE RESULTS ======================== //
+
+  let result = validate_and_save_results(SaveResultsArgs {
+    job,
+    deps: &deps,
+    job_args: &job_args,
+    comfy_deps,
+    comfy_args,
+    videos: &videos,
+    job_progress_reporter: &mut job_progress_reporter,
+    inference_duration,
+  }).await;
+
+  let media_file_token = match result {
+    Ok(token) => token,
+    Err(err) => {
+      error!("Error validating and saving results: {:?}", err);
+      safe_delete_all_input_videos(&videos);
+      safe_recursively_delete_files(&comfy_dirs.comfy_output_dir); // NB: Don't delete directory itself.
+
+      // NB: Forcing generic type to `&Path` with turbofish
+      safe_delete_possible_files_and_directories::<&Path>(&[
+        Some(&comfy_dirs.comfy_output_dir),
+        Some(&negative_prompt_file),
+        Some(&output_dir),
+        Some(&preview_frames_dir),
+        Some(&stderr_output_file),
+        Some(&stdout_output_file),
+        Some(&travel_prompt_file),
+        Some(&positive_prompt_file),
+        Some(work_temp_dir.path()),
+        Some(workflow_path.as_ref()),
+        global_ipa_image.as_ref().map(|ipa| ipa.ipa_image_path.as_path()),
+        maybe_lora_path.as_deref(),
+      ]);
+
+      return Err(err);
+    }
+  };
+
+  // ==================== (OPTIONAL) DEBUG SLEEP ==================== //
+
+  if let Some(sleep_millis) = comfy_args.sleep_millis {
+    info!("Sleeping for millis: {sleep_millis}");
+    thread::sleep(Duration::from_millis(sleep_millis));
+  }
+
+  // ==================== CLEANUP/ DELETE TEMP FILES ==================== //
+
+  info!("Cleaning up temporary files...");
+
+  safe_delete_all_input_videos(&videos);
+  safe_recursively_delete_files(&comfy_dirs.comfy_output_dir); // NB: Don't delete directory itself.
+
+  // NB: Forcing generic type to `&Path` with turbofish
+  safe_delete_possible_files_and_directories::<&Path>(&[
+    Some(&comfy_dirs.comfy_output_dir),
+    Some(&negative_prompt_file),
+    Some(&output_dir),
+    Some(&preview_frames_dir),
+    Some(&stderr_output_file),
+    Some(&stdout_output_file),
+    Some(&travel_prompt_file),
+    Some(&positive_prompt_file),
+    Some(work_temp_dir.path()),
+    Some(workflow_path.as_ref()),
+    global_ipa_image.as_ref().map(|ipa| ipa.ipa_image_path.as_path()),
+    maybe_lora_path.as_deref(),
+  ]);
+
+  // ==================== DONE ==================== //
+
+  info!("ComfyUI Done.");
+
+  job_progress_reporter.log_status("done")
+      .map_err(|e| ProcessSingleJobError::Other(e))?;
+
+  info!("Result video media token: {:?}", &media_file_token);
+
+  info!("Job {:?} complete success!", job.id);
+
+  Ok(JobSuccessResult {
+    maybe_result_entity: Some(ResultEntity {
+      entity_type: InferenceResultType::MediaFile,
+      entity_token: media_file_token.to_string(),
+    }),
+    inference_duration,
+  })
+}
+
+fn safe_delete_all_input_videos(videos: &VideoPathing) {
+  safe_delete_primary_videos(&videos.primary_video);
+  if let Some(depth) = &videos.maybe_depth {
+    safe_delete_secondary_videos(depth);
+  }
+  if let Some(normal) = &videos.maybe_normal {
+    safe_delete_secondary_videos(normal);
+  }
+  if let Some(outline) = &videos.maybe_outline {
+    safe_delete_secondary_videos(outline);
+  }
+}
+
+fn safe_delete_primary_videos(video: &PrimaryInputVideoAndPaths) {
+  safe_delete_possible_files_and_directories::<&Path>(&[
+    Some(&video.comfy_input_video_path),
+    Some(&video.comfy_output_video_path),
+    Some(&video.original_download_path),
+    Some(&video.trimmed_wav_audio_path),
+    Some(video.get_final_video_to_upload()),
+    Some(video.get_non_watermarked_video_to_upload()),
+    Some(video.video_to_watermark()),
+    video.audio_restored_video_path.as_deref(),
+    video.maybe_trimmed_resampled_path.as_deref(),
+    video.watermarked_video_path.as_deref(),
+  ]);
+}
+
+fn safe_delete_secondary_videos(video: &SecondaryInputVideoAndPaths) {
+  safe_delete_possible_files_and_directories::<&Path>(&[
+    Some(&video.original_download_path),
+    video.maybe_processed_path.as_deref(),
+  ]);
+}