From b3da60aa3a680ef60d5b5eccb3fec04931b2ab6d Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Wed, 10 Jul 2024 20:58:39 -0400 Subject: [PATCH] cleanup pathing (still wip) --- .../video_style_transfer/process_job.rs | 96 ++++++------- .../steps/download_input_videos.rs | 27 ++-- .../steps/post_process_add_watermark.rs | 8 +- .../steps/post_process_restore_audio.rs | 19 ++- .../steps/preprocess_save_audio.rs | 16 ++- .../preprocess_trim_and_resample_videos.rs | 36 ++--- .../steps/validate_and_save_results.rs | 19 ++- .../video_style_transfer/util/comfy_dirs.rs | 24 ++++ .../comfy_ui/video_style_transfer/util/mod.rs | 1 + .../util/video_pathing.rs | 131 ++++++++++-------- 10 files changed, 220 insertions(+), 157 deletions(-) create mode 100644 crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/util/comfy_dirs.rs 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 837b0b06e9..3e4452a1ee 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 @@ -52,8 +52,9 @@ use crate::job::job_types::workflow::comfy_ui::video_style_transfer::steps::post 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_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::util::comfy_dirs::ComfyDirs; use crate::job::job_types::workflow::comfy_ui::video_style_transfer::util::comfy_ui_inference_command::{InferenceArgs, InferenceDetails}; -use crate::job::job_types::workflow::comfy_ui::video_style_transfer::util::video_pathing::{SecondaryInputVideoAndPaths, VideoDownloads}; +use crate::job::job_types::workflow::comfy_ui::video_style_transfer::util::video_pathing::{PrimaryInputVideoAndPaths, SecondaryInputVideoAndPaths, VideoDownloads}; use crate::job::job_types::workflow::comfy_ui::video_style_transfer::util::video_pathing_deprecated::VideoPaths; use crate::job::job_types::workflow::comfy_ui::video_style_transfer::util::write_workflow_prompt::{WorkflowPromptArgs, write_workflow_prompt}; use crate::job_dependencies::JobDependencies; @@ -70,7 +71,7 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result) -> Result) -> Result) -> Result) -> Result) -> Result) -> Result) -> Result) -> Result) -> Result) -> Result) -> Result) -> Result) -> Result { pub job_args: &'a JobArgs<'a>, - pub videos: &'a VideoPaths, + pub comfy_dirs: &'a ComfyDirs, pub mysql_pool: &'a MySqlPool, pub remote_cloud_file_client: &'a RemoteCloudFileClient, } @@ -35,7 +37,7 @@ pub async fn download_input_videos( // TODO: Consolidate primary video download with secondary download logic. async fn download_primary_video( - args: &DownloadInputVideoArgs<'_> + args: &DownloadInputVideoArgs<'_>, ) -> Result { let input_media_file_token = match args.job_args.maybe_input_file { None => return Err(ProcessSingleJobError::InvalidJob(anyhow!("No primary input video file provided"))), @@ -81,7 +83,7 @@ async fn download_primary_video( // NB(bt,2024-07-09): This convention is muddled with the python side. // We may not have flexibility to change this pathing for a while. - let download_path = args.videos.comfy_input_dir.join("video.mp4"); + let download_path = args.comfy_dirs.comfy_input_dir.join("video.mp4"); info!("Downloading primary input file to {:?}", download_path); @@ -92,8 +94,13 @@ async fn download_primary_video( info!("Downloaded primary input video!"); + // TODO: This monumentally sucks. + // The upstream shouldn't be telling us what to do about this at all. + let job_output_path = args.job_args.output_path; + Ok(VideoDownloads { - input_video: PrimaryInputVideoAndPaths::new(input_media_file, download_path), + input_video: PrimaryInputVideoAndPaths::new( + input_media_file, &args.comfy_dirs, job_output_path), maybe_depth: None, maybe_normal: None, maybe_outline: None, @@ -169,9 +176,9 @@ async fn download_secondary_video( }; let download_path = match secondary_video_type { - SecondaryVideoType::Depth => args.videos.comfy_input_dir.join("depth_download.mp4"), - SecondaryVideoType::Normal => args.videos.comfy_input_dir.join("normal_download.mp4"), - SecondaryVideoType::Outline => args.videos.comfy_input_dir.join("outline_download.mp4"), + SecondaryVideoType::Depth => args.comfy_dirs.comfy_input_dir.join("depth_download.mp4"), + SecondaryVideoType::Normal => args.comfy_dirs.comfy_input_dir.join("normal_download.mp4"), + SecondaryVideoType::Outline => args.comfy_dirs.comfy_input_dir.join("outline_download.mp4"), }; let maybe_media_file = all_video_media_files @@ -202,7 +209,7 @@ async fn download_secondary_video( ).await?; Ok(Some(SecondaryInputVideoAndPaths { - record: VideoMediaFileRecord::Bulk(media_file.clone()), + record: media_file.clone(), original_download_path: download_path, maybe_processed_path: None, })) diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/steps/post_process_add_watermark.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/steps/post_process_add_watermark.rs index 128738fc1e..57e0f4cda9 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/steps/post_process_add_watermark.rs +++ b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/steps/post_process_add_watermark.rs @@ -3,12 +3,13 @@ use log::{error, info}; use filesys::check_file_exists::check_file_exists; use crate::job::job_types::workflow::comfy_ui::comfy_ui_dependencies::ComfyDependencies; +use crate::job::job_types::workflow::comfy_ui::video_style_transfer::util::video_pathing::VideoDownloads; use crate::job::job_types::workflow::comfy_ui::video_style_transfer::util::video_pathing_deprecated::VideoPaths; use crate::util::common_commands::ffmpeg_logo_watermark_command::WatermarkArgs; pub struct PostProcessAddWatermarkArgs<'a> { pub comfy_deps: &'a ComfyDependencies, - pub videos: &'a mut VideoPaths, + pub videos: &'a mut VideoDownloads, } /// NB: Purposefully infallible. @@ -21,13 +22,14 @@ pub fn post_process_add_watermark( info!("Adding watermark..."); let output_video_fs_path_watermark = args.videos + .input_video .comfy_output_video_path .with_extension("_watermark.mp4"); let command_exit_status = args.comfy_deps .ffmpeg_watermark_command .execute(WatermarkArgs { - video_path: args.videos.video_to_watermark(), + video_path: args.videos.input_video.video_to_watermark(), maybe_override_logo_path: None, alpha: 0.6, scale: 0.1, // NB: 0.1 is good for the Storyteller logo @ 2653x512 placed on 1024x576 output. @@ -49,6 +51,6 @@ pub fn post_process_add_watermark( if use_watermarked_file { info!("Success generating watermarked file: {:?}", output_video_fs_path_watermark); - args.videos.watermarked_video_path = Some(output_video_fs_path_watermark); + args.videos.input_video.watermarked_video_path = Some(output_video_fs_path_watermark); } } diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/steps/post_process_restore_audio.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/steps/post_process_restore_audio.rs index 2f3ddfff1f..88b2422ec4 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/steps/post_process_restore_audio.rs +++ b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/steps/post_process_restore_audio.rs @@ -12,7 +12,7 @@ use crate::util::common_commands::ffmpeg_audio_replace_args::FfmpegAudioReplaceA pub struct PostProcessRestoreVideoArgs<'a> { pub comfy_deps: &'a ComfyDependencies, - pub videos: &'a mut VideoPaths, + pub videos: &'a mut VideoDownloads, } /// NB: Purposefully infallible. @@ -21,7 +21,18 @@ pub fn post_process_restore_audio( ) -> () { info!("Restoring audio..."); + // Use the original downloaded video if we didn't trim and resample it. + let input_video_file = args + .videos + .input_video + .maybe_processed_path + .as_deref() + .unwrap_or( + &args.videos.input_video.original_download_path + ); + let output_video_fs_path_restored = args.videos + .input_video .comfy_output_video_path .with_extension("_restored.mp4"); @@ -30,8 +41,8 @@ pub fn post_process_restore_audio( .ffmpeg_command_runner .run_with_subprocess(RunAsSubprocessArgs { args: Box::new(&FfmpegAudioReplaceArgs { - input_video_file: &args.videos.comfy_output_video_path, - input_audio_file: &args.videos.trimmed_resampled_video_path, + input_video_file: &args.videos.input_video.comfy_output_video_path, + input_audio_file: &input_video_file, output_video_file: &output_video_fs_path_restored, }), stderr: StreamRedirection::None, @@ -53,6 +64,6 @@ pub fn post_process_restore_audio( if use_restored_audio { info!("Success generating restored audio file: {:?}", output_video_fs_path_restored); - args.videos.audio_restored_video_path = Some(output_video_fs_path_restored); + args.videos.input_video.audio_restored_video_path = Some(output_video_fs_path_restored); } } diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/steps/preprocess_save_audio.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/steps/preprocess_save_audio.rs index bf38c23493..1526f5d91c 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/steps/preprocess_save_audio.rs +++ b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/steps/preprocess_save_audio.rs @@ -7,13 +7,14 @@ use subprocess_common::command_runner::command_runner_args::{RunAsSubprocessArgs 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::util::video_pathing::VideoDownloads; use crate::job::job_types::workflow::comfy_ui::video_style_transfer::util::video_pathing_deprecated::VideoPaths; use crate::util::common_commands::ffmpeg_audio_replace_args::FfmpegAudioReplaceArgs; use crate::util::common_commands::ffmpeg_extract_audio_args::FfmpegExtractAudioArgs; pub struct ProcessSaveAudioArgs<'a> { pub comfy_deps: &'a ComfyDependencies, - pub videos: &'a mut VideoPaths, + pub videos: &'a mut VideoDownloads, } pub fn preprocess_save_audio( @@ -21,14 +22,23 @@ pub fn preprocess_save_audio( ) -> Result<(), ProcessSingleJobError> { info!("Extracting audio..."); + // Use the original downloaded video if we didn't trim and resample it. + let input_video_file = args + .videos + .input_video + .maybe_processed_path + .as_deref() + .unwrap_or( + &args.videos.input_video.original_download_path + ); let command_exit_status = args .comfy_deps .ffmpeg_command_runner .run_with_subprocess(RunAsSubprocessArgs { args: Box::new(&FfmpegExtractAudioArgs { - input_video_file: &args.videos.trimmed_resampled_video_path, - output_file: &args.videos.trimmed_audio_path + input_video_file, + output_file: &args.videos.input_video.trimmed_audio_path }), stderr: StreamRedirection::None, stdout: StreamRedirection::None, 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 index 44ae4f0806..e7c9b61938 100644 --- 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 @@ -13,21 +13,15 @@ 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::util::comfy_dirs::ComfyDirs; use crate::job::job_types::workflow::comfy_ui::video_style_transfer::util::video_pathing::{SecondaryInputVideoAndPaths, VideoDownloads}; use crate::job::job_types::workflow::comfy_ui::video_style_transfer::util::video_pathing_deprecated::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 VideoDownloads, + pub comfy_dirs: &'a ComfyDirs, + pub videos: &'a mut VideoDownloads, } pub fn preprocess_trim_and_resample_videos( @@ -57,16 +51,16 @@ pub fn preprocess_trim_and_resample_videos( preprocess_trim_and_resample_primary_video( &args.comfy_deps, + &args.comfy_dirs, &resample_details, skip_resampling_video, - args.download_videos, - args.primary_video_paths)?; + args.videos)?; if !skip_resampling_video { preprocess_trim_and_resample_secondary_videos( &args.comfy_deps, &resample_details, - args.download_videos)?; + args.videos)?; } Ok(()) @@ -80,25 +74,25 @@ struct ResampleDetails { fn preprocess_trim_and_resample_primary_video( comfy_deps: &ComfyDependencies, + comfy_dirs: &ComfyDirs, resample_details: &ResampleDetails, skip_process_video: bool, - download_videos: &mut VideoDownloads, - primary_video_paths: &VideoPaths, + videos: &mut VideoDownloads, ) -> Result<(), ProcessSingleJobError> { - let resampled_path = primary_video_paths.comfy_input_dir.join("trimmed.mp4"); + let resampled_path = comfy_dirs.comfy_input_dir.join("trimmed.mp4"); 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(&download_videos.input_video.original_download_path, &resampled_path) + std::fs::copy(&videos.input_video.original_download_path, &resampled_path) .map_err(|err| { error!("Error copying video (1): {:?}", err); ProcessSingleJobError::IoError(err) })?; - std::fs::copy(&download_videos.input_video.original_download_path, &primary_video_paths.comfy_output_video_path) + std::fs::copy(&videos.input_video.original_download_path, &videos.input_video.comfy_output_video_path) .map_err(|err| { error!("Error copying video (2): {:?}", err); ProcessSingleJobError::IoError(err) @@ -110,7 +104,7 @@ fn preprocess_trim_and_resample_primary_video( // 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. - let comfy_input_video_path = primary_video_paths.comfy_input_dir.join("input.mp4"); + let comfy_input_video_path = comfy_dirs.comfy_input_dir.join("input.mp4"); // shell out to python script let output = Command::new("python3") @@ -122,7 +116,7 @@ fn preprocess_trim_and_resample_primary_video( .arg(format!("{:?}", resample_details.trim_end_millis)) .arg(format!("{:?}", resample_details.target_fps)) .arg("--input") - .arg(path_to_string(&download_videos.input_video.original_download_path)) + .arg(path_to_string(&videos.input_video.original_download_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. @@ -153,14 +147,14 @@ fn preprocess_trim_and_resample_primary_video( })?; } - primary_video_paths.debug_print_paths_after_trim(); + //primary_video_paths.debug_print_paths_after_trim(); if let Ok(Some(dimensions)) = ffprobe_get_dimensions(&resampled_path) { info!("Trimmed / resampled video dimensions: {}x{}", dimensions.width, dimensions.height); } // NB(bt,2024-07-10): Even if we don't resample, the python side still expects certain pathing for now - download_videos.input_video.maybe_processed_path = Some(resampled_path); + videos.input_video.maybe_processed_path = Some(resampled_path); Ok(()) } 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 d42b151df6..8408505c27 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,8 +55,7 @@ pub struct SaveResultsArgs<'a> { pub comfy_deps: &'a ComfyDependencies, pub job_progress_reporter: &'a mut Box, - pub download_videos: &'a VideoDownloads, - pub videos: &'a VideoPaths, + pub videos: &'a VideoDownloads, pub inference_duration: Duration, pub comfy_args: &'a WorkflowArgs, @@ -71,7 +70,7 @@ pub async fn validate_and_save_results(args: SaveResultsArgs<'_>) -> Result) -> Result) -> Result) -> Result Self { + // TODO(bt,2024-04-21): The pathing for this job is complicated. ComfyUI + // was set up with some of these expectations, which is bad. Worse, the + // jobs enqueue the expected output path -- no idea why that was done. + // This should all be fixed. + let root_comfy_path = comfy_deps.inference_command.mounts_directory.as_path(); + Self { + comfy_input_dir: root_comfy_path.join("input"), + comfy_output_dir: root_comfy_path.join("output"), + } + } +} diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/util/mod.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/util/mod.rs index 0bc2d36338..54f15776b5 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/util/mod.rs +++ b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/util/mod.rs @@ -1,3 +1,4 @@ +pub mod comfy_dirs; pub mod comfy_ui_inference_command; pub mod video_pathing; pub mod video_pathing_deprecated; diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/util/video_pathing.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/util/video_pathing.rs index 81f602f022..da72dc2629 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/util/video_pathing.rs +++ b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/util/video_pathing.rs @@ -7,6 +7,7 @@ use mysql_queries::queries::media_files::get::get_media_file::MediaFile; use tokens::tokens::media_files::MediaFileToken; use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; +use crate::job::job_types::workflow::comfy_ui::video_style_transfer::util::comfy_dirs::ComfyDirs; /// Keep track of where we download videos to and where they end up after processing. pub struct VideoDownloads { @@ -22,17 +23,17 @@ pub struct VideoDownloads { /// This is the primary input video for style transfer. pub struct PrimaryInputVideoAndPaths { /// Media file record. - pub record: VideoMediaFileRecord, + pub record: MediaFile, /// Filesystem path of the downloaded original video /// This is the input. pub original_download_path: PathBuf, + // TODO: Rename back to `trimmed_resampled...` /// Filesystem path of the trimmed and resampled video /// This is the first output we generate. /// We'll use this downstream once it's available. - //pub trimmed_resampled_video_path: PathBuf, - pub maybe_processed_path: Option, + pub maybe_processed_path: Option, //pub trimmed_resampled_video_path: PathBuf, /// Filesystem path of the trimmed audio pub trimmed_audio_path: PathBuf, @@ -66,7 +67,7 @@ pub struct PrimaryInputVideoAndPaths { /// We attach metadata as these videos progress through the system (if they're present). pub struct SecondaryInputVideoAndPaths { /// Media file record. - pub record: VideoMediaFileRecord, + pub record: MediaFilesByTokensRecord, /// Filesystem path of the downloaded original video pub original_download_path: PathBuf, @@ -76,12 +77,12 @@ pub struct SecondaryInputVideoAndPaths { } // Adapter enum -pub enum VideoMediaFileRecord { - // Records returned by single lookup - Single(MediaFile), - // Records returned by batch query - Bulk(MediaFilesByTokensRecord), -} +// pub enum VideoMediaFileRecord { +// // Records returned by single lookup +// Single(MediaFile), +// // Records returned by batch query +// Bulk(MediaFilesByTokensRecord), +// } pub trait CommonVideoPathing { fn original_video_path(&self) -> &Path; @@ -100,22 +101,15 @@ impl CommonVideoPathing for SecondaryInputVideoAndPaths { } impl PrimaryInputVideoAndPaths { - pub fn new(media_file_record: MediaFile, root_comfy_path: &Path, job_output_path: &str) -> Self { - // TODO(bt,2024-04-21): The pathing for this job is complicated. ComfyUI - // was set up with some of these expectations, which is bad. Worse, the - // jobs enqueue the expected output path -- no idea why that was done. - // This should all be fixed. - let input_dir = root_comfy_path.join("input"); - let output_dir = root_comfy_path.join("output"); - - let original_video_path = input_dir.join("video.mp4"); - let trimmed_resampled_video_path = input_dir.join("trimmed.mp4"); - let trimmed_audio_path = input_dir.join("trimmed.wav"); - let comfy_input_video_path = input_dir.join("input.mp4"); - let comfy_output_video_path = output_dir.join(job_output_path); // TODO: This sucks. + pub fn new(media_file_record: MediaFile, comfy_dirs: &ComfyDirs, job_output_path: &str) -> Self { + let original_video_path = comfy_dirs.comfy_input_dir.join("video.mp4"); + let trimmed_resampled_video_path = comfy_dirs.comfy_input_dir.join("trimmed.mp4"); + let trimmed_audio_path = comfy_dirs.comfy_input_dir.join("trimmed.wav"); + let comfy_input_video_path = comfy_dirs.comfy_input_dir.join("input.mp4"); + let comfy_output_video_path = comfy_dirs.comfy_output_dir.join(job_output_path); // TODO: This sucks. Self { - record: VideoMediaFileRecord::Single(media_file_record), + record: media_file_record, original_download_path: original_video_path, maybe_processed_path: Some(trimmed_resampled_video_path), trimmed_audio_path, @@ -123,42 +117,61 @@ impl PrimaryInputVideoAndPaths { comfy_output_video_path, audio_restored_video_path: None, watermarked_video_path: None, - comfy_input_dir: input_dir, + comfy_input_dir: comfy_dirs.comfy_input_dir.clone(), } } + + pub fn video_to_watermark(&self) -> &PathBuf { + // Try to use the audio-restored video if it's available + self.audio_restored_video_path.as_ref() + .unwrap_or(&self.comfy_output_video_path) + } + + pub fn get_final_video_to_upload(&self) -> &PathBuf { + // This is the video to upload as the result and save in the media_files table. + self.watermarked_video_path.as_ref() + .or(self.audio_restored_video_path.as_ref()) + .unwrap_or(&self.comfy_output_video_path) + } + + pub fn get_non_watermarked_video_to_upload(&self) -> &PathBuf { + // We'll upload this for internal use and for premium users. + // Same as "video_to_watermark()" + self.audio_restored_video_path.as_ref() + .unwrap_or(&self.comfy_output_video_path) + } } - -impl VideoMediaFileRecord { - pub fn token(&self) -> &MediaFileToken { - match self { - VideoMediaFileRecord::Single(m) => &m.token, - VideoMediaFileRecord::Bulk(m) => &m.token, - } - } - - pub fn maybe_title(&self) -> Option<&str> { - match self { - VideoMediaFileRecord::Single(m) => m.maybe_title.as_deref(), - VideoMediaFileRecord::Bulk(m) => m.maybe_title.as_deref(), - } - } - - pub fn maybe_style_transfer_source_media_file_token(&self) -> Result, ProcessSingleJobError> { - // TODO(bt,2024-07-09): Future proofing this to deliberately explode in case I query the - // primary input media files with the bulk query. - match self { - VideoMediaFileRecord::Single(m) => Ok(m.maybe_style_transfer_source_media_file_token.as_ref()), - VideoMediaFileRecord::Bulk(_m) => Err(ProcessSingleJobError::Other(anyhow!("bad refactor?: failed to query foreign key"))), - } - } - - pub fn maybe_scene_source_media_file_token(&self) -> Result, ProcessSingleJobError> { - // TODO(bt,2024-07-09): Future proofing this to deliberately explode in case I query the - // primary input media files with the bulk query. - match self { - VideoMediaFileRecord::Single(m) => Ok(m.maybe_scene_source_media_file_token.as_ref()), - VideoMediaFileRecord::Bulk(_m) => Err(ProcessSingleJobError::Other(anyhow!("bad refactor?: failed to query foreign key"))), - } - } -} +// impl VideoMediaFileRecord { +// pub fn token(&self) -> &MediaFileToken { +// match self { +// VideoMediaFileRecord::Single(m) => &m.token, +// VideoMediaFileRecord::Bulk(m) => &m.token, +// } +// } +// +// pub fn maybe_title(&self) -> Option<&str> { +// match self { +// VideoMediaFileRecord::Single(m) => m.maybe_title.as_deref(), +// VideoMediaFileRecord::Bulk(m) => m.maybe_title.as_deref(), +// } +// } +// +// pub fn maybe_style_transfer_source_media_file_token(&self) -> Result, ProcessSingleJobError> { +// // TODO(bt,2024-07-09): Future proofing this to deliberately explode in case I query the +// // primary input media files with the bulk query. +// match self { +// VideoMediaFileRecord::Single(m) => Ok(m.maybe_style_transfer_source_media_file_token.as_ref()), +// VideoMediaFileRecord::Bulk(_m) => Err(ProcessSingleJobError::Other(anyhow!("bad refactor?: failed to query foreign key"))), +// } +// } +// +// pub fn maybe_scene_source_media_file_token(&self) -> Result, ProcessSingleJobError> { +// // TODO(bt,2024-07-09): Future proofing this to deliberately explode in case I query the +// // primary input media files with the bulk query. +// match self { +// VideoMediaFileRecord::Single(m) => Ok(m.maybe_scene_source_media_file_token.as_ref()), +// VideoMediaFileRecord::Bulk(_m) => Err(ProcessSingleJobError::Other(anyhow!("bad refactor?: failed to query foreign key"))), +// } +// } +// }