Merge pull request #286 from storytold/brandon-port-tts-queries

port TTS enqueue to weights

crates/lib/migration/src/text_to_speech/get_tts_model_for_inference_migration.rs

@@ -0,0 +1,74 @@
+use serde_derive::{Deserialize, Serialize};
+use sqlx::MySql;
+use sqlx::pool::PoolConnection;
+
+use enums::common::visibility::Visibility;
+use errors::AnyhowResult;
+use mysql_queries::queries::model_weights::get::get_weight_for_legacy_tts_enqueue::{get_weight_for_legacy_tts_enqueue_with_connection, ModelWeightForLegacyTtsInference};
+use mysql_queries::queries::tts::tts_models::get_tts_model::{get_tts_model_by_token_using_connection, TtsModelRecord};
+use tokens::tokens::model_weights::ModelWeightToken;
+
+/// Get TTS model
+/// This is for the tts inference page.
+pub async fn get_tts_model_for_inference_migration(
+  token: &str,
+  mysql_connection: &mut PoolConnection<MySql>,
+  can_see_deleted: bool,
+  use_weights_table: bool,
+) -> AnyhowResult<Option<TtsModelForInferenceMigration>> {
+  // NB: This is temporary migration code as we switch from the `tts_models` table to the `model_weights` table.
+  if use_weights_table {
+    let token = ModelWeightToken::new_from_str(token);
+
+    let maybe_model = get_weight_for_legacy_tts_enqueue_with_connection(
+      &token,
+      can_see_deleted,
+      mysql_connection
+    ).await?;
+
+    Ok(maybe_model.map(|model| TtsModelForInferenceMigration::ModelWeight(model)))
+
+  } else {
+
+    let maybe_model = get_tts_model_by_token_using_connection(
+      &token,
+      true,
+      mysql_connection
+    ).await?;
+
+    Ok(maybe_model.map(|model| TtsModelForInferenceMigration::LegacyTts(model)))
+  }
+}
+
+/// Union over the legacy table and the new table to support an easier migration.
+/// This enum can hold a record of either type and present a unified accessor interface.
+#[derive(Clone, Serialize, Deserialize)]
+pub enum TtsModelForInferenceMigration {
+  /// Old type from the `tts_models` table, on the way out
+  LegacyTts(TtsModelRecord),
+  /// New type, replacing the `tts_models` table.
+  ModelWeight(ModelWeightForLegacyTtsInference),
+}
+
+impl TtsModelForInferenceMigration {
+  pub fn token(&self) -> &str {
+    match self {
+      Self::LegacyTts(ref model) => model.model_token.as_str(),
+      Self::ModelWeight(ref model) => model.token.as_str(),
+    }
+  }
+
+  pub fn creator_user_token(&self) -> &str {
+    match self {
+      Self::LegacyTts(ref model) => &model.creator_user_token,
+      Self::ModelWeight(ref model) => model.creator_user_token.as_str(),
+    }
+  }
+
+  pub fn creator_set_visibility(&self) -> Visibility{
+    match self {
+      Self::LegacyTts(ref model) => model.creator_set_visibility,
+      Self::ModelWeight(ref model) => model.creator_set_visibility,
+    }
+  }
+}

crates/lib/migration/src/text_to_speech/mod.rs

@@ -1 +1,2 @@
-pub mod list_tts_models_for_migration;
+pub mod list_tts_models_for_migration;
+pub mod get_tts_model_for_inference_migration;

crates/lib/mysql_queries/src/queries/model_weights/get/get_weight_for_legacy_tts_enqueue.rs

