watermark branding from hostname

crates/lib/mysql_queries/src/payloads/generic_inference_args/common/mod.rs

@@ -0,0 +1 @@
+pub mod watermark_type;

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<WatermarkType> is None.
+  #[serde(rename = "n")] // NB: DO NOT CHANGE. It could break live jobs. Renamed to be fewer bytes.
+  NoWatermark,
+}

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<bool>,
 
+  #[serde(rename = "w")]
+  #[serde(skip_serializing_if = "Option::is_none")]
+  pub watermark_type: Option<WatermarkType>,
+
   /// This is a debugging flag.
   #[serde(rename = "sp")]
   #[serde(skip_serializing_if = "Option::is_none")]

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;

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<bool>,
 
+    #[serde(rename = "wt")]
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub watermark_type: Option<WatermarkType>,
+
     // 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")]

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,
   };

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,

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,

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...");

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...");

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

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;

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<DomainBranding> {
+  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));
+  }
+}

crates/service/web/storyteller_web/src/http_server/requests/mod.rs

@@ -0,0 +1 @@
+pub mod get_request_domain_branding;