download, process, and pass (depth,normal,outline) map videos

crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/comfy_ui_inference_command.rs

@@ -103,6 +103,10 @@ pub struct InferenceArgs<'s> {
 
     pub global_ipa_image_filename: Option<String>,
     pub global_ipa_strength: Option<f32>,
+
+    pub depth_video_path: Option<&'s Path>,
+    pub normal_video_path: Option<&'s Path>,
+    pub outline_video_path: Option<&'s Path>,
 }
 
 #[derive(Debug)]
@@ -335,6 +339,24 @@ impl ComfyInferenceCommand {
             command.push_str(" ");
         }
 
+        if let Some(depth_video_path) = args.depth_video_path {
+            command.push_str(" --depth_video_filename ");
+            command.push_str(&path_to_string(depth_video_path));
+            command.push_str(" ");
+        }
+
+        if let Some(normal_video_path) = args.normal_video_path {
+            command.push_str(" --normals_video_filename ");
+            command.push_str(&path_to_string(normal_video_path));
+            command.push_str(" ");
+        }
+
+        if let Some(outline_video_path) = args.outline_video_path {
+            command.push_str(" --outline_video_filename ");
+            command.push_str(&path_to_string(outline_video_path));
+            command.push_str(" ");
+        }
+
         if let Some(docker_options) = self.maybe_docker_options.as_ref() {
             command = docker_options.to_command_string(&command);
         }

crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/process_job.rs

@@ -51,7 +51,7 @@ use crate::job::job_types::workflow::comfy_ui::video_style_transfer::steps::down
 use crate::job::job_types::workflow::comfy_ui::video_style_transfer::steps::post_process_add_watermark::{post_process_add_watermark, PostProcessAddWatermarkArgs};
 use crate::job::job_types::workflow::comfy_ui::video_style_transfer::steps::post_process_restore_audio::{post_process_restore_audio, PostProcessRestoreVideoArgs};
 use crate::job::job_types::workflow::comfy_ui::video_style_transfer::steps::preprocess_save_audio::{preprocess_save_audio, ProcessSaveAudioArgs};
-use crate::job::job_types::workflow::comfy_ui::video_style_transfer::steps::preprocess_trim_and_resample_video::{preprocess_trim_and_resample_video, ProcessTrimAndResampleVideoArgs};
+use crate::job::job_types::workflow::comfy_ui::video_style_transfer::steps::preprocess_trim_and_resample_videos::{preprocess_trim_and_resample_videos, ProcessTrimAndResampleVideoArgs};
 use crate::job::job_types::workflow::comfy_ui::video_style_transfer::steps::validate_and_save_results::{SaveResultsArgs, validate_and_save_results};
 use crate::job::job_types::workflow::comfy_ui::video_style_transfer::video_paths::VideoPaths;
 use crate::job::job_types::workflow::comfy_ui::video_style_transfer::write_workflow_prompt::{WorkflowPromptArgs, write_workflow_prompt};
@@ -238,7 +238,7 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
 
     // ==================== DOWNLOAD VIDEO ==================== //
 
