From d518d544719b43bce28ea26dfb83b19942e407e3 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Mon, 18 Dec 2023 00:49:19 -0500 Subject: [PATCH] download job query --- Cargo.lock | 1 + ...ight_from_voice_conversion_download_job.rs | 142 ++++++++++++++++++ .../src/queries/model_weights/create/mod.rs | 1 + .../src/queries/model_weights/mod.rs | 1 + crates/service/job/download_job/Cargo.toml | 1 + .../rvc_v2/process_rvc_v2_model.rs | 34 ++++- .../so_vits_svc/process_so_vits_svc_model.rs | 34 ++++- 7 files changed, 210 insertions(+), 4 deletions(-) create mode 100644 crates/lib/mysql_queries/src/queries/model_weights/create/create_model_weight_from_voice_conversion_download_job.rs create mode 100644 crates/lib/mysql_queries/src/queries/model_weights/create/mod.rs diff --git a/Cargo.lock b/Cargo.lock index a6a2c3ec3d..10c9daad42 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1906,6 +1906,7 @@ dependencies = [ "subprocess", "subprocess_common", "tempdir", + "tokens", "tokio", "zip", ] diff --git a/crates/lib/mysql_queries/src/queries/model_weights/create/create_model_weight_from_voice_conversion_download_job.rs b/crates/lib/mysql_queries/src/queries/model_weights/create/create_model_weight_from_voice_conversion_download_job.rs new file mode 100644 index 0000000000..e066c68045 --- /dev/null +++ b/crates/lib/mysql_queries/src/queries/model_weights/create/create_model_weight_from_voice_conversion_download_job.rs @@ -0,0 +1,142 @@ +use std::path::Path; + +use anyhow::anyhow; +use sqlx::MySqlPool; +use enums::by_table::model_weights::weights_category::WeightsCategory; +use enums::by_table::model_weights::weights_types::WeightsType; + +use enums::by_table::voice_conversion_models::voice_conversion_model_type::VoiceConversionModelType; +use enums::common::visibility::Visibility; +use errors::AnyhowResult; +use tokens::tokens::model_weights::ModelWeightToken; +use tokens::tokens::users::UserToken; +use tokens::tokens::voice_conversion_models::VoiceConversionModelToken; + +pub struct CreateModelWeightArgs<'a> { + pub maybe_use_model_weights_token: Option, + pub maybe_old_voice_conversion_model_token: Option, + + pub weights_type: WeightsType, + pub weights_category: WeightsCategory, + + pub title: &'a str, + + pub original_download_url: &'a str, + pub original_filename: &'a str, + pub file_size_bytes: u64, + + pub creator_user_token: &'a UserToken, + pub creator_ip_address: &'a str, + pub creator_set_visibility: Visibility, + + /// (For rvc_v2 models) - whether the model has an associated `.index` file. + pub has_index_file: bool, + + pub public_bucket_hash: String, + pub maybe_public_bucket_prefix: Option, + pub maybe_public_bucket_extension: Option, + + pub mysql_pool: &'a MySqlPool, +} + +pub async fn create_model_weight_from_voice_conversion_download_job( + args: CreateModelWeightArgs<'_>, +) -> AnyhowResult<(u64, ModelWeightToken)> { + + let weights_token = args.maybe_use_model_weights_token + .map(|token| token.clone()) + .unwrap_or_else(|| ModelWeightToken::generate()); + + let mut transaction = args.mysql_pool.begin().await?; + + // NB: Not setting the following: + // - created_at + // - updated_at + // - user_deleted_at + // - mod_deleted_at + // - version = 0 + + let query_result = sqlx::query!( + r#" +INSERT INTO model_weights +SET + token = ?, + weights_type = ?, + weights_category = ?, + title = ?, + + maybe_thumbnail_token = NULL, + maybe_avatar_media_file_token = NULL, + maybe_cover_media_file_token = NULL, + description_markdown = NULL, + description_rendered_html = NULL, + + maybe_last_update_user_token = NULL, + original_download_url = ?, + original_filename = ?, + file_size_bytes = ?, + file_checksum_sha2 = "TODO", + + creator_user_token = ?, + creator_ip_address = ?, + creator_set_visibility = ?, + + private_bucket_hash = ?, + maybe_private_bucket_prefix = ?, + maybe_private_bucket_extension = ?, + + cached_user_ratings_total_count = 0, + cached_user_ratings_positive_count = 0, + cached_user_ratings_negative_count = 0, + maybe_cached_user_ratings_ratio = 0.0, + cached_user_ratings_last_updated_at = NOW(), + + maybe_migration_old_model_token = ? + "#, + &weights_token, + args.weights_type, + args.weights_category, + args.title, + + args.original_download_url, + args.original_filename, + args.file_size_bytes, + + args.creator_user_token, + args.creator_ip_address, + args.creator_set_visibility.to_str(), + + args.public_bucket_hash, + args.maybe_public_bucket_prefix, + args.maybe_public_bucket_extension, + + args.maybe_old_voice_conversion_model_token, + ) + .execute(&mut *transaction) + .await; + + let model_weights_record_id = match query_result { + Ok(res) => res.last_insert_id(), + Err(err) => return Err(anyhow!("Mysql error: {:?}", err)), + }; + + let query_result = sqlx::query!( + r#" +INSERT INTO model_weights_extension_voice_conversion_details +SET + model_weights_token = ?, + has_index_file = ? + "#, + &weights_token, + args.has_index_file + ) + .execute(&mut *transaction) + .await; + + match query_result { + Ok(_res) => {}, + Err(err) => return Err(anyhow!("Mysql error: {:?}", err)), + }; + + Ok((model_weights_record_id, weights_token)) +} diff --git a/crates/lib/mysql_queries/src/queries/model_weights/create/mod.rs b/crates/lib/mysql_queries/src/queries/model_weights/create/mod.rs new file mode 100644 index 0000000000..2b876c4c48 --- /dev/null +++ b/crates/lib/mysql_queries/src/queries/model_weights/create/mod.rs @@ -0,0 +1 @@ +pub mod create_model_weight_from_voice_conversion_download_job; \ No newline at end of file diff --git a/crates/lib/mysql_queries/src/queries/model_weights/mod.rs b/crates/lib/mysql_queries/src/queries/model_weights/mod.rs index f18223ec5c..e4ca37a05b 100644 --- a/crates/lib/mysql_queries/src/queries/model_weights/mod.rs +++ b/crates/lib/mysql_queries/src/queries/model_weights/mod.rs @@ -1,3 +1,4 @@ +pub mod create; pub mod create_weight; pub mod delete_weights; pub mod get_weight; diff --git a/crates/service/job/download_job/Cargo.toml b/crates/service/job/download_job/Cargo.toml index e97068915d..98ec26fcb2 100644 --- a/crates/service/job/download_job/Cargo.toml +++ b/crates/service/job/download_job/Cargo.toml @@ -32,6 +32,7 @@ mimetypes = { path = "../../../lib/files/mimetypes" } mysql_queries = { path = "../../../lib/mysql_queries" } reusable_types = { path = "../../../_types/reusable_types" } subprocess_common = { path = "../../../lib/subprocess_common" } +tokens = { path = "../../../_types/tokens" } # Workspace hack (faster compile times) my-workspace-hack = { version = "0.1", path = "../../../my-workspace-hack" } diff --git a/crates/service/job/download_job/src/job_types/voice_conversion/rvc_v2/process_rvc_v2_model.rs b/crates/service/job/download_job/src/job_types/voice_conversion/rvc_v2/process_rvc_v2_model.rs index a7f7fea4d3..9d2eaa3542 100644 --- a/crates/service/job/download_job/src/job_types/voice_conversion/rvc_v2/process_rvc_v2_model.rs +++ b/crates/service/job/download_job/src/job_types/voice_conversion/rvc_v2/process_rvc_v2_model.rs @@ -5,6 +5,8 @@ use log::{error, info, warn}; use tempdir::TempDir; use crockford::crockford_entropy_lower; +use enums::by_table::model_weights::weights_category::WeightsCategory; +use enums::by_table::model_weights::weights_types::WeightsType; use enums::by_table::voice_conversion_models::voice_conversion_model_type::VoiceConversionModelType; use enums::common::visibility::Visibility; use errors::AnyhowResult; @@ -15,7 +17,9 @@ use filesys::safe_delete_temp_directory::safe_delete_temp_directory; use filesys::safe_delete_temp_file::safe_delete_temp_file; use jobs_common::redis_job_status_logger::RedisJobStatusLogger; use mysql_queries::queries::generic_download::job::list_available_generic_download_jobs::AvailableDownloadJob; +use mysql_queries::queries::model_weights::create::create_model_weight_from_voice_conversion_download_job::{create_model_weight_from_voice_conversion_download_job, CreateModelWeightArgs}; use mysql_queries::queries::voice_conversion::models::insert_voice_conversion_model_from_download_job::{insert_voice_conversion_model_from_download_job, InsertVoiceConversionModelArgs}; +use tokens::tokens::model_weights::ModelWeightToken; use crate::job_loop::job_results::JobResults; use crate::job_types::voice_conversion::rvc_v2::extract_rvc_files::{DownloadedRvcFile, extract_rvc_files}; @@ -140,6 +144,8 @@ pub async fn process_rvc_v2_model<'a, 'b>( } } + // TODO(bt, 2023-12-18): Upload model weights files !! + // ==================== DELETE DOWNLOADED FILE ==================== // // NB: We should be using a tempdir, but to make absolutely certain we don't overflow the disk... @@ -153,7 +159,10 @@ pub async fn process_rvc_v2_model<'a, 'b>( info!("Saving Voice Conversion model record..."); - let (_id, model_token) = insert_voice_conversion_model_from_download_job(InsertVoiceConversionModelArgs { + // TODO(bt,2023-12-18): Once migration is done, move this back into the query code. + let model_weights_token = ModelWeightToken::generate(); + + let (_id, voice_conversion_model_token) = insert_voice_conversion_model_from_download_job(InsertVoiceConversionModelArgs { model_type: VoiceConversionModelType::RvcV2, title: &job.title, original_download_url: &job.download_url, @@ -168,6 +177,27 @@ pub async fn process_rvc_v2_model<'a, 'b>( mysql_pool: &job_state.mysql_pool, }).await?; + info!("Saving model weights record..."); + + let (_id, _model_weights_token) = create_model_weight_from_voice_conversion_download_job(CreateModelWeightArgs { + maybe_use_model_weights_token: Some(model_weights_token.clone()), + maybe_old_voice_conversion_model_token: Some(voice_conversion_model_token.clone()), + weights_type: WeightsType::RvcV2, + weights_category: WeightsCategory::VoiceConversion, + title: &job.title, + original_download_url: &job.download_url, + original_filename: &download_filename, + file_size_bytes, + creator_user_token: &job.creator_user_token, + creator_ip_address: &job.creator_ip_address, + creator_set_visibility: Visibility::Public, // TODO: All models default to public at start + has_index_file: maybe_original_model_index_file_path.is_some(), + public_bucket_hash: "TODO".to_string(), // TODO: This needs to be finished. + maybe_public_bucket_prefix: None, + maybe_public_bucket_extension: None, + mysql_pool: &job_state.mysql_pool, + }).await?; + job_state.badge_granter.maybe_grant_voice_conversion_model_uploads_badge(&job.creator_user_token) .await .map_err(|e| { @@ -176,7 +206,7 @@ pub async fn process_rvc_v2_model<'a, 'b>( })?; Ok(JobResults { - entity_token: Some(model_token.to_string()), + entity_token: Some(voice_conversion_model_token.to_string()), // TODO: Swap model token. entity_type: Some(VoiceConversionModelType::RvcV2.to_string()), // NB: This may be different from `GenericDownloadType` in the future! }) } diff --git a/crates/service/job/download_job/src/job_types/voice_conversion/so_vits_svc/process_so_vits_svc_model.rs b/crates/service/job/download_job/src/job_types/voice_conversion/so_vits_svc/process_so_vits_svc_model.rs index d7500c91f3..6502bdbeb6 100644 --- a/crates/service/job/download_job/src/job_types/voice_conversion/so_vits_svc/process_so_vits_svc_model.rs +++ b/crates/service/job/download_job/src/job_types/voice_conversion/so_vits_svc/process_so_vits_svc_model.rs @@ -3,6 +3,8 @@ use std::path::PathBuf; use anyhow::anyhow; use log::{error, info, warn}; use tempdir::TempDir; +use enums::by_table::model_weights::weights_category::WeightsCategory; +use enums::by_table::model_weights::weights_types::WeightsType; use enums::by_table::voice_conversion_models::voice_conversion_model_type::VoiceConversionModelType; use enums::common::visibility::Visibility; @@ -14,7 +16,9 @@ use filesys::safe_delete_temp_file::safe_delete_temp_file; use hashing::sha256::sha256_hash_file::sha256_hash_file; use jobs_common::redis_job_status_logger::RedisJobStatusLogger; use mysql_queries::queries::generic_download::job::list_available_generic_download_jobs::AvailableDownloadJob; +use mysql_queries::queries::model_weights::create::create_model_weight_from_voice_conversion_download_job::{create_model_weight_from_voice_conversion_download_job, CreateModelWeightArgs}; use mysql_queries::queries::voice_conversion::models::insert_voice_conversion_model_from_download_job::{insert_voice_conversion_model_from_download_job, InsertVoiceConversionModelArgs}; +use tokens::tokens::model_weights::ModelWeightToken; use crate::job_loop::job_results::JobResults; use crate::job_types::voice_conversion::so_vits_svc::so_vits_svc_model_check_command::{CheckArgs, Device}; @@ -86,6 +90,8 @@ pub async fn process_so_vits_svc_model<'a, 'b>( return Err(err); } + // TODO(bt, 2023-12-18): Upload model weights files !! + // ==================== DELETE DOWNLOADED FILE ==================== // // NB: We should be using a tempdir, but to make absolutely certain we don't overflow the disk... @@ -98,7 +104,10 @@ pub async fn process_so_vits_svc_model<'a, 'b>( info!("Saving Voice Conversion model record..."); - let (_id, model_token) = insert_voice_conversion_model_from_download_job(InsertVoiceConversionModelArgs { + // TODO(bt,2023-12-18): Once migration is done, move this back into the query code. + let model_weights_token = ModelWeightToken::generate(); + + let (_id, voice_conversion_model_token) = insert_voice_conversion_model_from_download_job(InsertVoiceConversionModelArgs { model_type: VoiceConversionModelType::SoVitsSvc, title: &job.title, original_download_url: &job.download_url, @@ -113,6 +122,27 @@ pub async fn process_so_vits_svc_model<'a, 'b>( mysql_pool: &job_state.mysql_pool, }).await?; + info!("Saving model weights record..."); + + let (_id, _model_weights_token) = create_model_weight_from_voice_conversion_download_job(CreateModelWeightArgs { + maybe_use_model_weights_token: Some(model_weights_token.clone()), + maybe_old_voice_conversion_model_token: Some(voice_conversion_model_token.clone()), + weights_type: WeightsType::SoVitsSvc, + weights_category: WeightsCategory::VoiceConversion, + title: &job.title, + original_download_url: &job.download_url, + original_filename: &download_filename, + file_size_bytes, + creator_user_token: &job.creator_user_token, + creator_ip_address: &job.creator_ip_address, + creator_set_visibility: Visibility::Public, // TODO: All models default to public at start + has_index_file: false, // SVC do not have index files. + public_bucket_hash: "TODO".to_string(), // TODO: This needs to be finished. + maybe_public_bucket_prefix: None, + maybe_public_bucket_extension: None, + mysql_pool: &job_state.mysql_pool, + }).await?; + job_state.badge_granter.maybe_grant_voice_conversion_model_uploads_badge(&job.creator_user_token) .await .map_err(|e| { @@ -121,7 +151,7 @@ pub async fn process_so_vits_svc_model<'a, 'b>( })?; Ok(JobResults { - entity_token: Some(model_token.to_string()), + entity_token: Some(voice_conversion_model_token.to_string()), // TODO: Swap model token. entity_type: Some(VoiceConversionModelType::SoVitsSvc.to_string()), // NB: This may be different from `GenericDownloadType` in the future! }) }