From cef5d79bb1ea79c88f0102da9b56f7d424ceca9d Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Fri, 26 Jul 2024 22:42:58 -0400 Subject: [PATCH] watermark branding from hostname --- .../generic_inference_args/common/mod.rs | 1 + .../common/watermark_type.rs | 16 ++++++ .../live_portrait_payload.rs | 5 ++ .../payloads/generic_inference_args/mod.rs | 1 + .../workflow_payload.rs | 5 ++ .../enqueue_live_portrait_workflow_handler.rs | 18 +++++- .../enqueue_studio_workflow_handler.rs | 3 +- ...e_video_style_transfer_workflow_handler.rs | 18 +++++- .../workflows/enqueue_comfy_ui_handler.rs | 1 + .../enqueue_video_style_transfer_handler.rs | 1 + .../enqueue_workflow_upload_request.rs | 1 + .../storyteller_web/src/http_server/mod.rs | 1 + .../requests/get_request_domain_branding.rs | 55 +++++++++++++++++++ .../src/http_server/requests/mod.rs | 1 + 14 files changed, 124 insertions(+), 3 deletions(-) create mode 100644 crates/lib/mysql_queries/src/payloads/generic_inference_args/common/mod.rs create mode 100644 crates/lib/mysql_queries/src/payloads/generic_inference_args/common/watermark_type.rs create mode 100644 crates/service/web/storyteller_web/src/http_server/requests/get_request_domain_branding.rs create mode 100644 crates/service/web/storyteller_web/src/http_server/requests/mod.rs diff --git a/crates/lib/mysql_queries/src/payloads/generic_inference_args/common/mod.rs b/crates/lib/mysql_queries/src/payloads/generic_inference_args/common/mod.rs new file mode 100644 index 0000000000..7713b3d99c --- /dev/null +++ b/crates/lib/mysql_queries/src/payloads/generic_inference_args/common/mod.rs @@ -0,0 +1 @@ +pub mod watermark_type; \ No newline at end of file diff --git a/crates/lib/mysql_queries/src/payloads/generic_inference_args/common/watermark_type.rs b/crates/lib/mysql_queries/src/payloads/generic_inference_args/common/watermark_type.rs new file mode 100644 index 0000000000..ed1dd8a121 --- /dev/null +++ b/crates/lib/mysql_queries/src/payloads/generic_inference_args/common/watermark_type.rs @@ -0,0 +1,16 @@ + +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)] +pub enum WatermarkType { + /// Show a FakeYou watermark + #[serde(rename = "f")] // NB: DO NOT CHANGE. It could break live jobs. Renamed to be fewer bytes. + FakeYou, + + /// Show a Storyteller watermark + #[serde(rename = "s")] // NB: DO NOT CHANGE. It could break live jobs. Renamed to be fewer bytes. + Storyteller, + + /// Show no watermark + /// Same as if Option is None. + #[serde(rename = "n")] // NB: DO NOT CHANGE. It could break live jobs. Renamed to be fewer bytes. + NoWatermark, +} diff --git a/crates/lib/mysql_queries/src/payloads/generic_inference_args/live_portrait_payload.rs b/crates/lib/mysql_queries/src/payloads/generic_inference_args/live_portrait_payload.rs index 39b13a9c9c..40f19358d7 100644 --- a/crates/lib/mysql_queries/src/payloads/generic_inference_args/live_portrait_payload.rs +++ b/crates/lib/mysql_queries/src/payloads/generic_inference_args/live_portrait_payload.rs @@ -1,4 +1,5 @@ use tokens::tokens::media_files::MediaFileToken; +use crate::payloads::generic_inference_args::common::watermark_type::WatermarkType; // **DO NOT CHANGE THE NAMES OF FIELDS WITHOUT A MIGRATION STRATEGY** #[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq)] @@ -21,6 +22,10 @@ pub struct LivePortraitPayload { #[serde(skip_serializing_if = "Option::is_none")] pub remove_watermark: Option, + #[serde(rename = "w")] + #[serde(skip_serializing_if = "Option::is_none")] + pub watermark_type: Option, + /// This is a debugging flag. #[serde(rename = "sp")] #[serde(skip_serializing_if = "Option::is_none")] diff --git a/crates/lib/mysql_queries/src/payloads/generic_inference_args/mod.rs b/crates/lib/mysql_queries/src/payloads/generic_inference_args/mod.rs index 45478e0100..f46ecca860 100644 --- a/crates/lib/mysql_queries/src/payloads/generic_inference_args/mod.rs +++ b/crates/lib/mysql_queries/src/payloads/generic_inference_args/mod.rs @@ -1,3 +1,4 @@ +pub mod common; pub mod generic_inference_args; pub mod image_generation_payload; pub mod lipsync_payload; 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 1d04de5579..e4da0372ed 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 @@ -4,6 +4,7 @@ use enums::common::visibility::Visibility; use enums::no_table::style_transfer::style_transfer_name::StyleTransferName; use tokens::tokens::media_files::MediaFileToken; use tokens::tokens::model_weights::ModelWeightToken; +use crate::payloads::generic_inference_args::common::watermark_type::WatermarkType; #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)] #[serde(untagged)] @@ -169,6 +170,10 @@ pub struct WorkflowArgs { #[serde(skip_serializing_if = "Option::is_none")] pub remove_watermark: Option, + #[serde(rename = "wt")] + #[serde(skip_serializing_if = "Option::is_none")] + pub watermark_type: Option, + // TODO(bt,2024-05-13): This is a temporary rollout flag to enable us to do Python-side mapping of job args #[deprecated(note = "This has been rolled out and is effectively dead.")] #[serde(rename = "pa")] diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue/enqueue_live_portrait_workflow_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue/enqueue_live_portrait_workflow_handler.rs index 07d96b02ff..69a8b0798f 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue/enqueue_live_portrait_workflow_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue/enqueue_live_portrait_workflow_handler.rs @@ -18,6 +18,7 @@ use enums::by_table::generic_inference_jobs::inference_model_type::InferenceMode use enums::common::visibility::Visibility; use enums::no_table::style_transfer::style_transfer_name::StyleTransferName; use http_server_common::request::get_request_ip::get_request_ip; +use mysql_queries::payloads::generic_inference_args::common::watermark_type::WatermarkType; use mysql_queries::payloads::generic_inference_args::generic_inference_args::{GenericInferenceArgs, InferenceCategoryAbbreviated, PolymorphicInferenceArgs}; use mysql_queries::payloads::generic_inference_args::live_portrait_payload::{CropDimensions, LivePortraitPayload}; use mysql_queries::queries::generic_inference::web::insert_generic_inference_job::{insert_generic_inference_job, InsertGenericInferenceArgs}; @@ -31,6 +32,7 @@ use crate::configs::plans::get_correct_plan_for_session::get_correct_plan_for_se use crate::configs::plans::plan_category::PlanCategory; use crate::http_server::headers::get_routing_tag_header::get_routing_tag_header; use crate::http_server::headers::has_debug_header::has_debug_header; +use crate::http_server::requests::get_request_domain_branding::{DomainBranding, get_request_domain_branding}; use crate::http_server::validations::validate_idempotency_token_format::validate_idempotency_token_format; use crate::http_server::web_utils::require_user_session::RequireUserSessionError; use crate::http_server::web_utils::require_user_session_using_connection::require_user_session_using_connection; @@ -204,7 +206,7 @@ pub async fn enqueue_live_portrait_workflow_handler( EnqueueLivePortraitWorkflowError::BadInput("invalid idempotency token".to_string()) })?; - // ==================== LOOK UP MODEL INFO ==================== // + // ==================== HANDLE REQUEST ==================== // // TODO(bt): CHECK DATABASE FOR TOKENS! @@ -216,6 +218,7 @@ pub async fn enqueue_live_portrait_workflow_handler( let has_paid_plan = plan.plan_slug() == "fakeyou_contributor" || plan.plan_category() == PlanCategory::Paid; let _is_allowed_expensive_generation = is_staff || has_paid_plan; + let is_allowed_no_watermark = is_staff || has_paid_plan; let maybe_crop = request.maybe_crop.as_ref() .map(|crop| { @@ -227,10 +230,23 @@ pub async fn enqueue_live_portrait_workflow_handler( } }); + let branding = get_request_domain_branding(&http_request); + + let mut watermark_type = match branding { + Some(DomainBranding::FakeYou) => Some(WatermarkType::FakeYou), + Some(DomainBranding::Storyteller) => Some(WatermarkType::Storyteller), + None => Some(WatermarkType::Storyteller), + }; + + if request.remove_watermark.unwrap_or(false) && is_allowed_no_watermark { + watermark_type = None; + } + let payload = LivePortraitPayload { portrait_media_file_token: empty_media_file_token_to_null(Some(&request.source_media_file_token)), driver_media_file_token: empty_media_file_token_to_null(Some(&request.face_driver_media_file_token)), crop: maybe_crop, + watermark_type, remove_watermark: None, sleep_millis: None, }; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue/enqueue_studio_workflow_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue/enqueue_studio_workflow_handler.rs index 7d7f975121..133fdc55fc 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue/enqueue_studio_workflow_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue/enqueue_studio_workflow_handler.rs @@ -273,9 +273,10 @@ pub async fn enqueue_studio_workflow_handler( use_cinematic: coordinated_args.use_cinematic, use_face_detailer: coordinated_args.use_face_detailer, use_upscaler: coordinated_args.use_upscaler, - remove_watermark: coordinated_args.remove_watermark, lipsync_enabled: coordinated_args.use_lipsync, enable_lipsync: coordinated_args.use_lipsync, // TODO(bt): We can stop writing this flag after we re-deploy the job. + remove_watermark: coordinated_args.remove_watermark, + watermark_type: None, // TODO: Get rid of the temporary flags. rollout_python_workflow_args: None, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue/enqueue_video_style_transfer_workflow_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue/enqueue_video_style_transfer_workflow_handler.rs index 225c3ca2c1..18df2dcee7 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue/enqueue_video_style_transfer_workflow_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue/enqueue_video_style_transfer_workflow_handler.rs @@ -19,6 +19,7 @@ use enums::common::visibility::Visibility; use enums::no_table::style_transfer::style_transfer_name::StyleTransferName; use http_server_common::request::get_request_header_optional::get_request_header_optional; use http_server_common::request::get_request_ip::get_request_ip; +use mysql_queries::payloads::generic_inference_args::common::watermark_type::WatermarkType; use mysql_queries::payloads::generic_inference_args::generic_inference_args::{GenericInferenceArgs, InferenceCategoryAbbreviated, PolymorphicInferenceArgs}; use mysql_queries::payloads::generic_inference_args::workflow_payload::{WorkflowArgs, WorkflowType}; use mysql_queries::queries::generic_inference::web::insert_generic_inference_job::{insert_generic_inference_job, InsertGenericInferenceArgs}; @@ -38,6 +39,7 @@ use crate::http_server::endpoints::workflows::enqueue::vst_common::vst_request:: use crate::http_server::endpoints::workflows::enqueue::vst_common::vst_response::VstSuccessResponse; use crate::http_server::headers::get_routing_tag_header::get_routing_tag_header; use crate::http_server::headers::has_debug_header::has_debug_header; +use crate::http_server::requests::get_request_domain_branding::{DomainBranding, get_request_domain_branding}; use crate::http_server::session::lookup::user_session_extended::UserSessionExtended; use crate::http_server::validations::validate_idempotency_token_format::validate_idempotency_token_format; use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; @@ -241,9 +243,22 @@ pub async fn enqueue_video_style_transfer_workflow_handler( }; let is_allowed_expensive_generation = is_staff || has_paid_plan; + let is_allowed_no_watermark = is_staff || has_paid_plan; let coordinated_args = coordinate_workflow_args(coordinated_args, is_allowed_expensive_generation); + let branding = get_request_domain_branding(&http_request); + + let mut watermark_type = match branding { + Some(DomainBranding::FakeYou) => Some(WatermarkType::FakeYou), + Some(DomainBranding::Storyteller) => Some(WatermarkType::Storyteller), + None => Some(WatermarkType::Storyteller), + }; + + if request.remove_watermark.unwrap_or(false) && is_allowed_no_watermark { + watermark_type = None; + } + let inference_args = WorkflowArgs { // Type of workflow workflow_type: Some(WorkflowType::VideoStyleTransfer), @@ -273,9 +288,10 @@ pub async fn enqueue_video_style_transfer_workflow_handler( use_cinematic: coordinated_args.use_cinematic, use_face_detailer: coordinated_args.use_face_detailer, use_upscaler: coordinated_args.use_upscaler, - remove_watermark: coordinated_args.remove_watermark, lipsync_enabled: coordinated_args.use_lipsync, enable_lipsync: coordinated_args.use_lipsync, // TODO(bt): We can stop writing this flag after we re-deploy the job. + remove_watermark: coordinated_args.remove_watermark, + watermark_type, // TODO: Get rid of the temporary flags. rollout_python_workflow_args: None, 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 3d90ef2e2e..14061a8d7a 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 @@ -287,6 +287,7 @@ pub async fn enqueue_comfy_ui_handler( use_cinematic: None, strength: None, frame_skip: None, + watermark_type: 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 b4538c8cb2..4598fff45c 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 @@ -440,6 +440,7 @@ pub async fn enqueue_video_style_transfer_handler( trim_start_seconds: None, trim_end_seconds: None, target_fps: None, + watermark_type: None, }; 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 b567fe7f7b..d34276b280 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 @@ -257,6 +257,7 @@ pub async fn enqueue_workflow_upload_request( use_cinematic: None, strength: None, frame_skip: None, + watermark_type: None, }; // create the inference args here diff --git a/crates/service/web/storyteller_web/src/http_server/mod.rs b/crates/service/web/storyteller_web/src/http_server/mod.rs index a23a3a22be..612daf5d06 100644 --- a/crates/service/web/storyteller_web/src/http_server/mod.rs +++ b/crates/service/web/storyteller_web/src/http_server/mod.rs @@ -4,6 +4,7 @@ pub mod cookies; pub mod endpoints; pub mod headers; pub mod middleware; +pub mod requests; pub mod routes; pub mod session; pub mod validations; diff --git a/crates/service/web/storyteller_web/src/http_server/requests/get_request_domain_branding.rs b/crates/service/web/storyteller_web/src/http_server/requests/get_request_domain_branding.rs new file mode 100644 index 0000000000..c6f72ece52 --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/requests/get_request_domain_branding.rs @@ -0,0 +1,55 @@ +use actix_web::HttpRequest; + +#[derive(Debug, PartialEq, Eq, Clone, Copy)] +pub enum DomainBranding { + FakeYou, + Storyteller, +} + +pub fn get_request_domain_branding(http_request: &HttpRequest) -> Option { + http_request.uri() + .host() + .and_then(|host| match host { + host if host.contains("fakeyou") => Some(DomainBranding::FakeYou), + host if host.contains("storyteller") => Some(DomainBranding::Storyteller), + _ => None, + }) +} + +#[cfg(test)] +mod tests { + use actix_web::test::TestRequest; + use crate::http_server::requests::get_request_domain_branding::{DomainBranding, get_request_domain_branding}; + + #[test] + fn fakeyou_dot_com() { + let request = TestRequest::get() + .uri("https://fakeyou.com") + .to_http_request(); + assert_eq!(get_request_domain_branding(&request), Some(DomainBranding::FakeYou)); + } + + #[test] + fn api_dot_fakeyou_dot_com() { + let request = TestRequest::get() + .uri("https://api.fakeyou.com") + .to_http_request(); + assert_eq!(get_request_domain_branding(&request), Some(DomainBranding::FakeYou)); + } + + #[test] + fn storyteller_dot_ai() { + let request = TestRequest::get() + .uri("https://storyteller.ai") + .to_http_request(); + assert_eq!(get_request_domain_branding(&request), Some(DomainBranding::Storyteller)); + } + + #[test] + fn api_dot_storyteller_dot_ai() { + let request = TestRequest::get() + .uri("https://api.storyteller.ai") + .to_http_request(); + assert_eq!(get_request_domain_branding(&request), Some(DomainBranding::Storyteller)); + } +} diff --git a/crates/service/web/storyteller_web/src/http_server/requests/mod.rs b/crates/service/web/storyteller_web/src/http_server/requests/mod.rs new file mode 100644 index 0000000000..3fa9521ab5 --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/requests/mod.rs @@ -0,0 +1 @@ +pub mod get_request_domain_branding; \ No newline at end of file