From 733c24aebb3e0517fcd72bb1c3c8a76cfa3e15b3 Mon Sep 17 00:00:00 2001 From: ArEnSc Date: Tue, 2 Jan 2024 10:50:24 -0500 Subject: [PATCH] Working on enqueuing image generation --- .../generic_inference_args.rs | 7 + .../image_generation_payload.rs | 54 ++++ .../payloads/generic_inference_args/mod.rs | 3 +- .../image_gen/enqueue_image_generation.rs | 232 ++++++++++++++++++ .../http_server/endpoints/image_gen/mod.rs | 1 + .../src/http_server/endpoints/mod.rs | 3 +- .../service/web/storyteller_web/src/routes.rs | 27 +- 7 files changed, 324 insertions(+), 3 deletions(-) create mode 100644 crates/lib/mysql_queries/src/payloads/generic_inference_args/image_generation_payload.rs create mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/enqueue_image_generation.rs create mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/mod.rs 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 a67dd0bd91..49da8fb98c 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 @@ -4,6 +4,7 @@ use errors::AnyhowResult; use crate::payloads::generic_inference_args::lipsync_payload::LipsyncArgs; use crate::payloads::generic_inference_args::tts_payload::TTSArgs; use crate::payloads::generic_inference_args::videofilter_payload::{RerenderArgs}; +use crate::payloads::generic_inference_args::image_generation_payload::{SDArgs}; /// Used to encode extra state for the `generic_inference_jobs` table. @@ -40,6 +41,10 @@ pub enum InferenceCategoryAbbreviated { #[serde(rename = "vf")] // NB: DO NOT CHANGE. It could break live jobs. Renamed to be fewer bytes. #[serde(alias = "video_filter")] VideoFilter, + + #[serde(rename = "ig")] // NB: DO NOT CHANGE. It could break live jobs. Renamed to be fewer bytes. + #[serde(alias = "image_generation")] + ImageGeneration } #[derive(Serialize, Deserialize, Clone, Copy, Debug)] @@ -113,6 +118,7 @@ impl InferenceCategoryAbbreviated { InferenceCategory::TextToSpeech => Self::TextToSpeech, InferenceCategory::VoiceConversion => Self::VoiceConversion, InferenceCategory::VideoFilter => Self::VideoFilter, + InferenceCategory::ImageGeneration => Self::ImageGeneration, } } @@ -122,6 +128,7 @@ impl InferenceCategoryAbbreviated { Self::TextToSpeech => InferenceCategory::TextToSpeech, Self::VoiceConversion => InferenceCategory::VoiceConversion, Self::VideoFilter => InferenceCategory::VideoFilter, + Self::ImageGeneration =>InferenceCategory::ImageGeneration } } } diff --git a/crates/lib/mysql_queries/src/payloads/generic_inference_args/image_generation_payload.rs b/crates/lib/mysql_queries/src/payloads/generic_inference_args/image_generation_payload.rs new file mode 100644 index 0000000000..47c1efa11b --- /dev/null +++ b/crates/lib/mysql_queries/src/payloads/generic_inference_args/image_generation_payload.rs @@ -0,0 +1,54 @@ +use tokens::tokens::model_weights::ModelWeightToken; + +/// Video sources can be one of several: +/// - F: media_files (todo) +/// - U: media_uploads (legacy) + +/// For image to image we probably want images +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] +pub enum SDSource { + // NB: DO NOT CHANGE. It could break live jobs. Renamed to be fewer bytes. + /// Media File Token (media_files table) + /// Serde cannot yet rename enum variants. + F(String), + + // NB: DO NOT CHANGE. It could break live jobs. Renamed to be fewer bytes. + /// Media Upload Token (media_uploads table) + /// Serde cannot yet rename enum variants. + U(String), +} + +#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq, Eq)] +pub struct SDArgs { + #[serde(rename = "vs")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_video_source: Option, + + #[serde(rename = "sd")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_sd_model_token: Option, + + #[serde(rename = "lora")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_lora_model_token: Option, + + #[serde(rename = "p")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_prompt: Option, + + #[serde(rename = "ap")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_a_prompt: Option, + + #[serde(rename = "np")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_n_prompt: Option, + + #[serde(rename = "se")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_seed: Option, + + #[serde(rename = "upload_path")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_upload_path: Option +} diff --git a/crates/lib/mysql_queries/src/payloads/generic_inference_args/mod.rs b/crates/lib/mysql_queries/src/payloads/generic_inference_args/mod.rs index e68fc5a71c..17b564fc16 100644 --- a/crates/lib/mysql_queries/src/payloads/generic_inference_args/mod.rs +++ b/crates/lib/mysql_queries/src/payloads/generic_inference_args/mod.rs @@ -1,4 +1,5 @@ pub mod generic_inference_args; pub mod lipsync_payload; pub mod tts_payload; -pub mod videofilter_payload; \ No newline at end of file +pub mod videofilter_payload; +mod image_generation_payload; \ No newline at end of file diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/enqueue_image_generation.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/enqueue_image_generation.rs new file mode 100644 index 0000000000..c0655f3071 --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/enqueue_image_generation.rs @@ -0,0 +1,232 @@ +#![forbid(unused_imports)] +#![forbid(unused_mut)] +#![forbid(unused_variables)] + +use std::fmt::Debug; +use std::sync::Arc; + +use actix_web::{HttpRequest, HttpResponse, web}; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use log::warn; +use serde::Deserialize; +use serde::Serialize; +use utoipa::ToSchema; + +use enums::by_table::generic_inference_jobs::inference_category::InferenceCategory; +use enums::by_table::generic_inference_jobs::inference_model_type::InferenceModelType; +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::tts_payload::TTSArgs; +use mysql_queries::queries::generic_inference::web::insert_generic_inference_job::{ + insert_generic_inference_job, + InsertGenericInferenceArgs, +}; + +use tokens::tokens::generic_inference_jobs::InferenceJobToken; +use tokens::tokens::users::UserToken; + +use crate::configs::plans::get_correct_plan_for_session::get_correct_plan_for_session; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::server_state::ServerState; + +/// 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"; + + +#[derive(Deserialize,ToSchema)] +pub struct EnqueueImageGenRequest { + uuid_idempotency_token: String, + text: String, + voice_token: String, +} + +#[derive(Serialize,ToSchema)] +pub struct EnqueueImageGenRequestSuccessResponse { + pub success: bool, + pub inference_job_token: InferenceJobToken, +} + +#[derive(Debug,ToSchema)] +pub enum EnqueueImageGenRequestError { + BadInput(String), + NotAuthorized, + ServerError, + RateLimited, +} + +impl ResponseError for EnqueueImageGenRequestError { + fn status_code(&self) -> StatusCode { + match *self { + EnqueueImageGenRequestError::BadInput(_) => StatusCode::BAD_REQUEST, + EnqueueImageGenRequestError::NotAuthorized => StatusCode::UNAUTHORIZED, + EnqueueImageGenRequestError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + EnqueueImageGenRequestError::RateLimited => StatusCode::TOO_MANY_REQUESTS, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + EnqueueImageGenRequestError::BadInput(reason) => reason.to_string(), + EnqueueImageGenRequestError::NotAuthorized => "unauthorized".to_string(), + EnqueueImageGenRequestError::ServerError => "server error".to_string(), + EnqueueImageGenRequestError::RateLimited => "rate limited".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +impl std::fmt::Display for EnqueueImageGenRequestError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{:?}", self) + } +} + +// Implementation for enqueuing a TTS request +// Reference enqueue_infer_tts_handler.rs for checks: rate limiting / user sessions +// insert generic inference job.rs +// Need to convert it to generic inference job. +#[utoipa::path( +post, +path = "/inference/enqueue_image_gen/", +responses( +(status = 200, description = "Enqueue TTS generically", body = EnqueueImageGenRequestSuccessResponse), +(status = 400, description = "Bad input", body = EnqueueImageGenRequestError), +(status = 401, description = "Not authorized", body = EnqueueImageGenRequestError), +(status = 429, description = "Rate limited", body = EnqueueImageGenRequestError), +(status = 500, description = "Server error", body = EnqueueImageGenRequestError) +), +params( +("request" = EnqueueImageGenRequest, description = "Payload for TTS Request") +) +)] +pub async fn enqueue_image_generation_request( + http_request: HttpRequest, + request: web::Json, + server_state: web::Data> +) -> Result { + + let mut maybe_user_token: Option = None; + + let mut mysql_connection = server_state.mysql_pool.acquire().await.map_err(|err| { + warn!("MySql pool error: {:?}", err); + EnqueueImageGenRequestError::ServerError + })?; + + // ==================== 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); + EnqueueImageGenRequestError::ServerError + })?; + + if let Some(user_session) = maybe_user_session.as_ref() { + maybe_user_token = Some(UserToken::new_from_str(&user_session.user_token)); + } + + // 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() + ); + + // Separate priority for animation. + let priority_level = plan.web_vc_base_priority_level(); + + // ==================== 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() + ); + + // ==================== BANNED USERS ==================== // + + if let Some(ref user) = maybe_user_session { + if user.role.is_banned { + return Err(EnqueueImageGenRequestError::NotAuthorized); + } + } + + // ==================== RATE LIMIT ==================== // + + let rate_limiter = match maybe_user_session { + None => &server_state.redis_rate_limiters.logged_out, + Some(ref _user) => &server_state.redis_rate_limiters.logged_in, + }; + + if let Err(_err) = rate_limiter.rate_limit_request(&http_request) { + return Err(EnqueueImageGenRequestError::RateLimited); + } + + // Get up IP address + let ip_address = get_request_ip(&http_request); + + // package as larger component args should always have an embedding token .. + let inference_args = TTSArgs { + voice_token: Some(request.voice_token.clone()), + dataset_token: None, + }; + + // create the inference args here + // enqueue a zero shot tts request here... + let maybe_avt_token = server_state.avt_cookie_manager.get_avt_token_from_request(&http_request); + + // create the job record here! + let query_result = insert_generic_inference_job(InsertGenericInferenceArgs { + uuid_idempotency_token: &request.uuid_idempotency_token, + inference_category: InferenceCategory::TextToSpeech, + maybe_model_type: Some(InferenceModelType::VallEX), // NB: Model is static during inference + maybe_model_token: None, // NB: Model is static during inference + maybe_input_source_token: None, + maybe_input_source_token_type: None, + maybe_raw_inference_text: Some(&request.text), + maybe_max_duration_seconds: None, + maybe_inference_args: Some(GenericInferenceArgs { + inference_category: Some(InferenceCategoryAbbreviated::TextToSpeech), + args: Some(PolymorphicInferenceArgs::Tts(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: enums::common::visibility::Visibility::Public, + priority_level, + requires_keepalive: true, // do we need this? I think so + 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(EnqueueImageGenRequestError::ServerError); + } + }; + + let response: EnqueueImageGenRequestSuccessResponse = EnqueueImageGenRequestSuccessResponse { + success: true, + inference_job_token: job_token, + }; + + let body = serde_json::to_string(&response).map_err(|_e| EnqueueImageGenRequestError::ServerError)?; + + // Error handling 101 rust result type returned like so. + Ok(HttpResponse::Ok().content_type("application/json").body(body)) +} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/mod.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/mod.rs new file mode 100644 index 0000000000..fd510c2b58 --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/mod.rs @@ -0,0 +1 @@ +pub mod enqueue_image_generation; 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 0a6ed9375d..5b92cab8d4 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 @@ -26,4 +26,5 @@ pub mod voice_conversion; pub mod voice_designer; pub mod w2l; pub mod user_bookmarks; -pub mod weights; \ No newline at end of file +pub mod weights; +pub mod image_gen; \ No newline at end of file diff --git a/crates/service/web/storyteller_web/src/routes.rs b/crates/service/web/storyteller_web/src/routes.rs index 6e85b5e805..ff66d62d46 100644 --- a/crates/service/web/storyteller_web/src/routes.rs +++ b/crates/service/web/storyteller_web/src/routes.rs @@ -162,7 +162,7 @@ use crate::http_server::endpoints::weights::list_featured_weights::list_featured use crate::http_server::endpoints::weights::list_weights_by_user::list_weights_by_user_handler; use crate::http_server::endpoints::weights::set_model_weight_cover_image::set_model_weight_cover_image_handler; use crate::http_server::endpoints::weights::update_weight::update_weight_handler; - +use crate::http_server::endpoints::image_gen::enqueue_image_generation:: pub fn add_routes (app: App, server_environment: ServerEnvironment) -> App where B: MessageBody, @@ -1257,6 +1257,31 @@ fn add_voice_designer_routes (app:App)-> App ) } + +fn add_image_gen_routes (app:App)-> App + where + B: MessageBody, + T: ServiceFactory< + ServiceRequest, + Config = (), + Response = ServiceResponse, + Error = Error, + InitError = (), + >, +{ + app.service( + web::scope("/v1/image_gen") + .service( + web::scope("/model") + .route("/upload", web::post().to(upload_zs_sample_handler)) + ) + .service( + web::scope("/inference") + .route("/enqueue_image_gen", web::post().to(enqueue_tts_request)) + ) + ) +} + // ==================== Weights ROUTES ==================== fn add_weights_routes(app: App) -> App where