@@ -0,0 +1,142 @@
+use anyhow::anyhow;
+use chrono::{DateTime, Utc};
+use log::error;
+use sqlx::{MySql, MySqlPool};
+use sqlx::pool::PoolConnection;
+
+use container_common::anyhow_result::AnyhowResult;
+use enums::by_table::model_weights::weights_category::WeightsCategory;
+use enums::by_table::model_weights::weights_types::WeightsType;
+use enums::common::visibility::Visibility;
+use tokens::tokens::{model_weights::ModelWeightToken, users::UserToken};
+
+#[derive(Clone, Serialize, Deserialize)]
+pub struct ModelWeightForLegacyTtsInference {
+    pub token: ModelWeightToken,
+    pub weights_type: WeightsType,
+    pub weights_category: WeightsCategory,
+
+    pub creator_user_token: UserToken,
+    pub creator_set_visibility: Visibility,
+
+    pub user_deleted_at: Option<DateTime<Utc>>,
+    pub mod_deleted_at: Option<DateTime<Utc>>,
+}
+
+pub async fn get_weight_for_legacy_tts_enqueue(
+    weight_token: &ModelWeightToken,
+    can_see_deleted: bool,
+    mysql_pool: &MySqlPool
+) -> AnyhowResult<Option<ModelWeightForLegacyTtsInference>> {
+    let mut connection = mysql_pool.acquire().await?;
+    get_weight_for_legacy_tts_enqueue_with_connection(weight_token, can_see_deleted, &mut connection).await
+}
+
+pub async fn get_weight_for_legacy_tts_enqueue_with_connection(
+    weight_token: &ModelWeightToken,
+    can_see_deleted: bool,
+    mysql_connection: &mut PoolConnection<MySql>
+) -> AnyhowResult<Option<ModelWeightForLegacyTtsInference>> {
+    let maybe_result = if can_see_deleted {
+        select_include_deleted(weight_token, mysql_connection).await
+    } else {
+        select_without_deleted(weight_token, mysql_connection).await
+    };
+
+    let record: RawWeight = match maybe_result {
+        Ok(record) => record,
+        Err(sqlx::Error::RowNotFound) => {
+            return Ok(None);
+        }
+        Err(err) => {
+            error!("Error fetching weights by token: {:?}", err);
+            return Err(anyhow!("Error fetching weights by token: {:?}", err));
+        }
+    };
+
+    // unwrap the result
+
+    Ok(
+        Some(ModelWeightForLegacyTtsInference {
+            token: record.token,
+            weights_type: record.weights_type,
+            weights_category: record.weights_category,
+            creator_user_token: record.creator_user_token,
+            creator_set_visibility: record.creator_set_visibility,
+            user_deleted_at: record.user_deleted_at,
+            mod_deleted_at: record.mod_deleted_at,
+        })
+    )
+}
+
+async fn select_include_deleted(
+    weight_token: &ModelWeightToken,
+    mysql_connection: &mut PoolConnection<MySql>
+) -> Result<RawWeight, sqlx::Error> {
+    sqlx
+        ::query_as!(
+            RawWeight,
+            r#"
+        SELECT
+        wt.token as `token: tokens::tokens::model_weights::ModelWeightToken`,
+        wt.weights_type as `weights_type: enums::by_table::model_weights::weights_types::WeightsType`,
+        wt.weights_category as `weights_category: enums::by_table::model_weights::weights_category::WeightsCategory`,
+
+        wt.creator_user_token as `creator_user_token: tokens::tokens::users::UserToken`,
+        wt.creator_set_visibility as `creator_set_visibility: enums::common::visibility::Visibility`,
+
+        wt.user_deleted_at,
+        wt.mod_deleted_at
+
+        FROM model_weights as wt
+        WHERE
+            wt.token = ?
+            "#,
+            weight_token.as_str()
+        )
+        .fetch_one(&mut **mysql_connection).await
+}
+
+async fn select_without_deleted(
+    weight_token: &ModelWeightToken,
+    mysql_connection: &mut PoolConnection<MySql>
+) -> Result<RawWeight, sqlx::Error> {
+    sqlx
+        ::query_as!(
+            RawWeight,
+            r#"
+        SELECT
+        wt.token as `token: tokens::tokens::model_weights::ModelWeightToken`,
+        wt.weights_type as `weights_type: enums::by_table::model_weights::weights_types::WeightsType`,
+        wt.weights_category as `weights_category: enums::by_table::model_weights::weights_category::WeightsCategory`,
+
+        wt.creator_user_token as `creator_user_token: tokens::tokens::users::UserToken`,
+        wt.creator_set_visibility as `creator_set_visibility: enums::common::visibility::Visibility`,
+
+        wt.user_deleted_at,
+        wt.mod_deleted_at
+
+        FROM model_weights as wt
+        WHERE
+            wt.token = ?
+            AND wt.user_deleted_at IS NULL
+            AND wt.mod_deleted_at IS NULL
+        "#,
+            weight_token.as_str()
+        )
+        .fetch_one(&mut **mysql_connection).await
+}
+
+// RawWeight is the struct that is returned from the database in raw form.
+#[derive(Serialize)]
+struct RawWeight {
+    pub token: ModelWeightToken,
+    pub weights_type: WeightsType,
+    pub weights_category: WeightsCategory,
+
+    pub creator_user_token: UserToken,
+    pub creator_set_visibility: Visibility,
+
+    pub user_deleted_at: Option<DateTime<Utc>>,
+    pub mod_deleted_at: Option<DateTime<Utc>>,
+}

