diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/comfy_ui_inference_command.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/comfy_ui_inference_command.rs index ae997192f2..befa786226 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/comfy_ui_inference_command.rs +++ b/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, pub global_ipa_strength: Option, + + 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); } diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/process_job.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/process_job.rs index 515e5b7437..97c7d5d349 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/process_job.rs +++ b/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) -> Result) -> Result) -> Result { - 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(()) -} - diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/steps/preprocess_trim_and_resample_videos.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/steps/preprocess_trim_and_resample_videos.rs new file mode 100644 index 0000000000..c2a719bd7e --- /dev/null +++ b/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 { + 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) +} diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/steps/validate_and_save_results.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/steps/validate_and_save_results.rs index 0f634c164e..5f7f2f9300 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/steps/validate_and_save_results.rs +++ b/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, - 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) -> Result