diff --git a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_job_type.rs b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_job_type.rs index ca9d401a7f..3f491f43f1 100644 --- a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_job_type.rs +++ b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_job_type.rs @@ -15,11 +15,18 @@ use strum::EnumIter; #[derive(Clone, Copy, Eq, PartialEq, Hash, Ord, PartialOrd, Deserialize, Serialize, Default)] #[serde(rename_all = "snake_case")] pub enum InferenceJobType { - /// Live Portrait Jobs. These still run in comfy, but we call it "live_portrait" + /// Storyteller Studio and Video Style Transfer Jobs (which we may want to split). + /// These run in Comfy. + /// TODO(bt,2024-07-15): We may segregate these two job types in the future + VideoRender, + + /// Live Portrait Jobs. + /// These run in Comfy. LivePortrait, /// Jobs that run ComfyUI workflows /// This is actually just for Video Style Transfer and Storyteller Studio. + #[deprecated(note = "Use VideoRender instead.")] ComfyUi, /// A job that turns "FBX" game engine files into "GLTF" files (Bevy-compatible). @@ -76,6 +83,7 @@ impl_mysql_enum_coders!(InferenceJobType); impl InferenceJobType { pub fn to_str(&self) -> &'static str { match self { + Self::VideoRender => "video_render", Self::LivePortrait => "live_portrait", Self::ComfyUi => "comfy_ui", Self::ConvertFbxToGltf => "convert_fbx_gltf", @@ -94,6 +102,7 @@ impl InferenceJobType { pub fn from_str(value: &str) -> Result { match value { + "video_render" => Ok(Self::VideoRender), "live_portrait" => Ok(Self::LivePortrait), "comfy_ui" => Ok(Self::ComfyUi), "convert_fbx_gltf" => Ok(Self::ConvertFbxToGltf), @@ -115,6 +124,7 @@ impl InferenceJobType { // NB: BTreeSet is sorted // NB: BTreeSet::from() isn't const, but not worth using LazyStatic, etc. BTreeSet::from([ + Self::VideoRender, Self::LivePortrait, Self::ComfyUi, Self::ConvertFbxToGltf, @@ -147,6 +157,7 @@ mod tests { #[test] fn test_serialization() { + assert_serialization(InferenceJobType::VideoRender, "video_render"); assert_serialization(InferenceJobType::LivePortrait, "live_portrait"); assert_serialization(InferenceJobType::ComfyUi, "comfy_ui"); assert_serialization(InferenceJobType::ConvertFbxToGltf, "convert_fbx_gltf"); @@ -164,6 +175,7 @@ mod tests { #[test] fn to_str() { + assert_eq!(InferenceJobType::VideoRender.to_str(), "video_render"); assert_eq!(InferenceJobType::LivePortrait.to_str(), "live_portrait"); assert_eq!(InferenceJobType::ComfyUi.to_str(), "comfy_ui"); assert_eq!(InferenceJobType::ConvertFbxToGltf.to_str(), "convert_fbx_gltf"); @@ -181,6 +193,7 @@ mod tests { #[test] fn from_str() { + assert_eq!(InferenceJobType::from_str("video_render").unwrap(), InferenceJobType::VideoRender); assert_eq!(InferenceJobType::from_str("live_portrait").unwrap(), InferenceJobType::LivePortrait); assert_eq!(InferenceJobType::from_str("comfy_ui").unwrap(), InferenceJobType::ComfyUi); assert_eq!(InferenceJobType::from_str("convert_fbx_gltf").unwrap(), InferenceJobType::ConvertFbxToGltf); @@ -200,7 +213,8 @@ mod tests { fn all_variants() { // Static check let mut variants = InferenceJobType::all_variants(); - assert_eq!(variants.len(), 13); + assert_eq!(variants.len(), 14); + assert_eq!(variants.pop_first(), Some(InferenceJobType::VideoRender)); assert_eq!(variants.pop_first(), Some(InferenceJobType::LivePortrait)); assert_eq!(variants.pop_first(), Some(InferenceJobType::ComfyUi)); assert_eq!(variants.pop_first(), Some(InferenceJobType::ConvertFbxToGltf)); 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 new file mode 100644 index 0000000000..7724fae930 --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue/enqueue_live_portrait_workflow_handler.rs @@ -0,0 +1,504 @@ +// NB: Incrementally getting rid of build warnings... +#![forbid(unused_imports)] +#![forbid(unused_mut)] +#![forbid(unused_variables)] + +use std::sync::Arc; + +use actix_web::{HttpRequest, HttpResponse, web}; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use log::{error, info, warn}; +use utoipa::ToSchema; + +use enums::by_table::generic_inference_jobs::inference_category::InferenceCategory; +use enums::by_table::generic_inference_jobs::inference_job_type::InferenceJobType; +use enums::by_table::generic_inference_jobs::inference_model_type::InferenceModelType; +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::generic_inference_args::{GenericInferenceArgs, InferenceCategoryAbbreviated, PolymorphicInferenceArgs}; +use mysql_queries::payloads::generic_inference_args::workflow_payload::WorkflowArgs; +use mysql_queries::queries::generic_inference::web::insert_generic_inference_job::{insert_generic_inference_job, InsertGenericInferenceArgs}; +use mysql_queries::queries::idepotency_tokens::insert_idempotency_token::insert_idempotency_token; +use primitives::lazy_any_option_true::lazy_any_option_true; +use primitives::str_to_bool::str_to_bool; +use primitives::traits::trim_or_emptyable::TrimOrEmptyable; +use primitives::try_str_to_num::try_str_to_num; +use tokens::tokens::generic_inference_jobs::InferenceJobToken; +use tokens::tokens::media_files::MediaFileToken; + +use crate::configs::plans::get_correct_plan_for_session::get_correct_plan_for_session; +use crate::configs::plans::plan_category::PlanCategory; +use crate::http_server::endpoints::workflows::coordinate_workflow_args::{coordinate_workflow_args, CoordinatedWorkflowArgs}; +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; +use crate::state::server_state::ServerState; +use crate::util::allowed_video_style_transfer_access::allowed_video_style_transfer_access; + +/// Debug requests can get routed to special "debug-only" workers, which can +/// be used to trial new code, run debugging, etc. +const DEBUG_HEADER_NAME : &str = "enable-debug-mode"; + +/// The routing tag header can send workloads to particular k8s hosts. +/// This is useful for catching the live logs or intercepting the job. +const ROUTING_TAG_HEADER_NAME : &str = "routing-tag"; + +/// The default ending trim point of a video if not supplied in the request. +const DEFAULT_TRIM_MILLISECONDS_END : u64 = 7_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) + uuid_idempotency_token: String, + + /// The name of the style to invoke (required) + style: StyleTransferName, + + /// The input video media file (required) + input_file: MediaFileToken, + + /// Optional: the depth video file + /// The underlying media file must be a video. + input_depth_file: Option, + + /// Optional: the normal map video file + /// The underlying media file must be a video. + input_normal_file: Option, + + /// Optional: the outline video file + /// The underlying media file must be a video. + input_outline_file: Option, + + /// Optional: Global IP-Adapter image. + /// The underlying media file must be an image. + global_ipa_media_token: Option, + + /// The positive prompt (optional) + prompt: Option, + + /// The negative prompt (optional) + negative_prompt: Option, + + /// Use Strength of the style transfer + /// Must be between 0.0 (match source) and 1.0 (maximum dreaming). + /// The default, if not sent, is 1.0. + use_strength: Option, + + /// Optional trim start in milliseconds + trim_start_millis: Option, + + /// Optional trim end in milliseconds + trim_end_millis: Option, + + /// Enable lipsyncing in the workflow + enable_lipsync: Option, + + /// Use lipsync in the workflow + #[deprecated(note = "use enable_lipsync")] + use_lipsync: Option, + + /// Remove watermark from the output + /// Only for premium accounts + remove_watermark: Option, + + /// Disable LCM + /// Don't let ordinary users do this. + /// Non-LCM workflows take a long time. + disable_lcm: Option, + + /// Use the cinematic workflow + /// Don't let ordinary users do this. + use_cinematic: Option, + + /// Use face detailer + /// Only for premium accounts + use_face_detailer: Option, + + /// Use video upscaler + /// Only for premium accounts + use_upscaler: Option, + + /// Optional visibility setting override. + creator_set_visibility: Option, + + /// To use prompt traveling, enqueue a prompt with this. + /// + /// An example of a prompt travel section: + /// + /// "50": "(explosion:1.4)", + /// "59": "(explosion:1.4)", + /// "60": "fire, smoke, ash", + /// "99": "fire, smoke, ash", + /// "100": "", + /// + /// On the interval 50-59 there will be explosion, on + /// the interval 60-99 there will be fire, smoke, ash. + /// You must set + /// + /// Alternatively, the main `prompt` field can have a + /// prompt travel section if you end the main prompt + /// with `---` and then include a prompt travel text. + travel_prompt: Option, + + /// If you'd like to skip frames. + /// + /// Meaning: + /// * 1 = not skipping frames. + /// * 2 = skipping every 2nd, etc. + /// + /// The default, if not specified, is 2, which means + /// every 2nd frame is skipped. + frame_skip: Option, +} + +#[derive(Serialize, ToSchema)] +pub struct EnqueueVideoStyleTransferSuccessResponse { + pub success: bool, + pub inference_job_token: InferenceJobToken, +} + +#[derive(Debug, ToSchema)] +pub enum EnqueueVideoStyleTransferError { + BadInput(String), + NotAuthorized, + ServerError, + RateLimited, +} + +impl ResponseError for EnqueueVideoStyleTransferError { + fn status_code(&self) -> StatusCode { + match *self { + EnqueueVideoStyleTransferError::BadInput(_) => StatusCode::BAD_REQUEST, + EnqueueVideoStyleTransferError::NotAuthorized => StatusCode::UNAUTHORIZED, + EnqueueVideoStyleTransferError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + EnqueueVideoStyleTransferError::RateLimited => StatusCode::TOO_MANY_REQUESTS, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + EnqueueVideoStyleTransferError::BadInput(reason) => reason.to_string(), + EnqueueVideoStyleTransferError::NotAuthorized => "unauthorized".to_string(), + EnqueueVideoStyleTransferError::ServerError => "server error".to_string(), + EnqueueVideoStyleTransferError::RateLimited => "rate limited".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl std::fmt::Display for EnqueueVideoStyleTransferError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{:?}", self) + } +} + +#[utoipa::path( + post, + tag = "Workflows", + path = "/v1/workflows/enqueue_live_portrait", + responses( + (status = 200, description = "Success", body = EnqueueVideoStyleTransferSuccessResponse), + (status = 400, description = "Bad input", body = EnqueueVideoStyleTransferError), + (status = 401, description = "Not authorized", body = EnqueueVideoStyleTransferError), + (status = 429, description = "Rate limited", body = EnqueueVideoStyleTransferError), + (status = 500, description = "Server error", body = EnqueueVideoStyleTransferError) + ), + params(("request" = EnqueueVideoStyleTransferRequest, description = "Payload for request")) +)] +pub async fn enqueue_video_style_transfer_workflow_handler( + http_request: HttpRequest, + request: web::Json, + server_state: web::Data>) -> Result +{ + // ==================== VALIDATION ==================== // + + match request.frame_skip { + None | Some(1) | Some(2) => {} // Allowed + _ => { + return Err(EnqueueVideoStyleTransferError::BadInput("Invalid frame skip value".to_string())); + } + } + + // ==================== DB ==================== // + + let mut mysql_connection = server_state.mysql_pool + .acquire() + .await + .map_err(|err| { + warn!("MySql pool error: {:?}", err); + EnqueueVideoStyleTransferError::ServerError + })?; + + let maybe_avt_token = server_state.avt_cookie_manager + .get_avt_token_from_request(&http_request); + + // ==================== USER SESSION ==================== // + + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session_extended_from_connection(&http_request, &mut mysql_connection) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + EnqueueVideoStyleTransferError::ServerError + })?; + + let maybe_user_token = maybe_user_session + .as_ref() + .map(|session| session.user_token_typed.clone()); + + // ==================== FEATURE FLAG CHECK ==================== // + + if !allowed_video_style_transfer_access(maybe_user_session.as_ref(), &server_state.flags) { + warn!("Video style transfer access is not permitted for user"); + return Err(EnqueueVideoStyleTransferError::NotAuthorized); + } + + // ==================== PAID PLAN + PRIORITY ==================== // + + // TODO: Plan should handle "first anonymous use" and "investor" cases. + let plan = get_correct_plan_for_session( + server_state.server_environment, + maybe_user_session.as_ref()); + + // TODO: Separate priority for animation. + let priority_level = plan.web_vc_base_priority_level(); + + let is_staff = maybe_user_session + .as_ref() + .map(|user| user.role.can_ban_users) + .unwrap_or(false); + + // ==================== DEBUG MODE + ROUTING TAG ==================== // + + let is_debug_request = + get_request_header_optional(&http_request, DEBUG_HEADER_NAME) + .is_some(); + + let maybe_routing_tag= + get_request_header_optional(&http_request, ROUTING_TAG_HEADER_NAME) + .map(|routing_tag| routing_tag.trim().to_string()); + + // ==================== RATE LIMIT ==================== // + + let rate_limiter = match maybe_user_session { + None => &server_state.redis_rate_limiters.logged_out, + Some(ref user) => { + if user.role.is_banned { + return Err(EnqueueVideoStyleTransferError::NotAuthorized); + } + &server_state.redis_rate_limiters.logged_in + }, + }; + + if let Err(_err) = rate_limiter.rate_limit_request(&http_request) { + return Err(EnqueueVideoStyleTransferError::RateLimited); + } + + // ==================== HANDLE IDEMPOTENCY ==================== // + + if let Err(reason) = validate_idempotency_token_format(&request.uuid_idempotency_token) { + return Err(EnqueueVideoStyleTransferError::BadInput(reason)); + } + + insert_idempotency_token(&request.uuid_idempotency_token, &mut *mysql_connection) + .await + .map_err(|err| { + error!("Error inserting idempotency token: {:?}", err); + EnqueueVideoStyleTransferError::BadInput("invalid idempotency token".to_string()) + })?; + + // ==================== LOOK UP MODEL INFO ==================== // + + // TODO(bt): CHECK DATABASE FOR TOKENS! + + let ip_address = get_request_ip(&http_request); + + let maybe_user_preferred_visibility : Option = maybe_user_session + .as_ref() + .map(|user_session: &UserSessionExtended| user_session.preferences.preferred_tts_result_visibility); // TODO: New setting for web-vc + + let set_visibility = request.creator_set_visibility + .or(maybe_user_preferred_visibility) + .unwrap_or(Visibility::Public); + + let mut trim_start_millis = request.trim_start_millis.unwrap_or(0); + let mut trim_end_millis = request.trim_end_millis.unwrap_or(DEFAULT_TRIM_MILLISECONDS_END); + + let has_paid_plan = plan.plan_slug() == "fakeyou_contributor" || plan.plan_category() == PlanCategory::Paid; + + // block trim too much + if has_paid_plan { + if trim_end_millis - trim_start_millis > MAXIMUM_DURATION_MILLIS { + trim_start_millis = 0; + trim_end_millis = MAXIMUM_DURATION_MILLIS; + } + } else { + if trim_end_millis - trim_start_millis > DEFAULT_TRIM_MILLISECONDS_END { + trim_start_millis = 0; + trim_end_millis = DEFAULT_TRIM_MILLISECONDS_END; + } + } + + 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()?; + + let coordinated_args = CoordinatedWorkflowArgs { + prompt: request.prompt.new_string_trim_or_empty(), + travel_prompt: request.travel_prompt.new_string_trim_or_empty(), + use_lipsync: lazy_any_option_true(&[ + Box::new(|| request.enable_lipsync), + Box::new(|| request.use_lipsync), + Box::new(|| get_request_header_optional(&http_request, "LIPSYNC-ENABLED") + .map(|value| str_to_bool(&value))) + ]), + disable_lcm: lazy_any_option_true(&[ + Box::new(|| request.disable_lcm), + Box::new(|| get_request_header_optional(&http_request, "DISABLE-LCM") + .map(|value| str_to_bool(&value))) + ]), + use_cinematic: lazy_any_option_true(&[ + Box::new(|| request.use_cinematic), + Box::new(|| get_request_header_optional(&http_request, "USE-CINEMATIC") + .map(|value| str_to_bool(&value))) + ]), + use_face_detailer: request.use_face_detailer, + use_upscaler: request.use_upscaler, + remove_watermark: request.remove_watermark, + }; + + let is_allowed_expensive_generation = is_staff || has_paid_plan; + + let coordinated_args = coordinate_workflow_args(coordinated_args, is_allowed_expensive_generation); + + let inference_args = WorkflowArgs { + // Main text prompts + positive_prompt: coordinated_args.prompt.new_string_trim_or_empty(), + negative_prompt: request.negative_prompt.new_string_trim_or_empty(), + travel_prompt: coordinated_args.travel_prompt.new_string_trim_or_empty(), + + // Input files + maybe_input_file: Some(request.input_file.clone()), + maybe_input_depth_file: empty_token_to_null(request.input_depth_file.as_ref()), + maybe_input_normal_file: empty_token_to_null(request.input_normal_file.as_ref()), + maybe_input_outline_file: empty_token_to_null(request.input_outline_file.as_ref()), + global_ip_adapter_token: empty_token_to_null(request.global_ipa_media_token.as_ref()), + + // Other inputs + style_name: Some(request.style), + creator_visibility: Some(set_visibility), + trim_start_milliseconds: Some(trim_start_millis), + trim_end_milliseconds: Some(trim_end_millis), + strength: maybe_strength, + frame_skip: request.frame_skip, + + // Flags + disable_lcm: coordinated_args.disable_lcm, + 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. + + // TODO: Get rid of the temporary flags. + rollout_python_workflow_args: None, + skip_process_video: get_request_header_optional(&http_request, "SKIP-PROCESS-VIDEO") + .map(|value| str_to_bool(&value)), + sleep_millis: get_request_header_optional(&http_request, "SLEEP-MILLIS") + .and_then(|value| try_str_to_num(&value).ok()), + + // The new, simplified enqueuing doesn't care about the following parameters: + maybe_lora_model: None, + maybe_json_modifications: None, + maybe_workflow_config: None, + maybe_output_path: None, + maybe_google_drive_link: None, + maybe_title: None, + maybe_commit_hash:None, + maybe_description:None, + trim_start_seconds: None, + trim_end_seconds: None, + target_fps: None, + }; + + info!("Creating ComfyUI job record..."); + + let query_result = insert_generic_inference_job(InsertGenericInferenceArgs { + uuid_idempotency_token: &request.uuid_idempotency_token, + job_type: InferenceJobType::ComfyUi, + inference_category: InferenceCategory::Workflow, + maybe_model_type: Some(InferenceModelType::ComfyUi), // NB: Model is static during inference + maybe_model_token: None, // NB: Model is static during inference + maybe_input_source_token: None, // TODO: Introduce a second foreign key ? + maybe_input_source_token_type: None, // TODO: Introduce a second foreign key ? + maybe_download_url: None, + maybe_cover_image_media_file_token: None, + maybe_raw_inference_text: None, // No text + maybe_max_duration_seconds: None, + maybe_inference_args: Some(GenericInferenceArgs { + inference_category: Some(InferenceCategoryAbbreviated::Workflow), + args: Some(PolymorphicInferenceArgs::Cu(inference_args)), + }), + maybe_creator_user_token: maybe_user_token.as_ref(), + maybe_avt_token: maybe_avt_token.as_ref(), + creator_ip_address: &ip_address, + creator_set_visibility: set_visibility, + priority_level, + requires_keepalive: plan.workflow_requires_frontend_keepalive(), + is_debug_request, + maybe_routing_tag: maybe_routing_tag.as_deref(), + mysql_pool: &server_state.mysql_pool, + }).await; + + let job_token = match query_result { + Ok((job_token, _id)) => job_token, + Err(err) => { + warn!("New generic inference job creation DB error: {:?}", err); + return Err(EnqueueVideoStyleTransferError::ServerError); + } + }; + + let response: EnqueueVideoStyleTransferSuccessResponse = EnqueueVideoStyleTransferSuccessResponse { + success: true, + inference_job_token: job_token, + }; + + let body = serde_json::to_string(&response) + .map_err(|_e| EnqueueVideoStyleTransferError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} + +// NB: Frontend is mistakenly sending empty string tokens - ignore those +fn empty_token_to_null(maybe_token: Option<&MediaFileToken>) -> Option { + maybe_token + .filter(|t| t.as_str().trim_or_empty().is_some()) + .map(|t| t.clone()) +} 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 new file mode 100644 index 0000000000..e3abcf5dce --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue/enqueue_studio_workflow_handler.rs @@ -0,0 +1,504 @@ +// NB: Incrementally getting rid of build warnings... +#![forbid(unused_imports)] +#![forbid(unused_mut)] +#![forbid(unused_variables)] + +use std::sync::Arc; + +use actix_web::{HttpRequest, HttpResponse, web}; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use log::{error, info, warn}; +use utoipa::ToSchema; + +use enums::by_table::generic_inference_jobs::inference_category::InferenceCategory; +use enums::by_table::generic_inference_jobs::inference_job_type::InferenceJobType; +use enums::by_table::generic_inference_jobs::inference_model_type::InferenceModelType; +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::generic_inference_args::{GenericInferenceArgs, InferenceCategoryAbbreviated, PolymorphicInferenceArgs}; +use mysql_queries::payloads::generic_inference_args::workflow_payload::WorkflowArgs; +use mysql_queries::queries::generic_inference::web::insert_generic_inference_job::{insert_generic_inference_job, InsertGenericInferenceArgs}; +use mysql_queries::queries::idepotency_tokens::insert_idempotency_token::insert_idempotency_token; +use primitives::lazy_any_option_true::lazy_any_option_true; +use primitives::str_to_bool::str_to_bool; +use primitives::traits::trim_or_emptyable::TrimOrEmptyable; +use primitives::try_str_to_num::try_str_to_num; +use tokens::tokens::generic_inference_jobs::InferenceJobToken; +use tokens::tokens::media_files::MediaFileToken; + +use crate::configs::plans::get_correct_plan_for_session::get_correct_plan_for_session; +use crate::configs::plans::plan_category::PlanCategory; +use crate::http_server::endpoints::workflows::coordinate_workflow_args::{coordinate_workflow_args, CoordinatedWorkflowArgs}; +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; +use crate::state::server_state::ServerState; +use crate::util::allowed_video_style_transfer_access::allowed_video_style_transfer_access; + +/// Debug requests can get routed to special "debug-only" workers, which can +/// be used to trial new code, run debugging, etc. +const DEBUG_HEADER_NAME : &str = "enable-debug-mode"; + +/// The routing tag header can send workloads to particular k8s hosts. +/// This is useful for catching the live logs or intercepting the job. +const ROUTING_TAG_HEADER_NAME : &str = "routing-tag"; + +/// The default ending trim point of a video if not supplied in the request. +const DEFAULT_TRIM_MILLISECONDS_END : u64 = 7_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) + uuid_idempotency_token: String, + + /// The name of the style to invoke (required) + style: StyleTransferName, + + /// The input video media file (required) + input_file: MediaFileToken, + + /// Optional: the depth video file + /// The underlying media file must be a video. + input_depth_file: Option, + + /// Optional: the normal map video file + /// The underlying media file must be a video. + input_normal_file: Option, + + /// Optional: the outline video file + /// The underlying media file must be a video. + input_outline_file: Option, + + /// Optional: Global IP-Adapter image. + /// The underlying media file must be an image. + global_ipa_media_token: Option, + + /// The positive prompt (optional) + prompt: Option, + + /// The negative prompt (optional) + negative_prompt: Option, + + /// Use Strength of the style transfer + /// Must be between 0.0 (match source) and 1.0 (maximum dreaming). + /// The default, if not sent, is 1.0. + use_strength: Option, + + /// Optional trim start in milliseconds + trim_start_millis: Option, + + /// Optional trim end in milliseconds + trim_end_millis: Option, + + /// Enable lipsyncing in the workflow + enable_lipsync: Option, + + /// Use lipsync in the workflow + #[deprecated(note = "use enable_lipsync")] + use_lipsync: Option, + + /// Remove watermark from the output + /// Only for premium accounts + remove_watermark: Option, + + /// Disable LCM + /// Don't let ordinary users do this. + /// Non-LCM workflows take a long time. + disable_lcm: Option, + + /// Use the cinematic workflow + /// Don't let ordinary users do this. + use_cinematic: Option, + + /// Use face detailer + /// Only for premium accounts + use_face_detailer: Option, + + /// Use video upscaler + /// Only for premium accounts + use_upscaler: Option, + + /// Optional visibility setting override. + creator_set_visibility: Option, + + /// To use prompt traveling, enqueue a prompt with this. + /// + /// An example of a prompt travel section: + /// + /// "50": "(explosion:1.4)", + /// "59": "(explosion:1.4)", + /// "60": "fire, smoke, ash", + /// "99": "fire, smoke, ash", + /// "100": "", + /// + /// On the interval 50-59 there will be explosion, on + /// the interval 60-99 there will be fire, smoke, ash. + /// You must set + /// + /// Alternatively, the main `prompt` field can have a + /// prompt travel section if you end the main prompt + /// with `---` and then include a prompt travel text. + travel_prompt: Option, + + /// If you'd like to skip frames. + /// + /// Meaning: + /// * 1 = not skipping frames. + /// * 2 = skipping every 2nd, etc. + /// + /// The default, if not specified, is 2, which means + /// every 2nd frame is skipped. + frame_skip: Option, +} + +#[derive(Serialize, ToSchema)] +pub struct EnqueueVideoStyleTransferSuccessResponse { + pub success: bool, + pub inference_job_token: InferenceJobToken, +} + +#[derive(Debug, ToSchema)] +pub enum EnqueueVideoStyleTransferError { + BadInput(String), + NotAuthorized, + ServerError, + RateLimited, +} + +impl ResponseError for EnqueueVideoStyleTransferError { + fn status_code(&self) -> StatusCode { + match *self { + EnqueueVideoStyleTransferError::BadInput(_) => StatusCode::BAD_REQUEST, + EnqueueVideoStyleTransferError::NotAuthorized => StatusCode::UNAUTHORIZED, + EnqueueVideoStyleTransferError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + EnqueueVideoStyleTransferError::RateLimited => StatusCode::TOO_MANY_REQUESTS, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + EnqueueVideoStyleTransferError::BadInput(reason) => reason.to_string(), + EnqueueVideoStyleTransferError::NotAuthorized => "unauthorized".to_string(), + EnqueueVideoStyleTransferError::ServerError => "server error".to_string(), + EnqueueVideoStyleTransferError::RateLimited => "rate limited".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl std::fmt::Display for EnqueueVideoStyleTransferError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{:?}", self) + } +} + +#[utoipa::path( + post, + tag = "Workflows", + path = "/v1/workflows/enqueue_studio", + responses( + (status = 200, description = "Success", body = EnqueueVideoStyleTransferSuccessResponse), + (status = 400, description = "Bad input", body = EnqueueVideoStyleTransferError), + (status = 401, description = "Not authorized", body = EnqueueVideoStyleTransferError), + (status = 429, description = "Rate limited", body = EnqueueVideoStyleTransferError), + (status = 500, description = "Server error", body = EnqueueVideoStyleTransferError) + ), + params(("request" = EnqueueVideoStyleTransferRequest, description = "Payload for request")) +)] +pub async fn enqueue_video_style_transfer_workflow_handler( + http_request: HttpRequest, + request: web::Json, + server_state: web::Data>) -> Result +{ + // ==================== VALIDATION ==================== // + + match request.frame_skip { + None | Some(1) | Some(2) => {} // Allowed + _ => { + return Err(EnqueueVideoStyleTransferError::BadInput("Invalid frame skip value".to_string())); + } + } + + // ==================== DB ==================== // + + let mut mysql_connection = server_state.mysql_pool + .acquire() + .await + .map_err(|err| { + warn!("MySql pool error: {:?}", err); + EnqueueVideoStyleTransferError::ServerError + })?; + + let maybe_avt_token = server_state.avt_cookie_manager + .get_avt_token_from_request(&http_request); + + // ==================== USER SESSION ==================== // + + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session_extended_from_connection(&http_request, &mut mysql_connection) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + EnqueueVideoStyleTransferError::ServerError + })?; + + let maybe_user_token = maybe_user_session + .as_ref() + .map(|session| session.user_token_typed.clone()); + + // ==================== FEATURE FLAG CHECK ==================== // + + if !allowed_video_style_transfer_access(maybe_user_session.as_ref(), &server_state.flags) { + warn!("Video style transfer access is not permitted for user"); + return Err(EnqueueVideoStyleTransferError::NotAuthorized); + } + + // ==================== PAID PLAN + PRIORITY ==================== // + + // TODO: Plan should handle "first anonymous use" and "investor" cases. + let plan = get_correct_plan_for_session( + server_state.server_environment, + maybe_user_session.as_ref()); + + // TODO: Separate priority for animation. + let priority_level = plan.web_vc_base_priority_level(); + + let is_staff = maybe_user_session + .as_ref() + .map(|user| user.role.can_ban_users) + .unwrap_or(false); + + // ==================== DEBUG MODE + ROUTING TAG ==================== // + + let is_debug_request = + get_request_header_optional(&http_request, DEBUG_HEADER_NAME) + .is_some(); + + let maybe_routing_tag= + get_request_header_optional(&http_request, ROUTING_TAG_HEADER_NAME) + .map(|routing_tag| routing_tag.trim().to_string()); + + // ==================== RATE LIMIT ==================== // + + let rate_limiter = match maybe_user_session { + None => &server_state.redis_rate_limiters.logged_out, + Some(ref user) => { + if user.role.is_banned { + return Err(EnqueueVideoStyleTransferError::NotAuthorized); + } + &server_state.redis_rate_limiters.logged_in + }, + }; + + if let Err(_err) = rate_limiter.rate_limit_request(&http_request) { + return Err(EnqueueVideoStyleTransferError::RateLimited); + } + + // ==================== HANDLE IDEMPOTENCY ==================== // + + if let Err(reason) = validate_idempotency_token_format(&request.uuid_idempotency_token) { + return Err(EnqueueVideoStyleTransferError::BadInput(reason)); + } + + insert_idempotency_token(&request.uuid_idempotency_token, &mut *mysql_connection) + .await + .map_err(|err| { + error!("Error inserting idempotency token: {:?}", err); + EnqueueVideoStyleTransferError::BadInput("invalid idempotency token".to_string()) + })?; + + // ==================== LOOK UP MODEL INFO ==================== // + + // TODO(bt): CHECK DATABASE FOR TOKENS! + + let ip_address = get_request_ip(&http_request); + + let maybe_user_preferred_visibility : Option = maybe_user_session + .as_ref() + .map(|user_session: &UserSessionExtended| user_session.preferences.preferred_tts_result_visibility); // TODO: New setting for web-vc + + let set_visibility = request.creator_set_visibility + .or(maybe_user_preferred_visibility) + .unwrap_or(Visibility::Public); + + let mut trim_start_millis = request.trim_start_millis.unwrap_or(0); + let mut trim_end_millis = request.trim_end_millis.unwrap_or(DEFAULT_TRIM_MILLISECONDS_END); + + let has_paid_plan = plan.plan_slug() == "fakeyou_contributor" || plan.plan_category() == PlanCategory::Paid; + + // block trim too much + if has_paid_plan { + if trim_end_millis - trim_start_millis > MAXIMUM_DURATION_MILLIS { + trim_start_millis = 0; + trim_end_millis = MAXIMUM_DURATION_MILLIS; + } + } else { + if trim_end_millis - trim_start_millis > DEFAULT_TRIM_MILLISECONDS_END { + trim_start_millis = 0; + trim_end_millis = DEFAULT_TRIM_MILLISECONDS_END; + } + } + + 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()?; + + let coordinated_args = CoordinatedWorkflowArgs { + prompt: request.prompt.new_string_trim_or_empty(), + travel_prompt: request.travel_prompt.new_string_trim_or_empty(), + use_lipsync: lazy_any_option_true(&[ + Box::new(|| request.enable_lipsync), + Box::new(|| request.use_lipsync), + Box::new(|| get_request_header_optional(&http_request, "LIPSYNC-ENABLED") + .map(|value| str_to_bool(&value))) + ]), + disable_lcm: lazy_any_option_true(&[ + Box::new(|| request.disable_lcm), + Box::new(|| get_request_header_optional(&http_request, "DISABLE-LCM") + .map(|value| str_to_bool(&value))) + ]), + use_cinematic: lazy_any_option_true(&[ + Box::new(|| request.use_cinematic), + Box::new(|| get_request_header_optional(&http_request, "USE-CINEMATIC") + .map(|value| str_to_bool(&value))) + ]), + use_face_detailer: request.use_face_detailer, + use_upscaler: request.use_upscaler, + remove_watermark: request.remove_watermark, + }; + + let is_allowed_expensive_generation = is_staff || has_paid_plan; + + let coordinated_args = coordinate_workflow_args(coordinated_args, is_allowed_expensive_generation); + + let inference_args = WorkflowArgs { + // Main text prompts + positive_prompt: coordinated_args.prompt.new_string_trim_or_empty(), + negative_prompt: request.negative_prompt.new_string_trim_or_empty(), + travel_prompt: coordinated_args.travel_prompt.new_string_trim_or_empty(), + + // Input files + maybe_input_file: Some(request.input_file.clone()), + maybe_input_depth_file: empty_token_to_null(request.input_depth_file.as_ref()), + maybe_input_normal_file: empty_token_to_null(request.input_normal_file.as_ref()), + maybe_input_outline_file: empty_token_to_null(request.input_outline_file.as_ref()), + global_ip_adapter_token: empty_token_to_null(request.global_ipa_media_token.as_ref()), + + // Other inputs + style_name: Some(request.style), + creator_visibility: Some(set_visibility), + trim_start_milliseconds: Some(trim_start_millis), + trim_end_milliseconds: Some(trim_end_millis), + strength: maybe_strength, + frame_skip: request.frame_skip, + + // Flags + disable_lcm: coordinated_args.disable_lcm, + 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. + + // TODO: Get rid of the temporary flags. + rollout_python_workflow_args: None, + skip_process_video: get_request_header_optional(&http_request, "SKIP-PROCESS-VIDEO") + .map(|value| str_to_bool(&value)), + sleep_millis: get_request_header_optional(&http_request, "SLEEP-MILLIS") + .and_then(|value| try_str_to_num(&value).ok()), + + // The new, simplified enqueuing doesn't care about the following parameters: + maybe_lora_model: None, + maybe_json_modifications: None, + maybe_workflow_config: None, + maybe_output_path: None, + maybe_google_drive_link: None, + maybe_title: None, + maybe_commit_hash:None, + maybe_description:None, + trim_start_seconds: None, + trim_end_seconds: None, + target_fps: None, + }; + + info!("Creating ComfyUI job record..."); + + let query_result = insert_generic_inference_job(InsertGenericInferenceArgs { + uuid_idempotency_token: &request.uuid_idempotency_token, + job_type: InferenceJobType::ComfyUi, + inference_category: InferenceCategory::Workflow, + maybe_model_type: Some(InferenceModelType::ComfyUi), // NB: Model is static during inference + maybe_model_token: None, // NB: Model is static during inference + maybe_input_source_token: None, // TODO: Introduce a second foreign key ? + maybe_input_source_token_type: None, // TODO: Introduce a second foreign key ? + maybe_download_url: None, + maybe_cover_image_media_file_token: None, + maybe_raw_inference_text: None, // No text + maybe_max_duration_seconds: None, + maybe_inference_args: Some(GenericInferenceArgs { + inference_category: Some(InferenceCategoryAbbreviated::Workflow), + args: Some(PolymorphicInferenceArgs::Cu(inference_args)), + }), + maybe_creator_user_token: maybe_user_token.as_ref(), + maybe_avt_token: maybe_avt_token.as_ref(), + creator_ip_address: &ip_address, + creator_set_visibility: set_visibility, + priority_level, + requires_keepalive: plan.workflow_requires_frontend_keepalive(), + is_debug_request, + maybe_routing_tag: maybe_routing_tag.as_deref(), + mysql_pool: &server_state.mysql_pool, + }).await; + + let job_token = match query_result { + Ok((job_token, _id)) => job_token, + Err(err) => { + warn!("New generic inference job creation DB error: {:?}", err); + return Err(EnqueueVideoStyleTransferError::ServerError); + } + }; + + let response: EnqueueVideoStyleTransferSuccessResponse = EnqueueVideoStyleTransferSuccessResponse { + success: true, + inference_job_token: job_token, + }; + + let body = serde_json::to_string(&response) + .map_err(|_e| EnqueueVideoStyleTransferError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} + +// NB: Frontend is mistakenly sending empty string tokens - ignore those +fn empty_token_to_null(maybe_token: Option<&MediaFileToken>) -> Option { + maybe_token + .filter(|t| t.as_str().trim_or_empty().is_some()) + .map(|t| t.clone()) +} 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 new file mode 100644 index 0000000000..a27039d7c9 --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue/enqueue_video_style_transfer_workflow_handler.rs @@ -0,0 +1,504 @@ +// NB: Incrementally getting rid of build warnings... +#![forbid(unused_imports)] +#![forbid(unused_mut)] +#![forbid(unused_variables)] + +use std::sync::Arc; + +use actix_web::{HttpRequest, HttpResponse, web}; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use log::{error, info, warn}; +use utoipa::ToSchema; + +use enums::by_table::generic_inference_jobs::inference_category::InferenceCategory; +use enums::by_table::generic_inference_jobs::inference_job_type::InferenceJobType; +use enums::by_table::generic_inference_jobs::inference_model_type::InferenceModelType; +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::generic_inference_args::{GenericInferenceArgs, InferenceCategoryAbbreviated, PolymorphicInferenceArgs}; +use mysql_queries::payloads::generic_inference_args::workflow_payload::WorkflowArgs; +use mysql_queries::queries::generic_inference::web::insert_generic_inference_job::{insert_generic_inference_job, InsertGenericInferenceArgs}; +use mysql_queries::queries::idepotency_tokens::insert_idempotency_token::insert_idempotency_token; +use primitives::lazy_any_option_true::lazy_any_option_true; +use primitives::str_to_bool::str_to_bool; +use primitives::traits::trim_or_emptyable::TrimOrEmptyable; +use primitives::try_str_to_num::try_str_to_num; +use tokens::tokens::generic_inference_jobs::InferenceJobToken; +use tokens::tokens::media_files::MediaFileToken; + +use crate::configs::plans::get_correct_plan_for_session::get_correct_plan_for_session; +use crate::configs::plans::plan_category::PlanCategory; +use crate::http_server::endpoints::workflows::coordinate_workflow_args::{coordinate_workflow_args, CoordinatedWorkflowArgs}; +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; +use crate::state::server_state::ServerState; +use crate::util::allowed_video_style_transfer_access::allowed_video_style_transfer_access; + +/// Debug requests can get routed to special "debug-only" workers, which can +/// be used to trial new code, run debugging, etc. +const DEBUG_HEADER_NAME : &str = "enable-debug-mode"; + +/// The routing tag header can send workloads to particular k8s hosts. +/// This is useful for catching the live logs or intercepting the job. +const ROUTING_TAG_HEADER_NAME : &str = "routing-tag"; + +/// The default ending trim point of a video if not supplied in the request. +const DEFAULT_TRIM_MILLISECONDS_END : u64 = 7_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) + uuid_idempotency_token: String, + + /// The name of the style to invoke (required) + style: StyleTransferName, + + /// The input video media file (required) + input_file: MediaFileToken, + + /// Optional: the depth video file + /// The underlying media file must be a video. + input_depth_file: Option, + + /// Optional: the normal map video file + /// The underlying media file must be a video. + input_normal_file: Option, + + /// Optional: the outline video file + /// The underlying media file must be a video. + input_outline_file: Option, + + /// Optional: Global IP-Adapter image. + /// The underlying media file must be an image. + global_ipa_media_token: Option, + + /// The positive prompt (optional) + prompt: Option, + + /// The negative prompt (optional) + negative_prompt: Option, + + /// Use Strength of the style transfer + /// Must be between 0.0 (match source) and 1.0 (maximum dreaming). + /// The default, if not sent, is 1.0. + use_strength: Option, + + /// Optional trim start in milliseconds + trim_start_millis: Option, + + /// Optional trim end in milliseconds + trim_end_millis: Option, + + /// Enable lipsyncing in the workflow + enable_lipsync: Option, + + /// Use lipsync in the workflow + #[deprecated(note = "use enable_lipsync")] + use_lipsync: Option, + + /// Remove watermark from the output + /// Only for premium accounts + remove_watermark: Option, + + /// Disable LCM + /// Don't let ordinary users do this. + /// Non-LCM workflows take a long time. + disable_lcm: Option, + + /// Use the cinematic workflow + /// Don't let ordinary users do this. + use_cinematic: Option, + + /// Use face detailer + /// Only for premium accounts + use_face_detailer: Option, + + /// Use video upscaler + /// Only for premium accounts + use_upscaler: Option, + + /// Optional visibility setting override. + creator_set_visibility: Option, + + /// To use prompt traveling, enqueue a prompt with this. + /// + /// An example of a prompt travel section: + /// + /// "50": "(explosion:1.4)", + /// "59": "(explosion:1.4)", + /// "60": "fire, smoke, ash", + /// "99": "fire, smoke, ash", + /// "100": "", + /// + /// On the interval 50-59 there will be explosion, on + /// the interval 60-99 there will be fire, smoke, ash. + /// You must set + /// + /// Alternatively, the main `prompt` field can have a + /// prompt travel section if you end the main prompt + /// with `---` and then include a prompt travel text. + travel_prompt: Option, + + /// If you'd like to skip frames. + /// + /// Meaning: + /// * 1 = not skipping frames. + /// * 2 = skipping every 2nd, etc. + /// + /// The default, if not specified, is 2, which means + /// every 2nd frame is skipped. + frame_skip: Option, +} + +#[derive(Serialize, ToSchema)] +pub struct EnqueueVideoStyleTransferSuccessResponse { + pub success: bool, + pub inference_job_token: InferenceJobToken, +} + +#[derive(Debug, ToSchema)] +pub enum EnqueueVideoStyleTransferError { + BadInput(String), + NotAuthorized, + ServerError, + RateLimited, +} + +impl ResponseError for EnqueueVideoStyleTransferError { + fn status_code(&self) -> StatusCode { + match *self { + EnqueueVideoStyleTransferError::BadInput(_) => StatusCode::BAD_REQUEST, + EnqueueVideoStyleTransferError::NotAuthorized => StatusCode::UNAUTHORIZED, + EnqueueVideoStyleTransferError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + EnqueueVideoStyleTransferError::RateLimited => StatusCode::TOO_MANY_REQUESTS, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + EnqueueVideoStyleTransferError::BadInput(reason) => reason.to_string(), + EnqueueVideoStyleTransferError::NotAuthorized => "unauthorized".to_string(), + EnqueueVideoStyleTransferError::ServerError => "server error".to_string(), + EnqueueVideoStyleTransferError::RateLimited => "rate limited".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl std::fmt::Display for EnqueueVideoStyleTransferError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{:?}", self) + } +} + +#[utoipa::path( + post, + tag = "Workflows", + path = "/v1/workflows/enqueue_vst", + responses( + (status = 200, description = "Success", body = EnqueueVideoStyleTransferSuccessResponse), + (status = 400, description = "Bad input", body = EnqueueVideoStyleTransferError), + (status = 401, description = "Not authorized", body = EnqueueVideoStyleTransferError), + (status = 429, description = "Rate limited", body = EnqueueVideoStyleTransferError), + (status = 500, description = "Server error", body = EnqueueVideoStyleTransferError) + ), + params(("request" = EnqueueVideoStyleTransferRequest, description = "Payload for request")) +)] +pub async fn enqueue_video_style_transfer_workflow_handler( + http_request: HttpRequest, + request: web::Json, + server_state: web::Data>) -> Result +{ + // ==================== VALIDATION ==================== // + + match request.frame_skip { + None | Some(1) | Some(2) => {} // Allowed + _ => { + return Err(EnqueueVideoStyleTransferError::BadInput("Invalid frame skip value".to_string())); + } + } + + // ==================== DB ==================== // + + let mut mysql_connection = server_state.mysql_pool + .acquire() + .await + .map_err(|err| { + warn!("MySql pool error: {:?}", err); + EnqueueVideoStyleTransferError::ServerError + })?; + + let maybe_avt_token = server_state.avt_cookie_manager + .get_avt_token_from_request(&http_request); + + // ==================== USER SESSION ==================== // + + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session_extended_from_connection(&http_request, &mut mysql_connection) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + EnqueueVideoStyleTransferError::ServerError + })?; + + let maybe_user_token = maybe_user_session + .as_ref() + .map(|session| session.user_token_typed.clone()); + + // ==================== FEATURE FLAG CHECK ==================== // + + if !allowed_video_style_transfer_access(maybe_user_session.as_ref(), &server_state.flags) { + warn!("Video style transfer access is not permitted for user"); + return Err(EnqueueVideoStyleTransferError::NotAuthorized); + } + + // ==================== PAID PLAN + PRIORITY ==================== // + + // TODO: Plan should handle "first anonymous use" and "investor" cases. + let plan = get_correct_plan_for_session( + server_state.server_environment, + maybe_user_session.as_ref()); + + // TODO: Separate priority for animation. + let priority_level = plan.web_vc_base_priority_level(); + + let is_staff = maybe_user_session + .as_ref() + .map(|user| user.role.can_ban_users) + .unwrap_or(false); + + // ==================== DEBUG MODE + ROUTING TAG ==================== // + + let is_debug_request = + get_request_header_optional(&http_request, DEBUG_HEADER_NAME) + .is_some(); + + let maybe_routing_tag= + get_request_header_optional(&http_request, ROUTING_TAG_HEADER_NAME) + .map(|routing_tag| routing_tag.trim().to_string()); + + // ==================== RATE LIMIT ==================== // + + let rate_limiter = match maybe_user_session { + None => &server_state.redis_rate_limiters.logged_out, + Some(ref user) => { + if user.role.is_banned { + return Err(EnqueueVideoStyleTransferError::NotAuthorized); + } + &server_state.redis_rate_limiters.logged_in + }, + }; + + if let Err(_err) = rate_limiter.rate_limit_request(&http_request) { + return Err(EnqueueVideoStyleTransferError::RateLimited); + } + + // ==================== HANDLE IDEMPOTENCY ==================== // + + if let Err(reason) = validate_idempotency_token_format(&request.uuid_idempotency_token) { + return Err(EnqueueVideoStyleTransferError::BadInput(reason)); + } + + insert_idempotency_token(&request.uuid_idempotency_token, &mut *mysql_connection) + .await + .map_err(|err| { + error!("Error inserting idempotency token: {:?}", err); + EnqueueVideoStyleTransferError::BadInput("invalid idempotency token".to_string()) + })?; + + // ==================== LOOK UP MODEL INFO ==================== // + + // TODO(bt): CHECK DATABASE FOR TOKENS! + + let ip_address = get_request_ip(&http_request); + + let maybe_user_preferred_visibility : Option = maybe_user_session + .as_ref() + .map(|user_session: &UserSessionExtended| user_session.preferences.preferred_tts_result_visibility); // TODO: New setting for web-vc + + let set_visibility = request.creator_set_visibility + .or(maybe_user_preferred_visibility) + .unwrap_or(Visibility::Public); + + let mut trim_start_millis = request.trim_start_millis.unwrap_or(0); + let mut trim_end_millis = request.trim_end_millis.unwrap_or(DEFAULT_TRIM_MILLISECONDS_END); + + let has_paid_plan = plan.plan_slug() == "fakeyou_contributor" || plan.plan_category() == PlanCategory::Paid; + + // block trim too much + if has_paid_plan { + if trim_end_millis - trim_start_millis > MAXIMUM_DURATION_MILLIS { + trim_start_millis = 0; + trim_end_millis = MAXIMUM_DURATION_MILLIS; + } + } else { + if trim_end_millis - trim_start_millis > DEFAULT_TRIM_MILLISECONDS_END { + trim_start_millis = 0; + trim_end_millis = DEFAULT_TRIM_MILLISECONDS_END; + } + } + + 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()?; + + let coordinated_args = CoordinatedWorkflowArgs { + prompt: request.prompt.new_string_trim_or_empty(), + travel_prompt: request.travel_prompt.new_string_trim_or_empty(), + use_lipsync: lazy_any_option_true(&[ + Box::new(|| request.enable_lipsync), + Box::new(|| request.use_lipsync), + Box::new(|| get_request_header_optional(&http_request, "LIPSYNC-ENABLED") + .map(|value| str_to_bool(&value))) + ]), + disable_lcm: lazy_any_option_true(&[ + Box::new(|| request.disable_lcm), + Box::new(|| get_request_header_optional(&http_request, "DISABLE-LCM") + .map(|value| str_to_bool(&value))) + ]), + use_cinematic: lazy_any_option_true(&[ + Box::new(|| request.use_cinematic), + Box::new(|| get_request_header_optional(&http_request, "USE-CINEMATIC") + .map(|value| str_to_bool(&value))) + ]), + use_face_detailer: request.use_face_detailer, + use_upscaler: request.use_upscaler, + remove_watermark: request.remove_watermark, + }; + + let is_allowed_expensive_generation = is_staff || has_paid_plan; + + let coordinated_args = coordinate_workflow_args(coordinated_args, is_allowed_expensive_generation); + + let inference_args = WorkflowArgs { + // Main text prompts + positive_prompt: coordinated_args.prompt.new_string_trim_or_empty(), + negative_prompt: request.negative_prompt.new_string_trim_or_empty(), + travel_prompt: coordinated_args.travel_prompt.new_string_trim_or_empty(), + + // Input files + maybe_input_file: Some(request.input_file.clone()), + maybe_input_depth_file: empty_token_to_null(request.input_depth_file.as_ref()), + maybe_input_normal_file: empty_token_to_null(request.input_normal_file.as_ref()), + maybe_input_outline_file: empty_token_to_null(request.input_outline_file.as_ref()), + global_ip_adapter_token: empty_token_to_null(request.global_ipa_media_token.as_ref()), + + // Other inputs + style_name: Some(request.style), + creator_visibility: Some(set_visibility), + trim_start_milliseconds: Some(trim_start_millis), + trim_end_milliseconds: Some(trim_end_millis), + strength: maybe_strength, + frame_skip: request.frame_skip, + + // Flags + disable_lcm: coordinated_args.disable_lcm, + 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. + + // TODO: Get rid of the temporary flags. + rollout_python_workflow_args: None, + skip_process_video: get_request_header_optional(&http_request, "SKIP-PROCESS-VIDEO") + .map(|value| str_to_bool(&value)), + sleep_millis: get_request_header_optional(&http_request, "SLEEP-MILLIS") + .and_then(|value| try_str_to_num(&value).ok()), + + // The new, simplified enqueuing doesn't care about the following parameters: + maybe_lora_model: None, + maybe_json_modifications: None, + maybe_workflow_config: None, + maybe_output_path: None, + maybe_google_drive_link: None, + maybe_title: None, + maybe_commit_hash:None, + maybe_description:None, + trim_start_seconds: None, + trim_end_seconds: None, + target_fps: None, + }; + + info!("Creating ComfyUI job record..."); + + let query_result = insert_generic_inference_job(InsertGenericInferenceArgs { + uuid_idempotency_token: &request.uuid_idempotency_token, + job_type: InferenceJobType::ComfyUi, + inference_category: InferenceCategory::Workflow, + maybe_model_type: Some(InferenceModelType::ComfyUi), // NB: Model is static during inference + maybe_model_token: None, // NB: Model is static during inference + maybe_input_source_token: None, // TODO: Introduce a second foreign key ? + maybe_input_source_token_type: None, // TODO: Introduce a second foreign key ? + maybe_download_url: None, + maybe_cover_image_media_file_token: None, + maybe_raw_inference_text: None, // No text + maybe_max_duration_seconds: None, + maybe_inference_args: Some(GenericInferenceArgs { + inference_category: Some(InferenceCategoryAbbreviated::Workflow), + args: Some(PolymorphicInferenceArgs::Cu(inference_args)), + }), + maybe_creator_user_token: maybe_user_token.as_ref(), + maybe_avt_token: maybe_avt_token.as_ref(), + creator_ip_address: &ip_address, + creator_set_visibility: set_visibility, + priority_level, + requires_keepalive: plan.workflow_requires_frontend_keepalive(), + is_debug_request, + maybe_routing_tag: maybe_routing_tag.as_deref(), + mysql_pool: &server_state.mysql_pool, + }).await; + + let job_token = match query_result { + Ok((job_token, _id)) => job_token, + Err(err) => { + warn!("New generic inference job creation DB error: {:?}", err); + return Err(EnqueueVideoStyleTransferError::ServerError); + } + }; + + let response: EnqueueVideoStyleTransferSuccessResponse = EnqueueVideoStyleTransferSuccessResponse { + success: true, + inference_job_token: job_token, + }; + + let body = serde_json::to_string(&response) + .map_err(|_e| EnqueueVideoStyleTransferError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} + +// NB: Frontend is mistakenly sending empty string tokens - ignore those +fn empty_token_to_null(maybe_token: Option<&MediaFileToken>) -> Option { + maybe_token + .filter(|t| t.as_str().trim_or_empty().is_some()) + .map(|t| t.clone()) +} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue/mod.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue/mod.rs new file mode 100644 index 0000000000..43ac116ddd --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue/mod.rs @@ -0,0 +1,3 @@ +pub mod enqueue_studio_workflow_handler; +pub mod enqueue_video_style_transfer_workflow_handler; +pub mod enqueue_live_portrait_workflow_handler; 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 1634133582..124e9fbf5a 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 @@ -28,11 +28,11 @@ use primitives::traits::trim_or_emptyable::TrimOrEmptyable; use primitives::try_str_to_num::try_str_to_num; use tokens::tokens::generic_inference_jobs::InferenceJobToken; use tokens::tokens::media_files::MediaFileToken; -use crate::http_server::session::lookup::user_session_extended::UserSessionExtended; use crate::configs::plans::get_correct_plan_for_session::get_correct_plan_for_session; use crate::configs::plans::plan_category::PlanCategory; use crate::http_server::endpoints::workflows::coordinate_workflow_args::{coordinate_workflow_args, CoordinatedWorkflowArgs}; +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; use crate::state::server_state::ServerState; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/mod.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/mod.rs index af4b75e83a..4de47c5eeb 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/mod.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/mod.rs @@ -1,4 +1,5 @@ -pub(crate) mod enqueue_comfy_ui_handler; -pub(crate) mod enqueue_workflow_upload_request; +mod coordinate_workflow_args; +pub mod enqueue; pub mod enqueue_video_style_transfer_handler; -mod coordinate_workflow_args; \ No newline at end of file +pub(crate) mod enqueue_comfy_ui_handler; +pub(crate) mod enqueue_workflow_upload_request; \ No newline at end of file