From 4ca8932591b79c6a183e0589f012e679796bc62d Mon Sep 17 00:00:00 2001
From: Brandon Thomas <echelon@gmail.com>
Date: Sat, 6 Jul 2024 18:37:50 -0400
Subject: [PATCH] inference job: add prompt travel and skip frame

---
 .../comfy_ui_inference_command.rs             | 21 ++++++++++++++++++-
 .../video_style_transfer/process_job.rs       | 18 +++++++++++++---
 2 files changed, 35 insertions(+), 4 deletions(-)

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 ffdac632e0..ae997192f2 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
@@ -99,6 +99,8 @@ pub struct InferenceArgs<'s> {
 
     pub maybe_strength: Option<f32>,
 
+    pub frame_skip: Option<u8>,
+
     pub global_ipa_image_filename: Option<String>,
     pub global_ipa_strength: Option<f32>,
 }
@@ -121,6 +123,11 @@ pub enum InferenceDetails<'s> {
         /// with the correct workflow args.
         maybe_negative_prompt_filename: Option<&'s Path>,
 
+        /// Travel prompt file.
+        /// Optional: If set, Python will be in charge of overwriting the prompt JSON file
+        /// with the correct workflow args.
+        maybe_travel_prompt_filename: Option<&'s Path>,
+
         /// Style name
         /// Optional: If set, Python will be in charge of overwriting the prompt JSON file
         /// with the correct workflow args.
@@ -254,6 +261,7 @@ impl ComfyInferenceCommand {
             InferenceDetails::NewPythonArgs {
                 maybe_positive_prompt_filename,
                 maybe_negative_prompt_filename,
+                maybe_travel_prompt_filename,
                 maybe_style
             } => {
                 if let Some(positive_prompt_filename) = maybe_positive_prompt_filename {
@@ -268,6 +276,12 @@ impl ComfyInferenceCommand {
                     command.push_str(" ");
                 }
 
+                if let Some(travel_prompt_filename) = maybe_travel_prompt_filename {
+                    command.push_str(" --travel_prompt_filename ");
+                    command.push_str(&path_to_string(travel_prompt_filename));
+                    command.push_str(" ");
+                }
+
                 if let Some(style) = maybe_style {
                     command.push_str(" --style ");
                     command.push_str(style.to_str());
@@ -302,6 +316,12 @@ impl ComfyInferenceCommand {
             command.push_str(" ");
         }
 
+        if let Some(frame_skip) = args.frame_skip {
+            command.push_str(" --frame_skip ");
+            command.push_str(&frame_skip.to_string());
+            command.push_str(" ");
+        }
+
         if let Some(global_ipa_image_filename) = &args.global_ipa_image_filename {
             // NB: This is very dangerous. We're not handling escaping well.
             command.push_str(" --global_ipa_image_filename ");
@@ -315,7 +335,6 @@ impl ComfyInferenceCommand {
             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 6dd3f970e6..b10ada9c27 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
@@ -284,18 +284,17 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
 
     // ==================== RUN COMFY INFERENCE ==================== //
 
-    info!("Ready for ComfyUI inference...");
+    info!("Preparing for ComfyUI inference...");
 
     job_progress_reporter.log_status("running inference")
         .map_err(|e| ProcessSingleJobError::Other(e))?;
 
-    info!("Running ComfyUI inference...");
-
     let stderr_output_file = work_temp_dir.path().join("stderr.txt");
     let stdout_output_file = work_temp_dir.path().join("stdout.txt");
 
     let positive_prompt_file = work_temp_dir.path().join("positive_prompt.txt");
     let negative_prompt_file = work_temp_dir.path().join("negative_prompt.txt");
+    let travel_prompt_file = work_temp_dir.path().join("travel_prompt.txt");
 
     let maybe_positive_prompt_filename = comfy_args.positive_prompt
         .as_deref()
@@ -315,14 +314,26 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
         })
         .transpose()?;
 
+    let maybe_travel_prompt_filename = comfy_args.travel_prompt
+        .as_deref()
+        .map(|prompt| {
+            std::fs::write(&travel_prompt_file, prompt)
+                .map(|_| travel_prompt_file.as_path())
+                .map_err(|e| ProcessSingleJobError::IoError(e))
+        })
+        .transpose()?;
+
     let inference_details = InferenceDetails::NewPythonArgs {
         maybe_style: comfy_args.style_name,
         maybe_positive_prompt_filename,
         maybe_negative_prompt_filename,
+        maybe_travel_prompt_filename,
     };
 
     let inference_start_time = Instant::now();
 
+    info!("Running ComfyUI inference...");
+
     let command_exit_status = model_dependencies
         .inference_command
         .execute_inference(InferenceArgs {
@@ -335,6 +346,7 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
             disable_lcm: comfy_args.disable_lcm.unwrap_or(false),
             use_cinematic: comfy_args.use_cinematic.unwrap_or(false),
             maybe_strength: comfy_args.strength,
+            frame_skip: comfy_args.frame_skip,
             global_ipa_image_filename: global_ipa_image
                 .as_ref()
                 .map(|image| path_to_string(&image.ipa_image_path)),