From d5f07376ef1133aff2277abc28d108bcb0fdcc16 Mon Sep 17 00:00:00 2001
From: Brandon Thomas <echelon@gmail.com>
Date: Mon, 13 May 2024 16:24:34 -0700
Subject: [PATCH] face detailer and upscaler args

---
 .../generic_inference_args/workflow_payload.rs        |  8 ++++++++
 .../workflow/comfy_ui/comfy_ui_inference_command.rs   | 11 +++++++++++
 .../job/job_types/workflow/comfy_ui/process_job.rs    |  2 ++
 .../endpoints/workflows/enqueue_comfy_ui_handler.rs   |  2 ++
 .../workflows/enqueue_video_style_transfer_handler.rs | 10 ++++++++++
 .../workflows/enqueue_workflow_upload_request.rs      |  2 ++
 6 files changed, 35 insertions(+)

diff --git a/crates/lib/mysql_queries/src/payloads/generic_inference_args/workflow_payload.rs b/crates/lib/mysql_queries/src/payloads/generic_inference_args/workflow_payload.rs
index 1231952b7b..1e5c6f7c3f 100644
--- a/crates/lib/mysql_queries/src/payloads/generic_inference_args/workflow_payload.rs
+++ b/crates/lib/mysql_queries/src/payloads/generic_inference_args/workflow_payload.rs
@@ -134,4 +134,12 @@ pub struct WorkflowArgs {
     #[serde(rename = "pa")]
     #[serde(skip_serializing_if = "Option::is_none")]
     pub rollout_python_workflow_args: Option<bool>,
+
+    #[serde(rename = "fd")]
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub use_face_detailer: Option<bool>,
+
+    #[serde(rename = "up")]
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub use_upscaler: Option<bool>,
 }
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 213c29dc80..b62110d84a 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
@@ -99,6 +99,9 @@ pub struct InferenceArgs<'s> {
     /// If set, Python will be in charge of overwriting the prompt JSON file
     /// with the correct workflow args.
     pub maybe_style: Option<StyleTransferName>,
+
+    pub face_detailer_enabled: bool,
+    pub upscaler_enabled: bool,
 }
 
 impl ComfyInferenceCommand {
@@ -234,6 +237,14 @@ impl ComfyInferenceCommand {
             command.push_str(" ");
         }
 
+        if args.face_detailer_enabled {
+            command.push_str(" --face-detailer-enabled=true ");
+        }
+
+        if args.upscaler_enabled {
+            command.push_str(" --upscaler-enabled=true ");
+        }
+
         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 d7707ccdd5..f0cc4e77c8 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
@@ -409,6 +409,8 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
             maybe_positive_prompt_filename,
             maybe_negative_prompt_filename,
             maybe_style,
+            face_detailer_enabled: comfy_args.use_face_detailer.unwrap_or(false),
+            upscaler_enabled: comfy_args.use_upscaler.unwrap_or(false),
         });
 
     let inference_duration = Instant::now().duration_since(inference_start_time);
diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_comfy_ui_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_comfy_ui_handler.rs
index f0d835f35d..6415221d6c 100644
--- a/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_comfy_ui_handler.rs
+++ b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_comfy_ui_handler.rs
@@ -279,6 +279,8 @@ pub async fn enqueue_comfy_ui_handler(
         enable_lipsync: None,
         rollout_python_workflow_args: get_request_header_optional(&http_request, "PYTHON-WORKFLOW-ARGS")
             .map(|value| str_to_bool(&value)),
+        use_face_detailer: None,
+        use_upscaler: None,
     };
 
     info!("Creating ComfyUI job record...");
diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_video_style_transfer_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_video_style_transfer_handler.rs
index f9d60813ed..12139e99bb 100644
--- a/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_video_style_transfer_handler.rs
+++ b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_video_style_transfer_handler.rs
@@ -78,6 +78,14 @@ pub struct EnqueueVideoStyleTransferRequest {
     /// Only for premium accounts
     remove_watermark: Option<bool>,
 
+    /// Use face detailer
+    /// Only for premium accounts
+    use_face_detailer: Option<bool>,
+
+    /// Use video upscaler
+    /// Only for premium accounts
+    use_upscaler: Option<bool>,
+
     /// Optional visibility setting override.
     creator_set_visibility: Option<Visibility>,
 }
@@ -290,6 +298,8 @@ pub async fn enqueue_video_style_transfer_handler(
         target_fps: None,
         rollout_python_workflow_args: get_request_header_optional(&http_request, "PYTHON-WORKFLOW-ARGS")
             .map(|value| str_to_bool(&value)),
+        use_face_detailer: request.use_face_detailer,
+        use_upscaler: request.use_upscaler,
     };
 
     info!("Creating ComfyUI job record...");
diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_workflow_upload_request.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_workflow_upload_request.rs
index e1aad7a188..28deea42b5 100644
--- a/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_workflow_upload_request.rs
+++ b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_workflow_upload_request.rs
@@ -248,6 +248,8 @@ pub async fn enqueue_workflow_upload_request(
         enable_lipsync: None,
         rollout_python_workflow_args: get_request_header_optional(&http_request, "PYTHON-WORKFLOW-ARGS")
             .map(|value| str_to_bool(&value)),
+        use_face_detailer: None,
+        use_upscaler: None,
     };
    
     // create the inference args here