Working on enqueuing image generation

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

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<VideofilterVideoSource>,
+
+    #[serde(rename = "sd")]
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub maybe_sd_model_token: Option<ModelWeightToken>,
+
+    #[serde(rename = "lora")]
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub maybe_lora_model_token: Option<ModelWeightToken>,
+
+    #[serde(rename = "p")]
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub maybe_prompt: Option<String>,
+
+    #[serde(rename = "ap")]
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub maybe_a_prompt: Option<String>,
+
+    #[serde(rename = "np")]
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub maybe_n_prompt: Option<String>,
+
+    #[serde(rename = "se")]
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub maybe_seed: Option<i32>,
+
+    #[serde(rename = "upload_path")]
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub maybe_upload_path: Option<String>
+}

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;
+pub mod videofilter_payload;
+mod image_generation_payload;

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<EnqueueImageGenRequest>,
+    server_state: web::Data<Arc<ServerState>>
+) -> Result<HttpResponse, EnqueueImageGenRequestError> {
+
+    let mut maybe_user_token: Option<UserToken> = 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))
+}

crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/mod.rs

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

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;
+pub mod weights;
+pub mod image_gen;

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<T, B> (app: App<T>, server_environment: ServerEnvironment) -> App<T>
   where
       B: MessageBody,
@@ -1257,6 +1257,31 @@ fn add_voice_designer_routes<T,B> (app:App<T>)-> App<T>
       )
 }
 
+
+fn add_image_gen_routes<T,B> (app:App<T>)-> App<T>
+    where
+        B: MessageBody,
+        T: ServiceFactory<
+            ServiceRequest,
+            Config = (),
+            Response = ServiceResponse<B>,
+            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<T, B>(app: App<T>) -> App<T>
     where