crates/lib/mysql_queries/src/queries/model_weights/get/mod.rs

@@ -0,0 +1,2 @@
+pub mod get_weight;
+pub mod get_weight_for_legacy_tts_enqueue;

crates/lib/mysql_queries/src/queries/model_weights/mod.rs

@@ -1,7 +1,7 @@
 pub mod create;
 pub mod delete_weights;
 pub mod edit;
-pub mod get_weight;
+pub mod get;
 pub mod inference;
 pub mod list;
 pub mod migration;

crates/lib/mysql_queries/src/queries/model_weights/tests.rs

@@ -27,7 +27,7 @@ mod tests {
         undelete_weights_as_mod,
         undelete_weights_as_user,
     };
-    use crate::queries::model_weights::get_weight::get_weight_by_token;
+    use crate::queries::model_weights::get::get_weight::get_weight_by_token;
     use crate::queries::model_weights::list::list_weights_by_user::{list_weights_by_creator_username, ListWeightsForUserArgs};
     use crate::queries::model_weights::list::list_weights_query_builder::ListWeightsQueryBuilder;
     use crate::queries::users::user::get_user_token_by_username::get_user_token_by_username;

crates/lib/redis_common/src/redis_cache_keys.rs

@@ -9,6 +9,10 @@ impl RedisCacheKeys {
     format!("get_tts_model:{}", model_token)
   }
 
+  pub fn get_tts_model_for_inference_migration_endpoint(model_token: &str) -> String {
+    format!("get_tts_inf:{}", model_token)
+  }
+
   pub fn get_unified_queue_stats_endpoint() -> &'static str {
     "q_stats_2" // NB: Schema changed, hence _2
   }

crates/service/web/storyteller_web/src/http_server/endpoints/tts/enqueue_infer_tts_handler/enqueue_infer_tts_handler.rs

@@ -16,15 +16,14 @@ use rand::seq::SliceRandom;
 
 use enums::by_table::generic_inference_jobs::inference_category::InferenceCategory;
 use enums::by_table::generic_inference_jobs::inference_model_type::InferenceModelType;
-use enums::by_table::tts_models::tts_model_type::TtsModelType;
 use enums::common::visibility::Visibility;
 use http_server_common::request::get_request_api_token::get_request_api_token;
 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 migration::text_to_speech::get_tts_model_for_inference_migration::TtsModelForInferenceMigration;
 use mysql_queries::payloads::generic_inference_args::generic_inference_args::{GenericInferenceArgs, InferenceCategoryAbbreviated};
 use mysql_queries::queries::generic_inference::web::insert_generic_inference_job::{insert_generic_inference_job, InsertGenericInferenceArgs};
 use mysql_queries::queries::tts::tts_inference_jobs::insert_tts_inference_job::TtsInferenceJobInsertBuilder;
-use mysql_queries::queries::tts::tts_models::get_tts_model::TtsModelRecord;
 use redis_common::redis_keys::RedisKeys;
 use tokens::tokens::tts_inference_jobs::TtsInferenceJobToken;
 use tokens::tokens::users::UserToken;
@@ -36,7 +35,7 @@ use crate::configs::app_startup::username_set::UsernameSet;
 use crate::configs::plans::get_correct_plan_for_session::get_correct_plan_for_session;
 use crate::configs::plans::plan::Plan;
 use crate::http_server::endpoints::investor_demo::demo_cookie::request_has_demo_cookie;
-use crate::http_server::endpoints::tts::enqueue_infer_tts_handler::get_model_with_caching::get_model_with_caching;
+use crate::http_server::endpoints::tts::enqueue_infer_tts_handler::get_tts_model_with_caching::get_tts_model_with_caching;
 use crate::http_server::web_utils::response_error_helpers::to_simple_json_error;
 use crate::server_state::ServerState;
 
