WIP VITS uploads

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

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<TtsModelType>,
+  tts_model_type: Option<SupportedTtsModelType>,
   creator_set_visibility: Option<Visibility>,
 }
 
@@ -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,