download job query

Cargo.lock

@@ -1906,6 +1906,7 @@ dependencies = [
  "subprocess",
  "subprocess_common",
  "tempdir",
+ "tokens",
  "tokio",
  "zip",
 ]

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<ModelWeightToken>,
+  pub maybe_old_voice_conversion_model_token: Option<VoiceConversionModelToken>,
+
+  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<String>,
+  pub maybe_public_bucket_extension: Option<String>,
+
+  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))
+}

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

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

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;

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

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

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