diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/job_outputs.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/job_outputs.rs index f97a319946..a7ca785fb8 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/job_outputs.rs +++ b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/job_outputs.rs @@ -1,69 +1,53 @@ -use std::path::PathBuf; +use std::path::{Path, PathBuf}; -// TODO(bt,2024-04-21): This doesn't need to be a state machine. We can generate the paths we need upfront. -pub struct JobOutputsStageOne { - /// Filesystem path of the downloaded original video - pub original_video_path: PathBuf, -} - -pub struct JobOutputsStageTwo { - /// Filesystem path of the downloaded original video - pub original_video_path: PathBuf, - - /// Filesystem path of the trimmed and resampled video - /// We'll use this downstream once it's available. - pub trimmed_resampled_video_path: PathBuf, -} - -pub struct JobOutputsStageThree { +pub struct JobOutputs { /// Filesystem path of the downloaded original video + /// This is the input. pub original_video_path: PathBuf, /// 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, /// Filesystem path of the style transfer output + /// This is the main purpose of the job, and the second output we generate. pub comfy_output_video_path: PathBuf, /// Filesystem path of the audio-restored output + /// This is the third output we generate. /// We'll want to upload this as a result, if available. pub audio_restored_video_path: Option, /// Watermarked final result + /// This is the fourth output we generate. /// We'll want to upload this as a result, if available. pub watermarked_video_path: Option, } -impl JobOutputsStageOne { - pub fn new(original_video_path: PathBuf) -> Self { + +impl JobOutputs { + pub fn new(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("input.mp4"); + let comfy_output_video_path = output_dir.join(job_output_path); // TODO: This sucks. + Self { original_video_path, - } - } - - pub fn with_trimmed_resampled_video(self, trimmed_resampled_video_path: PathBuf) -> JobOutputsStageTwo { - JobOutputsStageTwo { - original_video_path: self.original_video_path, trimmed_resampled_video_path, - } - } -} - -impl JobOutputsStageTwo { - pub fn add_comfy_output(self, comfy_output_video_path: PathBuf) -> JobOutputsStageThree { - JobOutputsStageThree { - original_video_path: self.original_video_path, - trimmed_resampled_video_path: self.trimmed_resampled_video_path, comfy_output_video_path, audio_restored_video_path: None, watermarked_video_path: None, } } -} - -impl JobOutputsStageThree { 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() @@ -84,4 +68,3 @@ impl JobOutputsStageThree { .unwrap_or(&self.comfy_output_video_path) } } - diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs index cb196fbb59..98deff3672 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs @@ -39,7 +39,7 @@ use tokens::tokens::prompts::PromptToken; 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::comfy_ui_inference_command::InferenceArgs; -use crate::job::job_types::workflow::comfy_ui::job_outputs::JobOutputsStageOne; +use crate::job::job_types::workflow::comfy_ui::job_outputs::JobOutputs; use crate::job::job_types::workflow::comfy_ui::validate_job::validate_job; use crate::job_dependencies::JobDependencies; use crate::util::common_commands::ffmpeg_audio_replace_args::FfmpegAudioReplaceArgs; @@ -265,7 +265,6 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result) -> Result = None; - //let mut maybe_original_input_path: Option = None; + let mut videos = JobOutputs::new(&root_comfy_path, job_args.output_path); // TODO(bt,2024-04-20): Clean up this mess. @@ -301,9 +296,6 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result) -> Result) -> Result {} } - let videos = videos.with_trimmed_resampled_video(input_dir.join("input.mp4")); - // make outputs dir if not exist let output_dir = root_comfy_path.join("output"); if !output_dir.exists() { @@ -408,14 +394,6 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result