From a1fc56eceee966860dd9a02bffcc43f05ab9fb33 Mon Sep 17 00:00:00 2001
From: Brandon Thomas <echelon@gmail.com>
Date: Sun, 4 Feb 2024 01:48:05 -0500
Subject: [PATCH] prompt argument for comfy

---
 .../workflow/comfy_ui/comfy_ui_inference_command.rs          | 5 +++++
 .../src/job/job_types/workflow/comfy_ui/process_job.rs       | 4 ++++
 2 files changed, 9 insertions(+)

diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_inference_command.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_inference_command.rs
index e9e1f0ada5..707932c08d 100644
--- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_inference_command.rs
+++ b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_inference_command.rs
@@ -67,6 +67,7 @@ pub enum ExecutableOrCommand {
 pub struct InferenceArgs<'s> {
     pub stderr_output_file: &'s Path,
     pub stdout_output_file: &'s Path,
+    pub prompt_location: &'s Path,
 }
 
 impl ComfyInferenceCommand {
@@ -156,6 +157,10 @@ impl ComfyInferenceCommand {
             }
         }
 
+        command.push_str(" --prompt ");
+        command.push_str(&path_to_string(args.prompt_location));
+        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/process_job.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs
index 17577cabbe..44a108e91e 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
@@ -324,16 +324,20 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
     info!("Running ComfyUI inference...");
 
     // ==================== RUN INFERENCE SCRIPT ==================== //
+
     let stderr_output_file = work_temp_dir.path().join("stderr.txt");
     let stdout_output_file = work_temp_dir.path().join("stdout.txt");
 
     let inference_start_time = Instant::now();
 
+    let prompt_path = PathBuf::from(&workflow_path);
+
     let command_exit_status = model_dependencies
         .inference_command
         .execute_inference(InferenceArgs {
             stderr_output_file: &stderr_output_file,
             stdout_output_file: &stdout_output_file,
+            prompt_location: &prompt_path,
         });
 
     let inference_duration = Instant::now().duration_since(inference_start_time);