From b3dedfc5a8d80b4d0cbaf667cebfbca5d9495636 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Sun, 2 Apr 2023 21:05:05 -0700 Subject: [PATCH] WIP VITS uploads --- .../src/workers/generic_download_type.rs | 10 +++ .../endpoints/tts/enqueue_upload_tts_model.rs | 86 ++++++++++++------- 2 files changed, 63 insertions(+), 33 deletions(-) diff --git a/crates/_types/enums/src/workers/generic_download_type.rs b/crates/_types/enums/src/workers/generic_download_type.rs index 87810a635a..b42f078ee4 100644 --- a/crates/_types/enums/src/workers/generic_download_type.rs +++ b/crates/_types/enums/src/workers/generic_download_type.rs @@ -36,6 +36,11 @@ pub enum GenericDownloadType { #[serde(rename = "tacotron2")] #[sqlx(rename = "tacotron2")] Tacotron2, + + /// VITS TTS models. + #[serde(rename = "vits")] + #[sqlx(rename = "vits")] + Vits, } /// NB: Legacy API for older code. @@ -46,6 +51,7 @@ impl GenericDownloadType { Self::HifiGanRocketVc => "hifigan_rocket_vc", Self::RocketVc => "rocket_vc", Self::Tacotron2 => "tacotron2", + Self::Vits => "vits", } } @@ -55,6 +61,7 @@ impl GenericDownloadType { "hifigan_rocket_vc" => Ok(Self::HifiGanRocketVc), "rocket_vc" => Ok(Self::RocketVc), "tacotron2" => Ok(Self::Tacotron2), + "vits" => Ok(Self::Vits), _ => Err(format!("invalid value: {:?}", value)), } } @@ -71,6 +78,7 @@ mod tests { assert_serialization(GenericDownloadType::HifiGanRocketVc, "hifigan_rocket_vc"); assert_serialization(GenericDownloadType::RocketVc, "rocket_vc"); assert_serialization(GenericDownloadType::Tacotron2, "tacotron2"); + assert_serialization(GenericDownloadType::Vits, "vits"); } #[test] @@ -79,6 +87,7 @@ mod tests { assert_eq!(GenericDownloadType::HifiGanRocketVc.to_str(), "hifigan_rocket_vc"); assert_eq!(GenericDownloadType::RocketVc.to_str(), "rocket_vc"); assert_eq!(GenericDownloadType::Tacotron2.to_str(), "tacotron2"); + assert_eq!(GenericDownloadType::Vits.to_str(), "vits"); } #[test] @@ -87,5 +96,6 @@ mod tests { assert_eq!(GenericDownloadType::from_str("hifigan_rocket_vc").unwrap(), GenericDownloadType::HifiGanRocketVc); assert_eq!(GenericDownloadType::from_str("rocket_vc").unwrap(), GenericDownloadType::RocketVc); assert_eq!(GenericDownloadType::from_str("tacotron2").unwrap(), GenericDownloadType::Tacotron2); + assert_eq!(GenericDownloadType::from_str("vits").unwrap(), GenericDownloadType::Vits); } } diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/tts/enqueue_upload_tts_model.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/tts/enqueue_upload_tts_model.rs index 1b7560f225..50b346da9f 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/tts/enqueue_upload_tts_model.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/tts/enqueue_upload_tts_model.rs @@ -16,17 +16,23 @@ use http_server_common::request::get_request_ip::get_request_ip; use log::warn; use std::fmt; use std::sync::Arc; +use enums::workers::generic_download_type::GenericDownloadType; #[derive(Deserialize)] -pub enum TtsModelType { +pub enum SupportedTtsModelType { /// tacotron2 Tacotron2, - /// glowtts - GlowTts, - /// glowtts-vocodes - GlowTts_Vocodes, - /// talknet - Talknet, + + /// VITS + Vits, + + // Not yet supported: + // /// glowtts + // GlowTts, + // /// glowtts-vocodes + // GlowTts_Vocodes, + // /// talknet + // Talknet, } #[derive(Deserialize)] @@ -34,7 +40,7 @@ pub struct UploadTtsModelRequest { idempotency_token: String, title: String, download_url: String, - tts_model_type: Option, + tts_model_type: Option, creator_set_visibility: Option, } @@ -133,35 +139,49 @@ pub async fn upload_tts_model_handler( } } - let tts_model_type = "tacotron2".to_string(); - let creator_set_visibility = "public".to_string(); + let mut supported_tts_model_type = request.tts_model_type + .clone() + .unwrap_or(SupportedTtsModelType::Tacotron2); - // This token is returned to the client. - let maybe_new_job_token = insert_tts_model_upload_job(InsertTtsModelUploadJobArgs { - uuid: &uuid, - creator_user_token: &user_session.user_token, - creator_ip_address: &ip_address, - creator_set_visibility: &creator_set_visibility, - title: &title, - tts_model_type: &tts_model_type, - download_url: &download_url, - mysql_pool: &server_state.mysql_pool, - }).await; + let mut job_token = ""; - let job_token = match maybe_new_job_token { - Ok(token) => token, - Err(err) => { - warn!("Error inserting new job record: {:?}", err); - return Err(UploadTtsModelError::ServerError); + match supported_tts_model_type { + SupportedTtsModelType::Tacotron2 => { + + let tts_model_type = "tacotron2".to_string(); + let creator_set_visibility = "public".to_string(); + + // This token is returned to the client. + let maybe_new_job_token = insert_tts_model_upload_job(InsertTtsModelUploadJobArgs { + uuid: &uuid, + creator_user_token: &user_session.user_token, + creator_ip_address: &ip_address, + creator_set_visibility: &creator_set_visibility, + title: &title, + tts_model_type: &tts_model_type, + download_url: &download_url, + mysql_pool: &server_state.mysql_pool, + }).await; + + let job_token = match maybe_new_job_token { + Ok(token) => token, + Err(err) => { + warn!("Error inserting new job record: {:?}", err); + return Err(UploadTtsModelError::ServerError); + } + }; + + server_state.firehose_publisher.enqueue_tts_model_upload(&user_session.user_token, &job_token) + .await + .map_err(|e| { + warn!("error publishing event: {:?}", e); + UploadTtsModelError::ServerError + })?; } - }; + SupportedTtsModelType::Vits => { - server_state.firehose_publisher.enqueue_tts_model_upload(&user_session.user_token, &job_token) - .await - .map_err(|e| { - warn!("error publishing event: {:?}", e); - UploadTtsModelError::ServerError - })?; + } + } let response = UploadTtsModelSuccessResponse { success: true,