@@ -286,10 +285,11 @@ pub async fn enqueue_infer_tts_handler(
 
   // ==================== CHECK TTS MODEL EXISTENCE / SETTINGS==================== //
 
-  let tts_model_lookup_result = get_model_with_caching(
+  let tts_model_lookup_result = get_tts_model_with_caching(
     &request.tts_model_token,
     &server_state.redis_ttl_cache,
-    &mut mysql_connection
+    &mut mysql_connection,
+    server_state.flags.switch_tts_to_model_weights,
   ).await;
 
   let tts_model = match tts_model_lookup_result {
@@ -377,18 +377,12 @@ pub async fn enqueue_infer_tts_handler(
 
   // ==================== ENQUEUE APPROPRIATE TTS TYPE ==================== //
 
-  let use_new_job_system = match tts_model.tts_model_type {
-    TtsModelType::Vits => true, // VITS is only on the new system.
-    TtsModelType::Tacotron2 => use_new_job_system_for_request || server_state.flags.enable_enqueue_generic_tts_job,
-  };
+  let use_new_job_system = use_new_job_system_for_request || server_state.flags.enable_enqueue_generic_tts_job;
 
   let job_token;
   let job_token_type;
 
-  let model_type = match tts_model.tts_model_type {
-    TtsModelType::Tacotron2 => InferenceModelType::Tacotron2,
-    TtsModelType::Vits => InferenceModelType::Vits,
-  };
+  let model_type = InferenceModelType::Tacotron2;
 
   let max_duration_seconds = plan.tts_max_duration_seconds();
 
@@ -559,10 +553,10 @@ pub fn validate_inference_text(text: &str, plan: &Plan) -> Result<(), String> {
 /// This is designed to fail closed (read: actually open!) rather than hit resources.
 async fn check_if_authorized_to_use_model(
   maybe_user_session: Option<&UserSessionExtended>,
-  tts_model: &TtsModelRecord,
+  tts_model: &TtsModelForInferenceMigration,
 ) -> bool {
 
-  match tts_model.creator_set_visibility {
+  match tts_model.creator_set_visibility() {
     Visibility::Public => return true,
     Visibility::Hidden => return true,
     Visibility::Private => {} // Fall through
@@ -570,7 +564,7 @@ async fn check_if_authorized_to_use_model(
 
   let is_authorized = maybe_user_session
       .map(|session| {
-        (session.user_token.as_str() == tts_model.creator_user_token.as_str()) // is_author
+        (session.user_token.as_str() == tts_model.creator_user_token()) // is_author
             || (session.role.can_ban_users) // is_mod
             || (session.user_token.as_str() == USER_FAKEYOU_USER_TOKEN) // is allowed user #1
             || (session.user_token.as_str() == USER_NEWS_STORY_USER_TOKEN) // is allowed user #2

crates/service/web/storyteller_web/src/http_server/endpoints/tts/enqueue_infer_tts_handler/get_tts_model_with_caching.rs

@@ -3,32 +3,34 @@ use sqlx::MySql;
 use sqlx::pool::PoolConnection;
 
 use errors::AnyhowResult;
-use mysql_queries::queries::tts::tts_models::get_tts_model::{get_tts_model_by_token_using_connection, TtsModelRecord};
+use migration::text_to_speech::get_tts_model_for_inference_migration::{get_tts_model_for_inference_migration, TtsModelForInferenceMigration};
 use redis_caching::redis_ttl_cache::RedisTtlCache;
 use redis_common::redis_cache_keys::RedisCacheKeys;
 
 // TODO: Read from in-memory TTS List cache first (!!!) for further performance savings
 
-pub async fn get_model_with_caching(
+pub async fn get_tts_model_with_caching(
   model_token: &str,
   redis_ttl_cache: &RedisTtlCache,
   mysql_connection: &mut PoolConnection<MySql>,
-) -> AnyhowResult<Option<TtsModelRecord>> {
-
+  use_weights_table: bool,
+) -> AnyhowResult<Option<TtsModelForInferenceMigration>> {
   // NB: Copy due to move. Not a cloned ref due to clippy lint
   let model_token2 = model_token.to_string();
 
   let get_tts_model = move || {
     // NB: async closures are not yet stable in Rust, so we include an async block.
     async move {
-      get_tts_model_by_token_using_connection(
+      get_tts_model_for_inference_migration(
         &model_token2,
+        mysql_connection,
         true,
-        mysql_connection).await
+        use_weights_table,
+      ).await
     }
   };
 
-  let cache_key = RedisCacheKeys::get_tts_model_endpoint(model_token);
+  let cache_key = RedisCacheKeys::get_tts_model_for_inference_migration_endpoint(model_token);
 
   let mut redis_ttl_cache_connection = match redis_ttl_cache.get_connection() {
     Ok(connection) => connection,
@@ -43,4 +45,3 @@ pub async fn get_model_with_caching(
     get_tts_model()
   }).await
 }
-

crates/service/web/storyteller_web/src/http_server/endpoints/tts/enqueue_infer_tts_handler/mod.rs

@@ -1,2 +1,2 @@
 pub mod enqueue_infer_tts_handler;
-pub mod get_model_with_caching;
+pub mod get_tts_model_with_caching;

crates/service/web/storyteller_web/src/http_server/endpoints/weights/delete_weight.rs

@@ -5,30 +5,22 @@ use actix_web::{HttpRequest, HttpResponse, web};
 use actix_web::error::ResponseError;
 use actix_web::http::StatusCode;
 use actix_web::web::Path;
-use log::{error, log, warn};
-use enums::by_table::audit_logs::audit_log_entity_action::AuditLogEntityAction::Delete;
+use log::warn;
+use utoipa::ToSchema;
 
-use http_server_common::request::get_request_ip::get_request_ip;
 use http_server_common::response::serialize_as_json_error::serialize_as_json_error;
 use mysql_queries::queries::model_weights::delete_weights::{
-  delete_weights_as_user,
   delete_weights_as_mod,
-  undelete_weights_as_user,
-  undelete_weights_as_mod
+  delete_weights_as_user,
+  undelete_weights_as_user
 };
-use mysql_queries::queries::model_weights::get_weight::get_weight_by_token;
-
+use mysql_queries::queries::model_weights::get::get_weight::get_weight_by_token;
 use tokens::tokens::model_weights::ModelWeightToken;
 
 use crate::http_server::web_utils::response_success_helpers::simple_json_success;
 use crate::server_state::ServerState;
 use crate::util::delete_role_disambiguation::{delete_role_disambiguation, DeleteRole};
 
-use enums::common::visibility::Visibility;
-use tokens::tokens::users::UserToken;
-
-use utoipa::ToSchema;
-
 #[derive(Deserialize,ToSchema)]
 pub struct DeleteWeightPathInfo {
   weight_token: String, 

crates/service/web/storyteller_web/src/http_server/endpoints/weights/get_weight.rs

@@ -13,10 +13,10 @@ use buckets::public::media_files::bucket_file_path::MediaFileBucketPath;
 use enums::by_table::model_weights::weights_category::WeightsCategory;
 use enums::by_table::model_weights::weights_types::WeightsType;
 use enums::common::visibility::Visibility;
-use mysql_queries::queries::model_weights::get_weight::get_weight_by_token;
+use mysql_queries::queries::model_weights::get::get_weight::get_weight_by_token;
 use tokens::tokens::model_weights::ModelWeightToken;
-use crate::http_server::common_responses::cover_image_details::CoverImageDetails;
 
+use crate::http_server::common_responses::cover_image_details::CoverImageDetails;
 use crate::http_server::common_responses::simple_entity_stats::SimpleEntityStats;
 use crate::http_server::common_responses::user_details_lite::UserDetailsLight;
 use crate::server_state::ServerState;

crates/service/web/storyteller_web/src/http_server/endpoints/weights/set_model_weight_cover_image.rs

@@ -15,7 +15,7 @@ use http_server_common::response::response_success_helpers::simple_json_success;
 use http_server_common::response::serialize_as_json_error::serialize_as_json_error;
 use mysql_queries::queries::media_files::get_media_file::get_media_file;
 use mysql_queries::queries::model_weights::edit::set_model_weight_cover_image::{set_model_weight_cover_image, UpdateArgs};
-use mysql_queries::queries::model_weights::get_weight::get_weight_by_token;
+use mysql_queries::queries::model_weights::get::get_weight::get_weight_by_token;
 use tokens::tokens::media_files::MediaFileToken;
 use tokens::tokens::model_weights::ModelWeightToken;
 

crates/service/web/storyteller_web/src/http_server/endpoints/weights/update_weight.rs

@@ -7,14 +7,14 @@ use actix_web::http::StatusCode;
 use actix_web::web::Path;
 use log::warn;
 use utoipa::ToSchema;
-use enums::by_table::media_files::media_file_type::MediaFileType;
 
+use enums::by_table::media_files::media_file_type::MediaFileType;
 use enums::common::visibility::Visibility;
 use http_server_common::response::response_success_helpers::simple_json_success;
 use http_server_common::response::serialize_as_json_error::serialize_as_json_error;
 use mysql_queries::queries::media_files::get_media_file::get_media_file;
 use mysql_queries::queries::model_weights::edit::update_weight::{CoverImageOption, update_weights, UpdateWeightArgs};
-use mysql_queries::queries::model_weights::get_weight::get_weight_by_token;
+use mysql_queries::queries::model_weights::get::get_weight::get_weight_by_token;
 use tokens::tokens::media_files::MediaFileToken;
 use tokens::tokens::model_weights::ModelWeightToken;
 use user_input_common::check_for_slurs::contains_slurs;