From 91695cde00a22c590cd99bbe8e44f149d5f95444 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Fri, 16 Dec 2022 02:41:40 -0500 Subject: [PATCH] downloader job: handle softvc hifigan --- .idea/dictionaries/bt.xml | 1 + .../up.sql | 1 + .../2022-08-05-202110_vocoder_models/up.sql | 4 +- .../src/mediators/badge_granter.rs | 86 ++++++++++- .../src/achievements/user_badge.rs | 60 +++++++- .../src/achievements/user_badge_type.rs | 3 + .../src/db/enums/vocoder_type.rs | 23 +++ .../src/bucket_path_unifier.rs | 1 + .../service/job/download_job/src/job_state.rs | 2 + .../src/job_steps/process_single_job.rs | 25 +++- .../hifigan/hifigan_model_check_command.rs | 2 +- .../hifigan_softvc_model_check_command.rs | 84 +++++++++++ .../src/job_types/hifigan_softvc/mod.rs | 2 + .../process_hifigan_softvc_vocoder.rs | 134 ++++++++++++++++++ .../job/download_job/src/job_types/mod.rs | 1 + 15 files changed, 421 insertions(+), 8 deletions(-) create mode 100644 crates/service/job/download_job/src/job_types/hifigan_softvc/hifigan_softvc_model_check_command.rs create mode 100644 crates/service/job/download_job/src/job_types/hifigan_softvc/mod.rs create mode 100644 crates/service/job/download_job/src/job_types/hifigan_softvc/process_hifigan_softvc_vocoder.rs diff --git a/.idea/dictionaries/bt.xml b/.idea/dictionaries/bt.xml index d3699ac475..562630d84e 100644 --- a/.idea/dictionaries/bt.xml +++ b/.idea/dictionaries/bt.xml @@ -9,6 +9,7 @@ fakeyou hifigan scrollback + softvc tacotron transcoded transcodings diff --git a/_migrations/2022-08-05-202017_generic_download_jobs/up.sql b/_migrations/2022-08-05-202017_generic_download_jobs/up.sql index afcf500440..2b84c5bde7 100644 --- a/_migrations/2022-08-05-202017_generic_download_jobs/up.sql +++ b/_migrations/2022-08-05-202017_generic_download_jobs/up.sql @@ -22,6 +22,7 @@ CREATE TABLE generic_download_jobs ( -- The type of the object will vary based on the type of the upload, -- and we may include heuristics that auto-detect types in the future + -- Example types: 'hifigan', 'hifigan_rocket_vc' (codename for softvc), 'tacotron2', ... on_success_downloaded_entity_type VARCHAR(32) DEFAULT NULL, -- ========== UPLOAD DETAILS ========== diff --git a/_migrations/2022-08-05-202110_vocoder_models/up.sql b/_migrations/2022-08-05-202110_vocoder_models/up.sql index 1ee7c74c58..d11793c780 100644 --- a/_migrations/2022-08-05-202110_vocoder_models/up.sql +++ b/_migrations/2022-08-05-202110_vocoder_models/up.sql @@ -10,7 +10,9 @@ CREATE TABLE vocoder_models ( token VARCHAR(32) NOT NULL, -- The architecture of the vocoder. - -- Currently supported values: 'hifigan' + -- Currently supported values: + -- * 'hifigan' - hifigan for Tacotron2 + -- * 'hifigan_rocket_vc' - hifigan for SoftVC ("rocket" is a codename) vocoder_type VARCHAR(32) NOT NULL, -- ========== MUTABLE METADATA ========== diff --git a/crates/lib/database_queries/src/mediators/badge_granter.rs b/crates/lib/database_queries/src/mediators/badge_granter.rs index 34dcee3975..9adb4a85d2 100644 --- a/crates/lib/database_queries/src/mediators/badge_granter.rs +++ b/crates/lib/database_queries/src/mediators/badge_granter.rs @@ -38,7 +38,11 @@ impl BadgeGranter { let mut maybe_badge = None; - if count >= 250 { + if count >= 1000 { + maybe_badge = Some(UserBadge::TtsModelUploader1000); + } else if count >= 500 { + maybe_badge = Some(UserBadge::TtsModelUploader500); + } else if count >= 250 { maybe_badge = Some(UserBadge::TtsModelUploader250); } else if count >= 200 { maybe_badge = Some(UserBadge::TtsModelUploader200); @@ -133,7 +137,11 @@ impl BadgeGranter { let mut maybe_badge = None; - if count >= 250 { + if count >= 1000 { + maybe_badge = Some(UserBadge::VocoderModelUploader1000); + } else if count >= 500 { + maybe_badge = Some(UserBadge::VocoderModelUploader500); + } else if count >= 250 { maybe_badge = Some(UserBadge::VocoderModelUploader250); } else if count >= 200 { maybe_badge = Some(UserBadge::VocoderModelUploader200); @@ -173,6 +181,56 @@ impl BadgeGranter { Ok(()) } + /// This needs to be called *after* successful upload. + pub async fn maybe_grant_softvc_vocoder_model_uploads_badge(&self, user_token: &str) -> AnyhowResult<()> { + let count = self.count_softvc_vocoder_models_uploaded(user_token).await?; + + let mut maybe_badge = None; + + if count >= 1000 { + maybe_badge = Some(UserBadge::VocoderRocketVcModelUploader1000); + } else if count >= 500 { + maybe_badge = Some(UserBadge::VocoderRocketVcModelUploader500); + } else if count >= 250 { + maybe_badge = Some(UserBadge::VocoderRocketVcModelUploader250); + } else if count >= 200 { + maybe_badge = Some(UserBadge::VocoderRocketVcModelUploader200); + } else if count >= 150 { + maybe_badge = Some(UserBadge::VocoderRocketVcModelUploader150); + } else if count >= 100 { + maybe_badge = Some(UserBadge::VocoderRocketVcModelUploader100); + } else if count >= 50 { + maybe_badge = Some(UserBadge::VocoderRocketVcModelUploader50); + } else if count >= 20 { + maybe_badge = Some(UserBadge::VocoderRocketVcModelUploader20); + } else if count >= 10 { + maybe_badge = Some(UserBadge::VocoderRocketVcModelUploader10); + } else if count >= 5 { + maybe_badge = Some(UserBadge::VocoderRocketVcModelUploader5); + } else if count >= 1 { + maybe_badge = Some(UserBadge::VocoderRocketVcModelUploader1); + } + + let badge = match maybe_badge { + Some(badge) => badge, + None => return Ok(()), + }; + + if self.has_badge(user_token, badge).await? { + return Ok(()) + } + + let _record_id = self.insert( + badge, + user_token, + ).await?; + + self.firehose_publisher.publish_user_badge_granted(user_token, badge.to_db_value()) + .await?; + + Ok(()) + } + // ======================================================================= pub async fn has_badge(&self, user_token: &str, user_badge: UserBadge) -> AnyhowResult { @@ -245,6 +303,30 @@ SELECT count(*) as count FROM vocoder_models WHERE creator_user_token = ? + AND vocoder_type = 'hifigan' +LIMIT 1 + "#, + user_token + ) + .fetch_one(&self.mysql_pool) + .await; + + self.handle_count_query(maybe_result) + } + + async fn count_softvc_vocoder_models_uploaded( + &self, + user_token: &str, + ) -> AnyhowResult { + // NB: This could get expensive! + let maybe_result = sqlx::query_as!( + CountRecord, + r#" +SELECT count(*) as count +FROM vocoder_models +WHERE + creator_user_token = ? + AND vocoder_type = 'hifigan_rocket_vc' LIMIT 1 "#, user_token diff --git a/crates/lib/reusable_types/src/achievements/user_badge.rs b/crates/lib/reusable_types/src/achievements/user_badge.rs index e45f42732a..d204af398c 100644 --- a/crates/lib/reusable_types/src/achievements/user_badge.rs +++ b/crates/lib/reusable_types/src/achievements/user_badge.rs @@ -17,6 +17,8 @@ pub enum UserBadge { TtsModelUploader150, TtsModelUploader200, TtsModelUploader250, + TtsModelUploader500, + TtsModelUploader1000, // Granted for vocoder models VocoderModelUploader1, @@ -28,6 +30,21 @@ pub enum UserBadge { VocoderModelUploader150, VocoderModelUploader200, VocoderModelUploader250, + VocoderModelUploader500, + VocoderModelUploader1000, + + // Granted for vocoder (softvc aka "rocket_vc") models + VocoderRocketVcModelUploader1, + VocoderRocketVcModelUploader5, + VocoderRocketVcModelUploader10, + VocoderRocketVcModelUploader20, + VocoderRocketVcModelUploader50, + VocoderRocketVcModelUploader100, + VocoderRocketVcModelUploader150, + VocoderRocketVcModelUploader200, + VocoderRocketVcModelUploader250, + VocoderRocketVcModelUploader500, + VocoderRocketVcModelUploader1000, // Granted for uploading templates W2lTemplateUploader1, @@ -57,6 +74,7 @@ pub enum UserBadge { } impl UserBadge { + /// This is the value that lives in the database as a VARCHAR(32) pub fn to_db_value(&self) -> &'static str { match self { UserBadge::EarlyUser => "early_user", @@ -69,6 +87,9 @@ impl UserBadge { UserBadge::TtsModelUploader150 => "tts_model_uploader_150", UserBadge::TtsModelUploader200 => "tts_model_uploader_200", UserBadge::TtsModelUploader250 => "tts_model_uploader_250", + UserBadge::TtsModelUploader500 => "tts_model_uploader_500", + UserBadge::TtsModelUploader1000 => "tts_model_uploader_1000", + UserBadge::VocoderModelUploader1 => "vocoder_model_uploader_1", UserBadge::VocoderModelUploader5 => "vocoder_model_uploader_5", UserBadge::VocoderModelUploader10 => "vocoder_model_uploader_10", @@ -78,6 +99,21 @@ impl UserBadge { UserBadge::VocoderModelUploader150 => "vocoder_model_uploader_150", UserBadge::VocoderModelUploader200 => "vocoder_model_uploader_200", UserBadge::VocoderModelUploader250 => "vocoder_model_uploader_250", + UserBadge::VocoderModelUploader500 => "vocoder_model_uploader_500", + UserBadge::VocoderModelUploader1000 => "vocoder_model_uploader_1000", + + UserBadge::VocoderRocketVcModelUploader1 => "vocoder_rocket_vc_uploader_1", + UserBadge::VocoderRocketVcModelUploader5 => "vocoder_rocket_vc_uploader_5", + UserBadge::VocoderRocketVcModelUploader10 => "vocoder_rocket_vc_uploader_10", + UserBadge::VocoderRocketVcModelUploader20 => "vocoder_rocket_vc_uploader_20", + UserBadge::VocoderRocketVcModelUploader50 => "vocoder_rocket_vc_uploader_50", + UserBadge::VocoderRocketVcModelUploader100 => "vocoder_rocket_vc_uploader_100", + UserBadge::VocoderRocketVcModelUploader150 => "vocoder_rocket_vc_uploader_150", + UserBadge::VocoderRocketVcModelUploader200 => "vocoder_rocket_vc_uploader_200", + UserBadge::VocoderRocketVcModelUploader250 => "vocoder_rocket_vc_uploader_250", + UserBadge::VocoderRocketVcModelUploader500 => "vocoder_rocket_vc_uploader_500", + UserBadge::VocoderRocketVcModelUploader1000 => "vocoder_rocket_vc_uploader_1000", // NB: 31 characters, table max is 32 + UserBadge::W2lTemplateUploader1 => "w2l_template_uploader_1", UserBadge::W2lTemplateUploader10 => "w2l_template_uploader_10", UserBadge::W2lTemplateUploader50 => "w2l_template_uploader_50", @@ -88,6 +124,7 @@ impl UserBadge { UserBadge::W2lTemplateUploader2000 => "w2l_template_uploader_2000", UserBadge::W2lTemplateUploader5000 => "w2l_template_uploader_5000", UserBadge::W2lTemplateUploader10000 => "w2l_template_uploader_10000", + //UserBadge::TtsInferenceUser100 => "tts_inference_100", //UserBadge::TtsInferenceUser500 => "tts_inference_500", //UserBadge::TtsInferenceUser1000 => "tts_inference_1000", @@ -113,7 +150,10 @@ impl UserBadge { | UserBadge::TtsModelUploader100 | UserBadge::TtsModelUploader150 | UserBadge::TtsModelUploader200 - | UserBadge::TtsModelUploader250 => UserBadgeType::TtsModelUploader, + | UserBadge::TtsModelUploader250 + | UserBadge::TtsModelUploader500 + | UserBadge::TtsModelUploader1000 + => UserBadgeType::TtsModelUploader, UserBadge::VocoderModelUploader1 | UserBadge::VocoderModelUploader5 @@ -123,7 +163,23 @@ impl UserBadge { | UserBadge::VocoderModelUploader100 | UserBadge::VocoderModelUploader150 | UserBadge::VocoderModelUploader200 - | UserBadge::VocoderModelUploader250 => UserBadgeType::VocoderModelUploader, + | UserBadge::VocoderModelUploader250 + | UserBadge::VocoderModelUploader500 + | UserBadge::VocoderModelUploader1000 + => UserBadgeType::VocoderModelUploader, + + UserBadge::VocoderRocketVcModelUploader1 + | UserBadge::VocoderRocketVcModelUploader5 + | UserBadge::VocoderRocketVcModelUploader10 + | UserBadge::VocoderRocketVcModelUploader20 + | UserBadge::VocoderRocketVcModelUploader50 + | UserBadge::VocoderRocketVcModelUploader100 + | UserBadge::VocoderRocketVcModelUploader150 + | UserBadge::VocoderRocketVcModelUploader200 + | UserBadge::VocoderRocketVcModelUploader250 + | UserBadge::VocoderRocketVcModelUploader500 + | UserBadge::VocoderRocketVcModelUploader1000 + => UserBadgeType::VocoderRocketVcModelUploader, UserBadge::W2lTemplateUploader1 | UserBadge::W2lTemplateUploader10 diff --git a/crates/lib/reusable_types/src/achievements/user_badge_type.rs b/crates/lib/reusable_types/src/achievements/user_badge_type.rs index 41426374ad..f5c3be7ecf 100644 --- a/crates/lib/reusable_types/src/achievements/user_badge_type.rs +++ b/crates/lib/reusable_types/src/achievements/user_badge_type.rs @@ -11,6 +11,9 @@ pub enum UserBadgeType { /// Granted for uploading vocoder models VocoderModelUploader, + /// Granted for uploading vocoder (softvc aka "rocket_vc") models + VocoderRocketVcModelUploader, + /// Granted for uploading W2L templates W2lTemplateUploader, diff --git a/crates/lib/reusable_types/src/db/enums/vocoder_type.rs b/crates/lib/reusable_types/src/db/enums/vocoder_type.rs index d262f48ad0..f25e432b22 100644 --- a/crates/lib/reusable_types/src/db/enums/vocoder_type.rs +++ b/crates/lib/reusable_types/src/db/enums/vocoder_type.rs @@ -1,8 +1,16 @@ #[derive(Clone, Copy, Eq, PartialEq, Debug, Deserialize, Serialize, sqlx::Type)] pub enum VocoderType { + /// NB: Note - this is hifigan for Tacotron2. + /// Some work will be needed to unify this with other hifigan types. #[serde(rename = "hifigan")] #[sqlx(rename = "hifigan")] HifiGan, + + /// NB: Note - this is hifigan for SoftVC (our internal codename is "rocketvc"). + /// Some work will need to be done to unify this with other hifigan types. + #[serde(rename = "hifigan_rocket_vc")] + #[sqlx(rename = "hifigan_rocket_vc")] + HifiGanRocketVc, } /// NB: Legacy API for older code. @@ -10,12 +18,14 @@ impl VocoderType { pub fn to_str(&self) -> &'static str { match self { Self::HifiGan=> "hifigan", + Self::HifiGanRocketVc => "hifigan_rocket_vc", } } pub fn from_str(value: &str) -> Result { match value { "hifigan" => Ok(Self::HifiGan), + "hifigan_rocket_vc" => Ok(Self::HifiGanRocketVc), _ => Err(format!("invalid value: {:?}", value)), } } @@ -29,5 +39,18 @@ mod tests { #[test] fn test_serialization() { assert_serialization(VocoderType::HifiGan, "hifigan"); + assert_serialization(VocoderType::HifiGanRocketVc, "hifigan_rocket_vc"); + } + + #[test] + fn to_str() { + assert_eq!(VocoderType::HifiGan.to_str(), "hifigan"); + assert_eq!(VocoderType::HifiGanRocketVc.to_str(), "hifigan_rocket_vc"); + } + + #[test] + fn from_str() { + assert_eq!(VocoderType::from_str("hifigan").unwrap(), VocoderType::HifiGan); + assert_eq!(VocoderType::from_str("hifigan_rocket_vc").unwrap(), VocoderType::HifiGanRocketVc); } } diff --git a/crates/lib/storage_buckets_common/src/bucket_path_unifier.rs b/crates/lib/storage_buckets_common/src/bucket_path_unifier.rs index a9611564d6..45dbc96c1d 100644 --- a/crates/lib/storage_buckets_common/src/bucket_path_unifier.rs +++ b/crates/lib/storage_buckets_common/src/bucket_path_unifier.rs @@ -66,6 +66,7 @@ impl BucketPathUnifier { } /// User-uploaded vocoders. + /// These can be HifiGan and HifiGanSoftVc pub fn vocoder_path(&self, vocoder_file_hash: &str) -> PathBuf { let hashed_path = Self::hashed_directory_path(vocoder_file_hash); let model_filename = format!("{}.pt", &vocoder_file_hash); diff --git a/crates/service/job/download_job/src/job_state.rs b/crates/service/job/download_job/src/job_state.rs index 18169265e3..82ded1e308 100644 --- a/crates/service/job/download_job/src/job_state.rs +++ b/crates/service/job/download_job/src/job_state.rs @@ -1,4 +1,5 @@ use crate::job_types::hifigan::hifigan_model_check_command::HifiGanModelCheckCommand; +use crate::job_types::hifigan_softvc::hifigan_softvc_model_check_command::HifiGanSoftVcModelCheckCommand; use crate::job_types::tacotron::tacotron_model_check_command::TacotronModelCheckCommand; use database_queries::mediators::badge_granter::BadgeGranter; use database_queries::mediators::firehose_publisher::FirehosePublisher; @@ -43,4 +44,5 @@ pub struct SidecarConfigs { pub google_drive_downloader: GoogleDriveDownloadCommand, pub tacotron_model_check_command: TacotronModelCheckCommand, pub hifigan_model_check_command: HifiGanModelCheckCommand, + pub hifigan_softvc_model_check_command: HifiGanSoftVcModelCheckCommand, } \ No newline at end of file diff --git a/crates/service/job/download_job/src/job_steps/process_single_job.rs b/crates/service/job/download_job/src/job_steps/process_single_job.rs index a18ae24b1d..b11d6daddd 100644 --- a/crates/service/job/download_job/src/job_steps/process_single_job.rs +++ b/crates/service/job/download_job/src/job_steps/process_single_job.rs @@ -4,6 +4,7 @@ use container_common::anyhow_result::AnyhowResult; use container_common::filesystem::safe_delete_temp_directory::safe_delete_temp_directory; use crate::JobState; use crate::job_types::hifigan::process_hifigan_vocoder::process_hifigan_vocoder; +use crate::job_types::hifigan_softvc::process_hifigan_softvc_vocoder::process_hifigan_softvc_vocoder; use crate::job_types::tacotron::process_tacotron_model::process_tacotron_model; use database_queries::queries::generic_download::job::list_available_generic_download_jobs::AvailableDownloadJob; use database_queries::queries::generic_download::job::mark_generic_download_job_done::mark_generic_download_job_done; @@ -69,6 +70,28 @@ pub async fn process_single_job(job_state: &JobState, job: &AvailableDownloadJob entity_token = results.entity_token.clone(); entity_type = results.entity_type.clone(); } + GenericDownloadType::HifiGanRocketVc => { + let results = process_hifigan_softvc_vocoder( + job_state, + job, + &temp_dir, + &download_filename, + &mut redis_logger, + ).await?; + entity_token = results.entity_token.clone(); + entity_type = results.entity_type.clone(); + } + GenericDownloadType::RocketVc => { + let results = process_softvc_model( + job_state, + job, + &temp_dir, + &download_filename, + &mut redis_logger, + ).await?; + entity_token = results.entity_token.clone(); + entity_type = results.entity_type.clone(); + } GenericDownloadType::Tacotron2 => { let results = process_tacotron_model( job_state, @@ -80,8 +103,6 @@ pub async fn process_single_job(job_state: &JobState, job: &AvailableDownloadJob entity_token = results.entity_token.clone(); entity_type = results.entity_type.clone(); } - GenericDownloadType::HifiGanRocketVc => {} - GenericDownloadType::RocketVc => {} } // ===================================================== diff --git a/crates/service/job/download_job/src/job_types/hifigan/hifigan_model_check_command.rs b/crates/service/job/download_job/src/job_types/hifigan/hifigan_model_check_command.rs index 37a8cc5d62..9bf645a400 100644 --- a/crates/service/job/download_job/src/job_types/hifigan/hifigan_model_check_command.rs +++ b/crates/service/job/download_job/src/job_types/hifigan/hifigan_model_check_command.rs @@ -1,7 +1,7 @@ use crate::AnyhowResult; use log::info; -use subprocess::{Popen, PopenConfig}; use std::path::Path; +use subprocess::{Popen, PopenConfig}; use subprocess_common::docker_options::DockerOptions; /// This command is used to check hifigan for being a real model diff --git a/crates/service/job/download_job/src/job_types/hifigan_softvc/hifigan_softvc_model_check_command.rs b/crates/service/job/download_job/src/job_types/hifigan_softvc/hifigan_softvc_model_check_command.rs new file mode 100644 index 0000000000..3cf34bb129 --- /dev/null +++ b/crates/service/job/download_job/src/job_types/hifigan_softvc/hifigan_softvc_model_check_command.rs @@ -0,0 +1,84 @@ +use crate::AnyhowResult; +use log::info; +use std::path::Path; +use subprocess::{Popen, PopenConfig}; +use subprocess_common::docker_options::DockerOptions; + +/// This command is used to check hifigan for being a real model +#[derive(Clone)] +pub struct HifiGanSoftVcModelCheckCommand { + /// Where the HifiGan code lives + hifigan_softvc_root_code_directory: String, + + /// eg. `source python/bin/activate` + maybe_virtual_env_activation_command: Option, + + hifigan_softvc_model_check_script_name: String, + + /// If this is run under Docker (eg. in development), these are the options. + maybe_docker_options: Option, +} + +impl HifiGanSoftVcModelCheckCommand { + pub fn new( + hifigan_root_code_directory: &str, + maybe_virtual_env_activation_command: Option<&str>, + hifigan_model_check_script_name: &str, + maybe_docker_options: Option, + ) -> Self { + Self { + hifigan_softvc_root_code_directory: hifigan_root_code_directory.to_string(), + maybe_virtual_env_activation_command: maybe_virtual_env_activation_command.map(|s| s.to_string()), + hifigan_softvc_model_check_script_name: hifigan_model_check_script_name.to_string(), + maybe_docker_options, + } + } + + pub fn execute>( + &self, + checkpoint_path: P, + output_metadata_filename: P, + ) -> AnyhowResult<()> { + + let mut command = String::new(); + command.push_str(&format!("cd {}", self.hifigan_softvc_root_code_directory)); + + if let Some(venv_command) = self.maybe_virtual_env_activation_command.as_deref() { + command.push_str(" && "); + command.push_str(venv_command); + command.push_str(" "); + } + + command.push_str(" && "); + command.push_str("python "); + command.push_str(&self.hifigan_softvc_model_check_script_name); + command.push_str(" --checkpoint_path "); + command.push_str(&checkpoint_path.as_ref().display().to_string()); + command.push_str(" --output_metadata_filename "); + command.push_str(&output_metadata_filename.as_ref().display().to_string()); + + if let Some(docker_options) = self.maybe_docker_options.as_ref() { + command = docker_options.to_command_string(&command); + } + + info!("Command: {:?}", command); + + let command_parts = [ + "bash", + "-c", + &command + ]; + + let mut p = Popen::create(&command_parts, PopenConfig { + ..Default::default() + })?; + + info!("Pid : {:?}", p.pid()); + + let exit_status = p.wait()?; + + info!("Exit status: {:?}", exit_status); + + Ok(()) + } +} diff --git a/crates/service/job/download_job/src/job_types/hifigan_softvc/mod.rs b/crates/service/job/download_job/src/job_types/hifigan_softvc/mod.rs new file mode 100644 index 0000000000..144e48ea3a --- /dev/null +++ b/crates/service/job/download_job/src/job_types/hifigan_softvc/mod.rs @@ -0,0 +1,2 @@ +pub mod hifigan_softvc_model_check_command; +pub mod process_hifigan_softvc_vocoder; diff --git a/crates/service/job/download_job/src/job_types/hifigan_softvc/process_hifigan_softvc_vocoder.rs b/crates/service/job/download_job/src/job_types/hifigan_softvc/process_hifigan_softvc_vocoder.rs new file mode 100644 index 0000000000..a2af6ca215 --- /dev/null +++ b/crates/service/job/download_job/src/job_types/hifigan_softvc/process_hifigan_softvc_vocoder.rs @@ -0,0 +1,134 @@ +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use container_common::filesystem::check_file_exists::check_file_exists; +use container_common::filesystem::safe_delete_temp_directory::safe_delete_temp_directory; +use container_common::filesystem::safe_delete_temp_file::safe_delete_temp_file; +use container_common::hashing::hash_file_sha2::hash_file_sha2; +use crate::JobState; +use crate::job_steps::job_results::JobResults; +use database_queries::queries::generic_download::job::list_available_generic_download_jobs::AvailableDownloadJob; +use database_queries::queries::vocoder::insert_vocoder_model::{Args, insert_vocoder_model}; +use jobs_common::redis_job_status_logger::RedisJobStatusLogger; +use log::{info, warn}; +use reusable_types::db::enums::vocoder_type::VocoderType; +use std::fs::File; +use std::io::Read; +use std::path::PathBuf; +use tempdir::TempDir; + +/// Returns the token of the entity. +pub async fn process_hifigan_softvc_vocoder<'a, 'b>( + job_state: &JobState, + job: &AvailableDownloadJob, + temp_dir: &TempDir, + download_filename: &str, + redis_logger: &'a mut RedisJobStatusLogger<'b>, +) -> AnyhowResult { + + // ==================== RUN MODEL CHECK ==================== // + + info!("Checking that model is valid..."); + + redis_logger.log_status("checking hifigan (softvc) model")?; + + let file_path = PathBuf::from(download_filename.clone()); + + let output_metadata_fs_path = temp_dir.path().join("metadata.json"); + + let model_check_result = job_state.sidecar_configs.hifigan_model_check_command.execute( + &file_path, + &output_metadata_fs_path + ); + + if let Err(e) = model_check_result { + safe_delete_temp_file(&file_path); + safe_delete_temp_directory(&temp_dir); + return Err(anyhow!("model check error: {:?}", e)); + } + + // ==================== CHECK ALL FILES EXIST AND GET METADATA ==================== // + + info!("Checking that metadata output file exists..."); + + check_file_exists(&output_metadata_fs_path)?; + + let file_metadata = match read_metadata_file(&output_metadata_fs_path) { + Ok(metadata) => metadata, + Err(e) => { + safe_delete_temp_file(&file_path); + safe_delete_temp_file(&output_metadata_fs_path); + safe_delete_temp_directory(&temp_dir); + return Err(e); + } + }; + + // ==================== UPLOAD MODEL FILE ==================== // + + info!("Uploading HifiGan (softvc) vocoder to GCS..."); + + let private_bucket_hash = hash_file_sha2(&download_filename)?; + + info!("File hash: {}", private_bucket_hash); + + let model_bucket_path = job_state.bucket_path_unifier.vocoder_path(&private_bucket_hash); + + info!("Destination bucket path: {:?}", &model_bucket_path); + + redis_logger.log_status("uploading hifigan model")?; + + if let Err(e) = job_state.bucket_client.upload_filename(&model_bucket_path, &file_path).await { + safe_delete_temp_file(&output_metadata_fs_path); + safe_delete_temp_file(&file_path); + safe_delete_temp_directory(&temp_dir); + return Err(e); + } + + // ==================== DELETE DOWNLOADED FILE ==================== // + + // NB: We should be using a tempdir, but to make absolutely certain we don't overflow the disk... + safe_delete_temp_file(&output_metadata_fs_path); + safe_delete_temp_file(&file_path); + safe_delete_temp_directory(&temp_dir); + + // ==================== SAVE RECORDS ==================== // + + info!("Saving model record..."); + let (_id, model_token) = insert_vocoder_model(Args { + vocoder_type: VocoderType::HifiGanRocketVc, // NB: "rocket_vc" is an internal codename for softvc. + title: &job.title, + original_download_url: &job.download_url, + original_filename: &download_filename, + file_size_bytes: file_metadata.file_size_bytes, + creator_user_token: &job.creator_user_token, + creator_ip_address: &job.creator_ip_address, + creator_set_visibility: job.creator_set_visibility, + private_bucket_hash: &private_bucket_hash, + private_bucket_object_name: &model_bucket_path, + mysql_pool: &job_state.mysql_pool + }).await?; + + job_state.badge_granter.maybe_grant_softvc_vocoder_model_uploads_badge(&job.creator_user_token) + .await + .map_err(|e| { + warn!("error maybe awarding badge: {:?}", e); + anyhow!("error maybe awarding badge") + })?; + + Ok(JobResults { + entity_token: Some(model_token), + entity_type: Some("hifigan_rocket_vc".to_string()), // NB: This may be different from `GenericDownloadType` in the future! + }) +} + +#[derive(Deserialize)] +struct FileMetadata { + pub file_size_bytes: u64, +} + +fn read_metadata_file(filename: &PathBuf) -> AnyhowResult { + let mut file = File::open(filename)?; + let mut buffer = String::new(); + file.read_to_string(&mut buffer)?; + Ok(serde_json::from_str(&buffer)?) +} + diff --git a/crates/service/job/download_job/src/job_types/mod.rs b/crates/service/job/download_job/src/job_types/mod.rs index 9f7e3cc1f8..c4ca220216 100644 --- a/crates/service/job/download_job/src/job_types/mod.rs +++ b/crates/service/job/download_job/src/job_types/mod.rs @@ -1,3 +1,4 @@ pub mod hifigan; +pub mod hifigan_softvc; pub mod tacotron; pub mod talknet; \ No newline at end of file