From f844aa7016128ebcccff2a129b2b74105969b164 Mon Sep 17 00:00:00 2001
From: ArEnSc <6252325+ArEnSc@users.noreply.github.com>
Date: Thu, 16 May 2024 06:30:33 -0700
Subject: [PATCH] Added strength parameters (#322)

Co-authored-by: Storyteller <storyteller@MacBook-Pro-de-Storyteller.local>
---
 .../workflow_payload.rs                       |  4 ++++
 .../comfy_ui/comfy_ui_inference_command.rs    |  7 ++++++
 .../workflow/comfy_ui/process_job.rs          |  2 ++
 .../workflows/enqueue_comfy_ui_handler.rs     |  1 +
 .../enqueue_video_style_transfer_handler.rs   | 23 +++++++++++++++++++
 .../enqueue_workflow_upload_request.rs        |  1 +
 6 files changed, 38 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 1e5c6f7c3f..1e4ccdcd34 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
@@ -142,4 +142,8 @@ pub struct WorkflowArgs {
     #[serde(rename = "up")]
     #[serde(skip_serializing_if = "Option::is_none")]
     pub use_upscaler: Option<bool>,
+
+    #[serde(rename = "s")]
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub strength: Option<f32>,
 }
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 f9ddce6d11..6df1560d91 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
@@ -86,6 +86,7 @@ pub struct InferenceArgs<'s> {
 
     pub face_detailer_enabled: bool,
     pub upscaler_enabled: bool,
+    pub maybe_strength: Option<f32>,
 }
 
 pub enum InferenceDetails<'s> {
@@ -262,6 +263,12 @@ impl ComfyInferenceCommand {
             command.push_str(" --upscaler-enabled ");
         }
 
+        if let Some(strength) = args.maybe_strength {
+            command.push_str(" --strength ");
+            command.push_str(&strength.to_string());
+            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 6a73de8cde..083238e3c2 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
@@ -399,8 +399,10 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
             inference_details,
             face_detailer_enabled: comfy_args.use_face_detailer.unwrap_or(false),
             upscaler_enabled: comfy_args.use_upscaler.unwrap_or(false),
+            maybe_strength: comfy_args.strength,
         });
 
+
     let inference_duration = Instant::now().duration_since(inference_start_time);
 
     info!("Inference command exited with status: {:?}", command_exit_status);
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 6415221d6c..1c11d0a258 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
@@ -281,6 +281,7 @@ pub async fn enqueue_comfy_ui_handler(
             .map(|value| str_to_bool(&value)),
         use_face_detailer: None,
         use_upscaler: None,
+        strength: 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 693ac30c9c..fd3f43e6b1 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
@@ -48,6 +48,15 @@ const DEFAULT_TRIM_MILLISECONDS_END : u64 = 3_000;
 /// This is the maximum duration (for premium users)
 const MAXIMUM_DURATION_MILLIS : u64 = 10_000;
 
+/// Strength of Diffusion:
+///  * Range (0.0 - 1.0)
+///  * 0.0 less dreamy
+///  * 1.0 dream more
+///  * The Python code will default to "1.0" if not supplied
+const MINIMUM_STRENGTH : f32 = 0.0;
+const MAXIMUM_STRENGTH : f32 = 1.0;
+const DEFAULT_STRENGTH : f32 = 1.0;
+
 #[derive(Deserialize, ToSchema)]
 pub struct EnqueueVideoStyleTransferRequest {
     /// Entropy for request de-duplication (required)
@@ -86,6 +95,9 @@ pub struct EnqueueVideoStyleTransferRequest {
     /// Only for premium accounts
     use_upscaler: Option<bool>,
 
+    /// Use Strength of the style transfer
+    use_strength: Option<f32>,
+
     /// Optional visibility setting override.
     creator_set_visibility: Option<Visibility>,
 
@@ -275,6 +287,16 @@ pub async fn enqueue_video_style_transfer_handler(
         }
     }
 
+    let maybe_strength = request.use_strength
+        .map(|strength| {
+            if strength < MINIMUM_STRENGTH || strength > MAXIMUM_STRENGTH {
+                Err(EnqueueVideoStyleTransferError::BadInput("Strength must be between 0.0 and 1.0".to_string()))
+            } else {
+                Ok(strength)
+            }
+        })
+        .transpose()?;
+
     // Must have paid plan to remove watermark
     let remove_watermark = request.remove_watermark
         .map(|remove_watermark| remove_watermark && has_paid_plan);
@@ -305,6 +327,7 @@ pub async fn enqueue_video_style_transfer_handler(
             .map(|value| str_to_bool(&value)),
         use_face_detailer: request.use_face_detailer,
         use_upscaler: request.use_upscaler,
+        strength: maybe_strength,
     };
 
     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 28deea42b5..5a27a5b294 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
@@ -250,6 +250,7 @@ pub async fn enqueue_workflow_upload_request(
             .map(|value| str_to_bool(&value)),
         use_face_detailer: None,
         use_upscaler: None,
+        strength: None,
     };
    
     // create the inference args here