From 7e1043b824dadb95b660e1fec40cb2b95bea5a65 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Sat, 6 Apr 2024 00:07:47 -0400 Subject: [PATCH] remove_watermark option on endpoint --- .../generic_inference_args/workflow_payload.rs | 4 ++++ .../workflows/enqueue_comfy_ui_handler.rs | 1 + .../enqueue_video_style_transfer_handler.rs | 15 +++++++++++++-- .../workflows/enqueue_workflow_upload_request.rs | 1 + 4 files changed, 19 insertions(+), 2 deletions(-) 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 ade78b212a..dca6bff67b 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 @@ -125,4 +125,8 @@ pub struct WorkflowArgs { #[serde(rename = "el")] #[serde(skip_serializing_if = "Option::is_none")] pub enable_lipsync: Option, + + #[serde(rename = "rm")] + #[serde(skip_serializing_if = "Option::is_none")] + pub remove_watermark: Option, } 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 ebabb0a4f3..4b3b92b293 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 @@ -262,6 +262,7 @@ pub async fn enqueue_comfy_ui_handler( trim_start_seconds: Some(trim_start_seconds), trim_end_seconds: Some(trim_end_seconds), target_fps:Some(target_fps), + remove_watermark: None, // NB: The following are unused: maybe_google_drive_link: None, maybe_title: None, 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 c11f46741d..0e91faecec 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 @@ -67,6 +67,10 @@ pub struct EnqueueVideoStyleTransferRequest { /// Enable lipsyncing in the workflow enable_lipsync: Option, + /// Remove watermark from the output + /// Only for premium accounts + remove_watermark: Option, + /// Optional visibility setting override. creator_set_visibility: Option, } @@ -236,9 +240,11 @@ pub async fn enqueue_video_style_transfer_handler( let mut trim_start_millis = request.trim_start_millis.unwrap_or(0); let mut trim_end_millis = request.trim_end_millis.unwrap_or(3_000); + let has_paid_plan = plan.plan_slug() == "fakeyou_contributor" || plan.plan_category() == PlanCategory::Paid; + // block trim too much - if plan.plan_slug() == "fakeyou_contributor" || plan.plan_category() == PlanCategory::Paid { - if trim_end_millis - trim_start_millis > 20_000 { + if has_paid_plan { + if trim_end_millis - trim_start_millis > 10_000 { trim_start_millis = 0; trim_end_millis = 3_000; } @@ -249,6 +255,10 @@ pub async fn enqueue_video_style_transfer_handler( } } + // Must have paid plan to remove watermark + let remove_watermark = request.remove_watermark + .map(|remove_watermark| remove_watermark && has_paid_plan); + let inference_args = WorkflowArgs { style_name: Some(request.style), creator_visibility: Some(set_visibility), @@ -258,6 +268,7 @@ pub async fn enqueue_video_style_transfer_handler( negative_prompt: request.negative_prompt.clone(), enable_lipsync: request.enable_lipsync, maybe_input_file: Some(request.input_file.clone()), + remove_watermark, // The new, simplified enqueuing doesn't care about the following parameters: maybe_lora_model: None, maybe_json_modifications: None, 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 301497b571..c18a5e5a1b 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 @@ -229,6 +229,7 @@ pub async fn enqueue_workflow_upload_request( maybe_description: Some(description), maybe_commit_hash: Some(commit_hash), creator_visibility: Some(visibility), + remove_watermark: None, // NB: The following args are irrelevant for uploading workflows maybe_lora_model: None, maybe_json_modifications: None,