-    let download_video = download_input_videos(DownloadInputVideoArgs {
+    let mut download_videos = download_input_videos(DownloadInputVideoArgs {
         job_args: &job_args,
         videos: &videos,
         mysql_pool: &deps.db.mysql_pool,
@@ -255,10 +255,11 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
 
     // ========================= TRIM AND PREPROCESS VIDEO ======================== //
 
-    preprocess_trim_and_resample_video(ProcessTrimAndResampleVideoArgs {
+    preprocess_trim_and_resample_videos(ProcessTrimAndResampleVideoArgs {
         comfy_args,
         comfy_deps: model_dependencies,
-        videos: &videos,
+        primary_video_paths: &videos,
+        download_videos: &mut download_videos,
     })?;
 
     // ========================= PREPROCESS AUDIO ======================== //
@@ -351,6 +352,15 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
                 .as_ref()
                 .map(|image| path_to_string(&image.ipa_image_path)),
             global_ipa_strength: None, // TODO: Expose a UI slider
+            depth_video_path: download_videos.maybe_depth.as_ref()
+                .map(|v| v.maybe_processed_path.as_deref())
+                .flatten(),
+            normal_video_path: download_videos.maybe_normal.as_ref()
+                .map(|v| v.maybe_processed_path.as_deref())
+                .flatten(),
+            outline_video_path: download_videos.maybe_outline.as_ref()
+                .map(|v| v.maybe_processed_path.as_deref())
+                .flatten(),
         });
 
     let inference_duration = Instant::now().duration_since(inference_start_time);
@@ -439,7 +449,7 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
         comfy_args,
         videos: &videos,
         job_progress_reporter: &mut job_progress_reporter,
-        download_video,
+        download_videos,
         inference_duration,
     }).await?;
 

crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/steps/mod.rs

@@ -5,5 +5,5 @@ pub mod download_input_videos;
 pub mod post_process_add_watermark;
 pub mod post_process_restore_audio;
 pub mod preprocess_save_audio;
-pub mod preprocess_trim_and_resample_video;
+pub mod preprocess_trim_and_resample_videos;
 pub mod validate_and_save_results;

crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/steps/preprocess_trim_and_resample_video.rs

@@ -1,106 +0,0 @@
-use std::path::Path;
-use std::process::{Command, Stdio};
-
-use anyhow::anyhow;
-use log::{error, info};
-use sqlx::MySqlPool;
-
-use cloud_storage::remote_file_manager::remote_cloud_file_manager::RemoteCloudFileClient;
-use filesys::path_to_string::path_to_string;
-use mysql_queries::payloads::generic_inference_args::workflow_payload::WorkflowArgs;
-use tokens::tokens::media_files::MediaFileToken;
-use videos::ffprobe_get_dimensions::ffprobe_get_dimensions;
-
-use crate::job::job_loop::process_single_job_error::ProcessSingleJobError;
-use crate::job::job_types::workflow::comfy_ui::comfy_ui_dependencies::ComfyDependencies;
-use crate::job::job_types::workflow::comfy_ui::video_style_transfer::video_paths::VideoPaths;
-
-pub struct ProcessTrimAndResampleVideoArgs<'a> {
-  pub comfy_args: &'a WorkflowArgs,
-  pub comfy_deps: &'a ComfyDependencies,
-  pub videos: &'a VideoPaths,
-}
-
-pub fn preprocess_trim_and_resample_video(
-  args: ProcessTrimAndResampleVideoArgs<'_>
-) -> Result<(), ProcessSingleJobError> {
-  let target_fps = args.comfy_args.target_fps.unwrap_or(24);
-
-  let trim_start_millis = args.comfy_args.trim_start_milliseconds
-      .or_else(|| args.comfy_args.trim_start_seconds.map(|s| s as u64 * 1_000))
-      .unwrap_or(0);
-
-  let trim_end_millis = args.comfy_args.trim_end_milliseconds
-      .or_else(|| args.comfy_args.trim_end_seconds.map(|s| s as u64 * 1_000))
-      .unwrap_or(3_000);
-
-  info!("trim start millis: {trim_start_millis}");
-  info!("trim end millis: {trim_end_millis}");
-  info!("target FPS: {target_fps}");
-
-  let skip_process_video = args.comfy_args.skip_process_video.unwrap_or(false);
-
-  if skip_process_video {
-    info!("Skipping video trim / resample...");
-    info!("(This might break if we need to copy the video path. Salt's code implicitly expects videos to be in certain places, but doesn't allow passing of config, and that's horrible.)");
-
-    std::fs::copy(&args.videos.original_video_path, &args.videos.trimmed_resampled_video_path)
-        .map_err(|err| {
-          error!("Error copying video (1): {:?}", err);
-          ProcessSingleJobError::IoError(err)
-        })?;
-
-    std::fs::copy(&args.videos.original_video_path, &args.videos.comfy_output_video_path)
-        .map_err(|err| {
-          error!("Error copying video (2): {:?}", err);
-          ProcessSingleJobError::IoError(err)
-        })?;
-
-  } else {
-    info!("Calling video trim / resample...");
-    info!("Script: {:?}", &args.comfy_deps.inference_command.processing_script);
-
-    // shell out to python script
-    let output = Command::new("python3")
-        .stdout(Stdio::inherit()) // NB: This should emit to the rust job's stdout
-        .stderr(Stdio::inherit()) // NB: This should emit to the rust job's stderr
-        .arg(path_to_string(&args.comfy_deps.inference_command.processing_script))
-        .arg(path_to_string(&args.videos.original_video_path))
-        .arg(format!("{:?}", trim_start_millis))
-        .arg(format!("{:?}", trim_end_millis))
-        .arg(format!("{:?}", target_fps))
-        .output()
-        .map_err(|e| {
-          error!("Error running inference: {:?}", e);
-          ProcessSingleJobError::Other(e.into())
-        })?;
-
-    // check if the command was successful
-    if !output.status.success() {
-      // print stdout and stderr
-      error!("Video processing failed: {:?}", output.status);
-      error!("stdout: {}", String::from_utf8_lossy(&output.stdout));
-      error!("stderr: {}", String::from_utf8_lossy(&output.stderr));
-      return Err(ProcessSingleJobError::Other(anyhow!("Command failed: {:?}", output.status)));
-    }
-
-    info!("Finished video trim / resample.");
-
-    // NB: The process video script implicitly saves the above video as "input.mp4"
-    // Comfy sometimes overwrites this, so we need to make a copy.
-    std::fs::copy(&args.videos.comfy_input_video_path, &args.videos.trimmed_resampled_video_path)
-        .map_err(|err| {
-          error!("Error copying trimmed video: {:?}", err);
-          ProcessSingleJobError::IoError(err)
-        })?;
-  }
-
-  args.videos.debug_print_paths_after_trim();
-
-  if let Ok(Some(dimensions)) = ffprobe_get_dimensions(&args.videos.trimmed_resampled_video_path) {
-    info!("Trimmed / resampled video dimensions: {}x{}", dimensions.width, dimensions.height);
-  }
-
-  Ok(())
-}
-

crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/steps/preprocess_trim_and_resample_videos.rs

@@ -0,0 +1,227 @@
+use std::path::{Path, PathBuf};
+use std::process::{Command, Stdio};
+
+use anyhow::anyhow;
+use log::{error, info};
+use sqlx::MySqlPool;
+
+use cloud_storage::remote_file_manager::remote_cloud_file_manager::RemoteCloudFileClient;
+use filesys::path_to_string::path_to_string;
+use mysql_queries::payloads::generic_inference_args::workflow_payload::WorkflowArgs;
+use tokens::tokens::media_files::MediaFileToken;
+use videos::ffprobe_get_dimensions::ffprobe_get_dimensions;
+
+use crate::job::job_loop::process_single_job_error::ProcessSingleJobError;
+use crate::job::job_types::workflow::comfy_ui::comfy_ui_dependencies::ComfyDependencies;
+use crate::job::job_types::workflow::comfy_ui::video_style_transfer::steps::download_input_videos::VideoDownloadDetails;
+use crate::job::job_types::workflow::comfy_ui::video_style_transfer::steps::secondary_video_and_paths::SecondaryVideoAndPaths;
+use crate::job::job_types::workflow::comfy_ui::video_style_transfer::video_paths::VideoPaths;
+
+pub struct ProcessTrimAndResampleVideoArgs<'a> {
+  pub comfy_args: &'a WorkflowArgs,
+  pub comfy_deps: &'a ComfyDependencies,
+
+  // TODO(bt,2024-07-09): This is going to need refatoring and consolidation for clarity.
+  //  'VideoPaths' was for the original video and all of its resamplings, outputs,
+  //  watermarkings, etc. Now that we have depth, normal, and outline map videos (and
+  //  possibly more), we're going to need to group these. Those files are referenced by
+  //  `VideoDownloadDetails`, which is a better system to keep track of the videos as they
+  //  pass through the system.
+  pub primary_video_paths: &'a VideoPaths,
+  pub download_videos: &'a mut VideoDownloadDetails,
+}
+
+pub fn preprocess_trim_and_resample_videos(
+  args: ProcessTrimAndResampleVideoArgs<'_>
+) -> Result<(), ProcessSingleJobError> {
+  let target_fps = args.comfy_args.target_fps.unwrap_or(24);
+
+  let trim_start_millis = args.comfy_args.trim_start_milliseconds
+      .or_else(|| args.comfy_args.trim_start_seconds.map(|s| s as u64 * 1_000))
+      .unwrap_or(0);
+
+  let trim_end_millis = args.comfy_args.trim_end_milliseconds
+      .or_else(|| args.comfy_args.trim_end_seconds.map(|s| s as u64 * 1_000))
+      .unwrap_or(3_000);
+
+  info!("trim start millis: {trim_start_millis}");
+  info!("trim end millis: {trim_end_millis}");
+  info!("target FPS: {target_fps}");
+
+  let resample_details = ResampleDetails {
+    target_fps,
+    trim_start_millis,
+    trim_end_millis,
+  };
+
+  let skip_resampling_video = args.comfy_args.skip_process_video.unwrap_or(false);
+
+  preprocess_trim_and_resample_primary_video(&args, &resample_details, skip_resampling_video)?;
+
+  if !skip_resampling_video {
+    preprocess_trim_and_resample_secondary_videos(&args.comfy_deps, &resample_details, args.download_videos)?;
+  }
+
+  Ok(())
+}
+
+struct ResampleDetails {
+  target_fps: u32,
+  trim_start_millis: u64,
+  trim_end_millis: u64,
+}
+
+fn preprocess_trim_and_resample_primary_video(
+  args: &ProcessTrimAndResampleVideoArgs<'_>,
+  resample_details: &ResampleDetails,
+  skip_process_video: bool,
+) -> Result<(), ProcessSingleJobError> {
+
+  if skip_process_video {
+    info!("Skipping video trim / resample...");
+    info!("(This might break if we need to copy the video path. Salt's code implicitly expects videos to be in certain places, but doesn't allow passing of config, and that's horrible.)");
+
+    std::fs::copy(&args.primary_video_paths.original_video_path, &args.primary_video_paths.trimmed_resampled_video_path)
+        .map_err(|err| {
+          error!("Error copying video (1): {:?}", err);
+          ProcessSingleJobError::IoError(err)
+        })?;
+
+    std::fs::copy(&args.primary_video_paths.original_video_path, &args.primary_video_paths.comfy_output_video_path)
+        .map_err(|err| {
+          error!("Error copying video (2): {:?}", err);
+          ProcessSingleJobError::IoError(err)
+        })?;
+
+  } else {
+    info!("Calling video trim / resample...");
+    info!("Script: {:?}", &args.comfy_deps.inference_command.processing_script);
+
+    // shell out to python script
+    let output = Command::new("python3")
+        .stdout(Stdio::inherit()) // NB: This should emit to the rust job's stdout
+        .stderr(Stdio::inherit()) // NB: This should emit to the rust job's stderr
+        .arg(path_to_string(&args.comfy_deps.inference_command.processing_script))
+        .arg("DO_NOT_USE_THIS_ARG_ANYMORE")
+        .arg(format!("{:?}", resample_details.trim_start_millis))
+        .arg(format!("{:?}", resample_details.trim_end_millis))
+        .arg(format!("{:?}", resample_details.target_fps))
+        .arg("--input")
+        .arg(path_to_string(&args.primary_video_paths.original_video_path))
+        .arg("--output")
+        // NB(bt,2024-07-09): Despite what the comments on this field say, the script `format_video.py` writes
+        // to a file named 'input.mp4', not 'trimmed.mp4'. This pathing really needs to be cleaned up.
+        .arg(path_to_string(&args.primary_video_paths.comfy_input_video_path))
+        .output()
+        .map_err(|e| {
+          error!("Error running inference: {:?}", e);
+          ProcessSingleJobError::Other(e.into())
+        })?;
+
+    // check if the command was successful
+    if !output.status.success() {
+      // print stdout and stderr
+      error!("Video processing failed: {:?}", output.status);
+      error!("stdout: {}", String::from_utf8_lossy(&output.stdout));
+      error!("stderr: {}", String::from_utf8_lossy(&output.stderr));
+      return Err(ProcessSingleJobError::Other(anyhow!("Command failed: {:?}", output.status)));
+    }
+
+    info!("Finished video trim / resample.");
+
+    // NB: The process video script implicitly saves the above video as "input.mp4"
+    // Comfy sometimes overwrites this, so we need to make a copy.
+    std::fs::copy(&args.primary_video_paths.comfy_input_video_path, &args.primary_video_paths.trimmed_resampled_video_path)
+        .map_err(|err| {
+          error!("Error copying trimmed video: {:?}", err);
+          ProcessSingleJobError::IoError(err)
+        })?;
+  }
+
+  args.primary_video_paths.debug_print_paths_after_trim();
+
+  if let Ok(Some(dimensions)) = ffprobe_get_dimensions(&args.primary_video_paths.trimmed_resampled_video_path) {
+    info!("Trimmed / resampled video dimensions: {}x{}", dimensions.width, dimensions.height);
+  }
+
+  Ok(())
+}
+
+fn preprocess_trim_and_resample_secondary_videos(
+  comfy_deps: &ComfyDependencies,
+  resample_details: &ResampleDetails,
+  secondary_videos: &mut VideoDownloadDetails,
+) -> Result<(), ProcessSingleJobError> {
+
+  if let Some(depth) = secondary_videos.maybe_depth.as_mut() {
+    let resampled_path = preprocess_trim_and_resample_secondary_video(comfy_deps, resample_details, depth)?;
+    depth.maybe_processed_path = Some(resampled_path);
+  }
+
+  if let Some(normal) = secondary_videos.maybe_normal.as_mut() {
+    let resampled_path = preprocess_trim_and_resample_secondary_video(comfy_deps, resample_details, normal)?;
+    normal.maybe_processed_path = Some(resampled_path);
+  }
+
+  if let Some(outline) = secondary_videos.maybe_outline.as_mut() {
+    let resampled_path = preprocess_trim_and_resample_secondary_video(comfy_deps, resample_details, outline)?;
+    outline.maybe_processed_path = Some(resampled_path);
+  }
+
+  Ok(())
+}
+
+fn preprocess_trim_and_resample_secondary_video(
+  comfy_deps: &ComfyDependencies,
+  resample_details: &ResampleDetails,
+  video: &SecondaryVideoAndPaths,
+) -> Result<PathBuf, ProcessSingleJobError> {
+  info!("Calling video trim / resample...");
+  info!("Script: {:?}", &comfy_deps.inference_command.processing_script);
+
+  let output_path = {
+    let filename = video.original_download_path
+        .file_stem()
+        .map(|s| s.to_str())
+        .flatten()
+        .ok_or_else(|| {
+          ProcessSingleJobError::Other(anyhow!("Error getting filename: {:?}", video.original_download_path))
+        })?;
+
+    let filename = format!("{}_resampled.mp4", filename);
+
+    video.original_download_path.with_file_name(filename)
+  };
+
+  // shell out to python script
+  let output = Command::new("python3")
+      .stdout(Stdio::inherit()) // NB: This should emit to the rust job's stdout
+      .stderr(Stdio::inherit()) // NB: This should emit to the rust job's stderr
+      .arg(path_to_string(&comfy_deps.inference_command.processing_script))
+      .arg(path_to_string("DO_NOT_USE_THIS_ARG_ANYMORE"))
+      .arg(format!("{:?}", resample_details.trim_start_millis))
+      .arg(format!("{:?}", resample_details.trim_end_millis))
+      .arg(format!("{:?}", resample_details.target_fps))
+      .arg("--input")
+      .arg(format!("{:?}", &video.original_download_path))
+      .arg("--output")
+      .arg(format!("{:?}", &output_path))
+      .output()
+      .map_err(|e| {
+        error!("Error running inference: {:?}", e);
+        ProcessSingleJobError::Other(e.into())
+      })?;
+
+  // check if the command was successful
+  if !output.status.success() {
+    // print stdout and stderr
+    error!("Video processing failed: {:?}", output.status);
+    error!("stdout: {}", String::from_utf8_lossy(&output.stdout));
+    error!("stderr: {}", String::from_utf8_lossy(&output.stderr));
+    return Err(ProcessSingleJobError::Other(anyhow!("Command failed: {:?}", output.status)));
+  }
+
+  info!("Finished video trim / resample.");
+
+  Ok(output_path)
+}

crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/steps/validate_and_save_results.rs

@@ -55,7 +55,7 @@ pub struct SaveResultsArgs<'a> {
   pub comfy_deps: &'a ComfyDependencies,
   pub job_progress_reporter: &'a mut Box<dyn JobProgressReporter>,
 
-  pub download_video: VideoDownloadDetails,
+  pub download_videos: VideoDownloadDetails,
   pub videos: &'a VideoPaths,
   pub inference_duration: Duration,
 
@@ -200,11 +200,11 @@ pub async fn validate_and_save_results(args: SaveResultsArgs<'_>) -> Result<Medi
   // This also lets us hide the engine renders from users.
   // This shouldn't ever become a deeply nested tree of children, but rather a single root
   // with potentially many direct children.
-  let style_transfer_source_media_file_token = args.download_video
+  let style_transfer_source_media_file_token = args.download_videos
       .input_video_media_file
       .maybe_style_transfer_source_media_file_token
       .as_ref()
-      .unwrap_or_else(|| &args.download_video.input_video_media_file.token);
+      .unwrap_or_else(|| &args.download_videos.input_video_media_file.token);
 
   let prompt_token = PromptToken::generate();
 
@@ -212,9 +212,12 @@ pub async fn validate_and_save_results(args: SaveResultsArgs<'_>) -> Result<Medi
     pool: &args.deps.db.mysql_pool,
     job: &args.job,
     maybe_mime_type: Some(&mimetype),
-    maybe_title: args.download_video.input_video_media_file.maybe_title.as_deref(),
+    maybe_title: args.download_videos.input_video_media_file.maybe_title.as_deref(),
     maybe_style_transfer_source_media_file_token: Some(&style_transfer_source_media_file_token),
-    maybe_scene_source_media_file_token: args.download_video.input_video_media_file.maybe_scene_source_media_file_token.as_ref(),
+    maybe_scene_source_media_file_token: args.download_videos
+        .input_video_media_file
+        .maybe_scene_source_media_file_token
+        .as_ref(),
     file_size_bytes,
     sha256_checksum: &file_checksum,
     maybe_prompt_token: Some(&prompt_token),