From ed7a9905ee233e20dcdabb2c46166afa5291baff Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 16 Jan 2025 20:35:01 -0500 Subject: [PATCH] add studio gen2 enqueue endpoint --- .../inference_job_product_category.rs | 12 +- .../inference_job_type.rs | 11 + .../public_inference_job_type.rs | 6 + .../generic_inference_args.rs | 4 + .../inner_payloads/mod.rs | 1 + .../inner_payloads/studio_gen2_payload.rs | 24 ++ .../src/http_server/endpoints/mod.rs | 3 +- .../enqueue_studio_gen2_handler.rs | 295 ++++++++++++++++++ .../http_server/endpoints/studio_gen2/mod.rs | 1 + .../src/http_server/routes/add_routes.rs | 2 + .../routes/add_studio_gen2_routes.rs | 24 ++ .../src/http_server/routes/mod.rs | 1 + 12 files changed, 382 insertions(+), 2 deletions(-) create mode 100644 crates/lib/mysql_queries/src/payloads/generic_inference_args/inner_payloads/studio_gen2_payload.rs create mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/studio_gen2/enqueue_studio_gen2_handler.rs create mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/studio_gen2/mod.rs create mode 100644 crates/service/web/storyteller_web/src/http_server/routes/add_studio_gen2_routes.rs diff --git a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_job_product_category.rs b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_job_product_category.rs index 5bc00a938d..a838621803 100644 --- a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_job_product_category.rs +++ b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_job_product_category.rs @@ -63,6 +63,9 @@ pub enum InferenceJobProductCategory { /// Video: Studio VidStudio, + /// Video: Studio Gen 2 + VidStudioGen2, + /// Video: Style Transfer VidStyleTransfer, @@ -122,6 +125,7 @@ impl InferenceJobProductCategory { Self::VidLivePortrait => "vid_live_portrait", Self::VidLivePortraitWebcam => "vid_live_portrait_webcam", Self::VidStudio => "vid_studio", + Self::VidStudioGen2 => "vid_studio_gen2", Self::VidStyleTransfer => "vid_style_transfer", Self::LipsyncFaceFusion => "lipsync_face_fusion", Self::LipsyncSadTalker => "lipsync_sad_talker", @@ -149,6 +153,7 @@ impl InferenceJobProductCategory { "vid_live_portrait" => Ok(Self::VidLivePortrait), "vid_live_portrait_webcam" => Ok(Self::VidLivePortraitWebcam), "vid_studio" => Ok(Self::VidStudio), + "vid_studio_gen2" => Ok(Self::VidStudioGen2), "vid_style_transfer" => Ok(Self::VidStyleTransfer), "lipsync_face_fusion" => Ok(Self::LipsyncFaceFusion), "lipsync_sad_talker" => Ok(Self::LipsyncSadTalker), @@ -179,6 +184,7 @@ impl InferenceJobProductCategory { Self::VidLivePortrait, Self::VidLivePortraitWebcam, Self::VidStudio, + Self::VidStudioGen2, Self::VidStyleTransfer, Self::LipsyncFaceFusion, Self::LipsyncSadTalker, @@ -215,6 +221,7 @@ mod tests { assert_serialization(InferenceJobProductCategory::VidLivePortrait, "vid_live_portrait"); assert_serialization(InferenceJobProductCategory::VidLivePortraitWebcam, "vid_live_portrait_webcam"); assert_serialization(InferenceJobProductCategory::VidStudio, "vid_studio"); + assert_serialization(InferenceJobProductCategory::VidStudioGen2, "vid_studio_gen2"); assert_serialization(InferenceJobProductCategory::VidStyleTransfer, "vid_style_transfer"); assert_serialization(InferenceJobProductCategory::LipsyncFaceFusion, "lipsync_face_fusion"); assert_serialization(InferenceJobProductCategory::LipsyncSadTalker, "lipsync_sad_talker"); @@ -240,6 +247,7 @@ mod tests { assert_eq!(InferenceJobProductCategory::VidLivePortrait.to_str(), "vid_live_portrait"); assert_eq!(InferenceJobProductCategory::VidLivePortraitWebcam.to_str(), "vid_live_portrait_webcam"); assert_eq!(InferenceJobProductCategory::VidStudio.to_str(), "vid_studio"); + assert_eq!(InferenceJobProductCategory::VidStudioGen2.to_str(), "vid_studio_gen2"); assert_eq!(InferenceJobProductCategory::VidStyleTransfer.to_str(), "vid_style_transfer"); assert_eq!(InferenceJobProductCategory::LipsyncFaceFusion.to_str(), "lipsync_face_fusion"); assert_eq!(InferenceJobProductCategory::LipsyncSadTalker.to_str(), "lipsync_sad_talker"); @@ -265,6 +273,7 @@ mod tests { assert_eq!(InferenceJobProductCategory::from_str("vid_live_portrait").unwrap(), InferenceJobProductCategory::VidLivePortrait); assert_eq!(InferenceJobProductCategory::from_str("vid_live_portrait_webcam").unwrap(), InferenceJobProductCategory::VidLivePortraitWebcam); assert_eq!(InferenceJobProductCategory::from_str("vid_studio").unwrap(), InferenceJobProductCategory::VidStudio); + assert_eq!(InferenceJobProductCategory::from_str("vid_studio_gen2").unwrap(), InferenceJobProductCategory::VidStudioGen2); assert_eq!(InferenceJobProductCategory::from_str("vid_style_transfer").unwrap(), InferenceJobProductCategory::VidStyleTransfer); assert_eq!(InferenceJobProductCategory::from_str("lipsync_face_fusion").unwrap(), InferenceJobProductCategory::LipsyncFaceFusion); assert_eq!(InferenceJobProductCategory::from_str("lipsync_sad_talker").unwrap(), InferenceJobProductCategory::LipsyncSadTalker); @@ -280,7 +289,7 @@ mod tests { fn all_variants() { // Static check let mut variants = InferenceJobProductCategory::all_variants(); - assert_eq!(variants.len(), 20); + assert_eq!(variants.len(), 21); assert_eq!(variants.pop_first(), Some(InferenceJobProductCategory::DownloadGptSoVits)); assert_eq!(variants.pop_first(), Some(InferenceJobProductCategory::TtsGptSoVits)); assert_eq!(variants.pop_first(), Some(InferenceJobProductCategory::TtsStyleTts2)); @@ -293,6 +302,7 @@ mod tests { assert_eq!(variants.pop_first(), Some(InferenceJobProductCategory::VidLivePortrait)); assert_eq!(variants.pop_first(), Some(InferenceJobProductCategory::VidLivePortraitWebcam)); assert_eq!(variants.pop_first(), Some(InferenceJobProductCategory::VidStudio)); + assert_eq!(variants.pop_first(), Some(InferenceJobProductCategory::VidStudioGen2)); assert_eq!(variants.pop_first(), Some(InferenceJobProductCategory::VidStyleTransfer)); assert_eq!(variants.pop_first(), Some(InferenceJobProductCategory::LipsyncFaceFusion)); assert_eq!(variants.pop_first(), Some(InferenceJobProductCategory::LipsyncSadTalker)); 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 f08e785883..76ad4e0ec9 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 @@ -36,6 +36,10 @@ pub enum InferenceJobType { #[deprecated(note = "Use VideoRender instead.")] ComfyUi, + /// Second gen studio + #[serde(rename = "studio_gen2")] + StudioGen2, + /// A job that turns "FBX" game engine files into "GLTF" files (Bevy-compatible). #[serde(rename = "convert_fbx_gltf")] ConvertFbxToGltf, @@ -103,6 +107,7 @@ impl InferenceJobType { Self::F5TTS => "f5_tts", Self::GptSovits => "gpt_sovits", Self::ComfyUi => "comfy_ui", + Self::StudioGen2 => "studio_gen2", Self::ConvertFbxToGltf => "convert_fbx_gltf", Self::MocapNet => "mocap_net", Self::RvcV2 => "rvc_v2", @@ -126,6 +131,7 @@ impl InferenceJobType { "f5_tts" => Ok(Self::F5TTS), "gpt_sovits" => Ok(Self::GptSovits), "comfy_ui" => Ok(Self::ComfyUi), + "studio_gen2" => Ok(Self::StudioGen2), "convert_fbx_gltf" => Ok(Self::ConvertFbxToGltf), "mocap_net" => Ok(Self::MocapNet), "rvc_v2" => Ok(Self::RvcV2), @@ -152,6 +158,7 @@ impl InferenceJobType { Self::F5TTS, Self::GptSovits, Self::ComfyUi, + Self::StudioGen2, Self::ConvertFbxToGltf, Self::MocapNet, Self::RvcV2, @@ -189,6 +196,7 @@ mod tests { assert_serialization(InferenceJobType::F5TTS, "f5_tts"); assert_serialization(InferenceJobType::GptSovits, "gpt_sovits"); assert_serialization(InferenceJobType::ComfyUi, "comfy_ui"); + assert_serialization(InferenceJobType::StudioGen2, "studio_gen2"); assert_serialization(InferenceJobType::ConvertFbxToGltf, "convert_fbx_gltf"); assert_serialization(InferenceJobType::MocapNet, "mocap_net"); assert_serialization(InferenceJobType::RvcV2, "rvc_v2"); @@ -211,6 +219,7 @@ mod tests { assert_eq!(InferenceJobType::F5TTS.to_str(), "f5_tts"); assert_eq!(InferenceJobType::GptSovits.to_str(), "gpt_sovits"); assert_eq!(InferenceJobType::ComfyUi.to_str(), "comfy_ui"); + assert_eq!(InferenceJobType::StudioGen2.to_str(), "studio_gen2"); assert_eq!(InferenceJobType::ConvertFbxToGltf.to_str(), "convert_fbx_gltf"); assert_eq!(InferenceJobType::MocapNet.to_str(), "mocap_net"); assert_eq!(InferenceJobType::RvcV2.to_str(), "rvc_v2"); @@ -233,6 +242,7 @@ mod tests { assert_eq!(InferenceJobType::from_str("f5_tts").unwrap(), InferenceJobType::F5TTS); assert_eq!(InferenceJobType::from_str("gpt_sovits").unwrap(), InferenceJobType::GptSovits); assert_eq!(InferenceJobType::from_str("comfy_ui").unwrap(), InferenceJobType::ComfyUi); + assert_eq!(InferenceJobType::from_str("studio_gen2").unwrap(), InferenceJobType::StudioGen2); assert_eq!(InferenceJobType::from_str("convert_fbx_gltf").unwrap(), InferenceJobType::ConvertFbxToGltf); assert_eq!(InferenceJobType::from_str("mocap_net").unwrap(), InferenceJobType::MocapNet); assert_eq!(InferenceJobType::from_str("rvc_v2").unwrap(), InferenceJobType::RvcV2); @@ -258,6 +268,7 @@ mod tests { assert_eq!(variants.pop_first(), Some(InferenceJobType::F5TTS)); assert_eq!(variants.pop_first(), Some(InferenceJobType::GptSovits)); assert_eq!(variants.pop_first(), Some(InferenceJobType::ComfyUi)); + assert_eq!(variants.pop_first(), Some(InferenceJobType::StudioGen2)); assert_eq!(variants.pop_first(), Some(InferenceJobType::ConvertFbxToGltf)); assert_eq!(variants.pop_first(), Some(InferenceJobType::MocapNet)); assert_eq!(variants.pop_first(), Some(InferenceJobType::RvcV2)); diff --git a/crates/_types/enums_public/src/by_table/generic_inference_jobs/public_inference_job_type.rs b/crates/_types/enums_public/src/by_table/generic_inference_jobs/public_inference_job_type.rs index 7dc537770a..c525c05253 100644 --- a/crates/_types/enums_public/src/by_table/generic_inference_jobs/public_inference_job_type.rs +++ b/crates/_types/enums_public/src/by_table/generic_inference_jobs/public_inference_job_type.rs @@ -39,6 +39,10 @@ pub enum PublicInferenceJobType { #[deprecated(note = "Use VideoRender instead.")] ComfyUi, + /// Second gen studio + #[serde(rename = "studio_gen2")] + StudioGen2, + /// A job that turns "FBX" game engine files into "GLTF" files (Bevy-compatible). #[serde(rename = "convert_fbx_gltf")] ConvertFbxToGltf, @@ -102,6 +106,7 @@ impl PublicInferenceJobType { InferenceJobType::GptSovits => Self::GptSovits, InferenceJobType::F5TTS => Self::F5TTS, InferenceJobType::ComfyUi => Self::ComfyUi, + InferenceJobType::StudioGen2 => Self::StudioGen2, InferenceJobType::ConvertFbxToGltf => Self::ConvertFbxToGltf, InferenceJobType::MocapNet => Self::MocapNet, InferenceJobType::RvcV2 => Self::RvcV2, @@ -127,6 +132,7 @@ impl PublicInferenceJobType { Self::GptSovits => InferenceJobType::GptSovits, Self::F5TTS => InferenceJobType::F5TTS, Self::ComfyUi => InferenceJobType::ComfyUi, + Self::StudioGen2 => InferenceJobType::StudioGen2, Self::ConvertFbxToGltf => InferenceJobType::ConvertFbxToGltf, Self::MocapNet => InferenceJobType::MocapNet, Self::RvcV2 => InferenceJobType::RvcV2, diff --git a/crates/lib/mysql_queries/src/payloads/generic_inference_args/generic_inference_args.rs b/crates/lib/mysql_queries/src/payloads/generic_inference_args/generic_inference_args.rs index 27a0a98bf3..ce822f98d1 100644 --- a/crates/lib/mysql_queries/src/payloads/generic_inference_args/generic_inference_args.rs +++ b/crates/lib/mysql_queries/src/payloads/generic_inference_args/generic_inference_args.rs @@ -7,6 +7,7 @@ use crate::payloads::generic_inference_args::inner_payloads::live_portrait_paylo use crate::payloads::generic_inference_args::inner_payloads::mocap_payload::MocapArgs; use crate::payloads::generic_inference_args::inner_payloads::render_engine_scene_to_video_payload::RenderEngineSceneToVideoArgs; use crate::payloads::generic_inference_args::inner_payloads::seed_vc_payload::SeedVcPayload; +use crate::payloads::generic_inference_args::inner_payloads::studio_gen2_payload::StudioGen2Payload; use crate::payloads::generic_inference_args::inner_payloads::tts_payload::TTSArgs; use crate::payloads::generic_inference_args::inner_payloads::videofilter_payload::RerenderArgs; use crate::payloads::generic_inference_args::inner_payloads::workflow_payload::WorkflowArgs; @@ -148,6 +149,9 @@ pub enum PolymorphicInferenceArgs { /// ComfyUI (Short name to save space when serializing.) Cu(WorkflowArgs), + /// Studio Gen2 (Short name to save space when serializing.) + S2(StudioGen2Payload), + /// ComfyUI (Short name to save space when serializing.) Lp(LivePortraitPayload), diff --git a/crates/lib/mysql_queries/src/payloads/generic_inference_args/inner_payloads/mod.rs b/crates/lib/mysql_queries/src/payloads/generic_inference_args/inner_payloads/mod.rs index 6fde89e11d..d7ced92c19 100644 --- a/crates/lib/mysql_queries/src/payloads/generic_inference_args/inner_payloads/mod.rs +++ b/crates/lib/mysql_queries/src/payloads/generic_inference_args/inner_payloads/mod.rs @@ -5,6 +5,7 @@ pub mod lipsync_payload; pub mod live_portrait_payload; pub mod mocap_payload; pub mod render_engine_scene_to_video_payload; +pub mod studio_gen2_payload; pub mod tts_payload; pub mod videofilter_payload; pub mod workflow_payload; diff --git a/crates/lib/mysql_queries/src/payloads/generic_inference_args/inner_payloads/studio_gen2_payload.rs b/crates/lib/mysql_queries/src/payloads/generic_inference_args/inner_payloads/studio_gen2_payload.rs new file mode 100644 index 0000000000..eb4b878645 --- /dev/null +++ b/crates/lib/mysql_queries/src/payloads/generic_inference_args/inner_payloads/studio_gen2_payload.rs @@ -0,0 +1,24 @@ +use enums::common::visibility::Visibility; +use crate::payloads::generic_inference_args::common::watermark_type::WatermarkType; +use tokens::tokens::media_files::MediaFileToken; + +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)] +pub struct StudioGen2Payload { + /// The input image media file (required) + #[serde(rename = "i")] + #[serde(skip_serializing_if = "Option::is_none")] + pub image_file: Option, + + /// The input video media file (required) + #[serde(rename = "v")] + #[serde(skip_serializing_if = "Option::is_none")] + pub video_file: Option, + + #[serde(rename = "wt")] + #[serde(skip_serializing_if = "Option::is_none")] + pub watermark_type: Option, + + #[serde(rename = "cv")] + #[serde(skip_serializing_if = "Option::is_none")] + pub creator_visibility: Option, +} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/mod.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/mod.rs index e90d9b068b..9300769c41 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/mod.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/mod.rs @@ -25,13 +25,14 @@ pub mod prompts; pub mod service; pub mod stats; pub mod stubs; +pub mod studio_gen2; pub mod subscriptions; pub mod tags; pub mod trending; pub mod tts; +pub mod users; pub mod user_bookmarks; pub mod user_ratings; -pub mod users; pub mod vocoders; pub mod voice_clone_requests; pub mod voice_conversion; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/studio_gen2/enqueue_studio_gen2_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/studio_gen2/enqueue_studio_gen2_handler.rs new file mode 100644 index 0000000000..b90806c73c --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/studio_gen2/enqueue_studio_gen2_handler.rs @@ -0,0 +1,295 @@ +// NB: Incrementally getting rid of build warnings... +// #![forbid(unused_imports)] +// #![forbid(unused_mut)] +// #![forbid(unused_variables)] + +use std::sync::Arc; + +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::Json; +use actix_web::{web, HttpRequest, HttpResponse}; +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::common::watermark_type::WatermarkType; +use mysql_queries::payloads::generic_inference_args::generic_inference_args::{GenericInferenceArgs, InferenceCategoryAbbreviated, PolymorphicInferenceArgs}; +use mysql_queries::payloads::generic_inference_args::inner_payloads::studio_gen2_payload::StudioGen2Payload; +use mysql_queries::payloads::generic_inference_args::inner_payloads::workflow_payload::{WorkflowArgs, WorkflowType}; +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::enqueue::vst_common::vst_request::VstRequest; +use crate::http_server::endpoints::workflows::enqueue::vst_common::vst_response::VstSuccessResponse; +use crate::http_server::requests::request_headers::get_routing_tag_header::get_routing_tag_header; +use crate::http_server::requests::request_headers::has_debug_header::has_debug_header; +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; +use crate::util::cleaners::empty_media_file_token_to_null::empty_media_file_token_to_null; + +#[derive(Deserialize, ToSchema)] +pub struct EnqueueStudioGen2Request { + /// Entropy for request de-duplication (required) + pub uuid_idempotency_token: String, + + /// The input image media file (required) + pub image_file: MediaFileToken, + + /// The input video media file (required) + pub video_file: MediaFileToken, + + /// Remove watermark from the output + /// Only for premium accounts + pub remove_watermark: Option, + + /// Optional visibility setting override. + pub creator_set_visibility: Option, +} + +#[derive(Serialize, ToSchema)] +pub struct EnqueueStudioGen2Response { + pub success: bool, + pub inference_job_token: InferenceJobToken, +} + +#[derive(Debug, ToSchema)] +pub enum EnqueueStudioGen2Error { + BadInput(String), + NotAuthorized, + ServerError, + RateLimited, +} + +impl ResponseError for EnqueueStudioGen2Error { + fn status_code(&self) -> StatusCode { + match *self { + EnqueueStudioGen2Error::BadInput(_) => StatusCode::BAD_REQUEST, + EnqueueStudioGen2Error::NotAuthorized => StatusCode::UNAUTHORIZED, + EnqueueStudioGen2Error::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + EnqueueStudioGen2Error::RateLimited => StatusCode::TOO_MANY_REQUESTS, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + EnqueueStudioGen2Error::BadInput(reason) => reason.to_string(), + EnqueueStudioGen2Error::NotAuthorized => "unauthorized".to_string(), + EnqueueStudioGen2Error::ServerError => "server error".to_string(), + EnqueueStudioGen2Error::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 EnqueueStudioGen2Error { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{:?}", self) + } +} + +/// Enqueue Studio Gen2 jobs. +#[utoipa::path( + post, + tag = "Studio Gen2", + path = "/v1/studio_gen2/enqueue", + responses( + (status = 200, description = "Success", body = VstSuccessResponse), + (status = 400, description = "Bad input", body = EnqueueStudioGen2Error), + (status = 401, description = "Not authorized", body = EnqueueStudioGen2Error), + (status = 429, description = "Rate limited", body = EnqueueStudioGen2Error), + (status = 500, description = "Server error", body = EnqueueStudioGen2Error) + ), + params(("request" = VstRequest, description = "Payload for request")) +)] +pub async fn enqueue_studio_gen2_handler( + http_request: HttpRequest, + request: Json, + server_state: web::Data> +) -> Result, EnqueueStudioGen2Error> { + + let a = 1; + // ==================== DB ==================== // + + let mut mysql_connection = server_state.mysql_pool + .acquire() + .await + .map_err(|err| { + warn!("MySql pool error: {:?}", err); + EnqueueStudioGen2Error::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); + EnqueueStudioGen2Error::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(EnqueueStudioGen2Error::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 = has_debug_header(&http_request); + + let maybe_routing_tag= get_routing_tag_header(&http_request); + + // ==================== 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(EnqueueStudioGen2Error::NotAuthorized); + } + &server_state.redis_rate_limiters.logged_in + }, + }; + + if let Err(_err) = rate_limiter.rate_limit_request(&http_request) { + return Err(EnqueueStudioGen2Error::RateLimited); + } + + // ==================== HANDLE IDEMPOTENCY ==================== // + + if let Err(reason) = validate_idempotency_token_format(&request.uuid_idempotency_token) { + return Err(EnqueueStudioGen2Error::BadInput(reason)); + } + + insert_idempotency_token(&request.uuid_idempotency_token, &mut *mysql_connection) + .await + .map_err(|err| { + error!("Error inserting idempotency token: {:?}", err); + EnqueueStudioGen2Error::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 has_paid_plan = plan.plan_slug() == "fakeyou_contributor" || plan.plan_category() == PlanCategory::Paid; + + let watermark_type ; + + let is_requested_to_remove_watermark = request.remove_watermark.unwrap_or(false); + let is_allowed_to_remove_watermark = is_staff; // TODO: Paid permission instead + + if is_requested_to_remove_watermark && is_allowed_to_remove_watermark { + watermark_type = None; + } else { + watermark_type = Some(WatermarkType::Storyteller); // Studio is only on Storyteller. + } + + let inference_args = StudioGen2Payload { + image_file: Some(request.image_file.clone()), + video_file: Some(request.video_file.clone()), + creator_visibility: Some(set_visibility), + watermark_type, + }; + + info!("Creating ComfyUI job record..."); + + let query_result = insert_generic_inference_job(InsertGenericInferenceArgs { + uuid_idempotency_token: &request.uuid_idempotency_token, + job_type: InferenceJobType::StudioGen2, + inference_category: InferenceCategory::DeprecatedField, + maybe_model_type: None, + maybe_product_category: None, + maybe_model_token: None, + 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::S2(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); + if err.had_duplicate_idempotency_token() { + return Err(EnqueueStudioGen2Error::BadInput("Duplicate idempotency token".to_string())); + } + return Err(EnqueueStudioGen2Error::ServerError); + } + }; + + Ok(Json(EnqueueStudioGen2Response { + success: true, + inference_job_token: job_token, + })) +} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/studio_gen2/mod.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/studio_gen2/mod.rs new file mode 100644 index 0000000000..0ab15cfd67 --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/studio_gen2/mod.rs @@ -0,0 +1 @@ +pub mod enqueue_studio_gen2_handler; \ No newline at end of file diff --git a/crates/service/web/storyteller_web/src/http_server/routes/add_routes.rs b/crates/service/web/storyteller_web/src/http_server/routes/add_routes.rs index 84b2eeccde..9242c46e17 100644 --- a/crates/service/web/storyteller_web/src/http_server/routes/add_routes.rs +++ b/crates/service/web/storyteller_web/src/http_server/routes/add_routes.rs @@ -123,6 +123,7 @@ use crate::http_server::endpoints::w2l::get_w2l_template_use_count::get_w2l_temp use crate::http_server::endpoints::w2l::get_w2l_upload_template_job_status::get_w2l_upload_template_job_status_handler; use crate::http_server::endpoints::w2l::list_w2l_templates::list_w2l_templates_handler; use crate::http_server::endpoints::w2l::set_w2l_template_mod_approval::set_w2l_template_mod_approval_handler; +use crate::http_server::routes::add_studio_gen2_routes::add_studio_gen2_routes; use crate::http_server::routes::beta_key_routes::add_beta_key_routes; use crate::http_server::routes::job_routes::add_job_routes; use crate::http_server::routes::media_files_routes::add_media_file_routes; @@ -169,6 +170,7 @@ pub fn add_routes (app: App, server_environment: ServerEnvironment) -> app = add_tag_routes(app); /* /v1/tags */ app = add_model_download_routes(app); app = add_workflow_routes(app); + app = add_studio_gen2_routes(app); app = add_job_routes(app); app = add_engine_routes(app); /* /v1/engine/... */ diff --git a/crates/service/web/storyteller_web/src/http_server/routes/add_studio_gen2_routes.rs b/crates/service/web/storyteller_web/src/http_server/routes/add_studio_gen2_routes.rs new file mode 100644 index 0000000000..3d11fd93f1 --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/routes/add_studio_gen2_routes.rs @@ -0,0 +1,24 @@ +use actix_http::body::MessageBody; +use actix_service::ServiceFactory; +use actix_web::{App, Error, HttpResponse, web}; +use actix_web::dev::{ServiceRequest, ServiceResponse}; +use crate::http_server::endpoints::studio_gen2::enqueue_studio_gen2_handler::enqueue_studio_gen2_handler; + +pub fn add_studio_gen2_routes (app: App) -> App +where + B: MessageBody, + T: ServiceFactory< + ServiceRequest, + Config = (), + Response = ServiceResponse, + Error = Error, + InitError = (), + >, +{ + app.service(web::scope("/v1/studio_gen2") + .service(web::resource("/enqueue") + .route(web::post().to(enqueue_studio_gen2_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + ) +} diff --git a/crates/service/web/storyteller_web/src/http_server/routes/mod.rs b/crates/service/web/storyteller_web/src/http_server/routes/mod.rs index 60d6221669..ffeeaab40a 100644 --- a/crates/service/web/storyteller_web/src/http_server/routes/mod.rs +++ b/crates/service/web/storyteller_web/src/http_server/routes/mod.rs @@ -1,4 +1,5 @@ pub mod add_routes; +pub mod add_studio_gen2_routes; pub mod beta_key_routes; pub mod job_routes; pub mod media_files_routes;