cleanup job code

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<PathBuf>,
 
   /// 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<PathBuf>,
 }
 
-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)
   }
 }
-

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<JobSuccessResu
     }
 
     let input_dir = root_comfy_path.join("input");
-    let original_video_input_path = input_dir.join("video.mp4");
 
     if !input_dir.exists() {
         std::fs::create_dir_all(&input_dir)
@@ -273,11 +272,7 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
     }
 
     // Keep track of all the video files we generate
-    let videos = JobOutputsStageOne::new(original_video_input_path);
-
-    // Download Input file if specified
-    //let mut maybe_input_path: Option<PathBuf> = None;
-    //let mut maybe_original_input_path: Option<PathBuf> = 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<JobSuccessResu
 
             info!("Input media file cloud bucket path: {:?}", media_file_bucket_path.get_full_object_path_str());
 
-            //let input_filename = "video.mp4".to_string();
-            //let input_path = path_to_string(input_dir.join(input_filename));
-
             info!("Downloading input file to {:?}", videos.original_video_path);
 
             remote_cloud_file_client.download_media_file(
@@ -313,8 +305,6 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
 
             info!("Downloaded!");
 
-            //maybe_original_input_path = Some(input_path.parse().unwrap());
-
             // ========================= PROCESS VIDEO ======================== //
 
             let target_fps = comfy_args.target_fps.unwrap_or(24);
@@ -356,14 +346,10 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
                 error!("stderr: {}", String::from_utf8_lossy(&output.stderr));
                 return Err(ProcessSingleJobError::Other(anyhow!("Command failed: {:?}", output.status)));
             }
-
-            //maybe_input_path = Some(input_dir.join("input.mp4"));
         }
         None => {}
     }
 
-    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<JobSuccessResu
 
     // ==================== CHECK OUTPUT FILE ======================== //
 
-    // TODO(bt,2024-04-21): We should be responsible for these paths, not relying on the downstream comfy server
-    let mut videos = {
-        let mut comfy_output_video_file = root_comfy_path.join("output");
-        comfy_output_video_file = comfy_output_video_file.join(job_args.output_path);
-
-        videos.add_comfy_output(comfy_output_video_file)
-    };
-
     if let Err(err) = check_file_exists(&videos.comfy_output_video_path) {
         error!("Output file does not  exist: {:?}", err);