From 450c4ebe0a711e7e505bfcc15e39da90bca665f9 Mon Sep 17 00:00:00 2001
From: Brandon Thomas <echelon@gmail.com>
Date: Tue, 26 May 2020 20:57:40 -0400
Subject: [PATCH] Add /speak and /speakers endpoints

---
 tts_service/src/config.rs             |  7 ++-
 tts_service/src/endpoints/mod.rs      |  4 +-
 tts_service/src/endpoints/speak.rs    | 77 +++++++++++++++++++++++++++
 tts_service/src/endpoints/speakers.rs | 30 +++++++++++
 tts_service/src/main.rs               |  8 +++
 5 files changed, 124 insertions(+), 2 deletions(-)
 create mode 100644 tts_service/src/endpoints/speak.rs
 create mode 100644 tts_service/src/endpoints/speakers.rs

diff --git a/tts_service/src/config.rs b/tts_service/src/config.rs
index 76f4ef3ca9..ad9ccab453 100644
--- a/tts_service/src/config.rs
+++ b/tts_service/src/config.rs
@@ -4,7 +4,7 @@ use std::path::Path;
 
 #[derive(Serialize, Deserialize, Debug, Clone)]
 pub struct ModelConfigs {
-  speakers: Vec<Speaker>,
+  pub speakers: Vec<Speaker>,
   pub model_locations: Vec<ModelLocation>,
   tacotron: Vec<ModelDetails>,
   melgan: Vec<ModelDetails>,
@@ -84,4 +84,9 @@ impl ModelConfigs {
           .map(|s| s.to_string());
     }
   }
+
+  pub fn find_speaker_by_slug(&self, slug: &str) -> Option<&Speaker> {
+    self.speakers.iter()
+        .find(|speaker| speaker.slug == slug)
+  }
 }
diff --git a/tts_service/src/endpoints/mod.rs b/tts_service/src/endpoints/mod.rs
index d39adab1bc..e9865c74aa 100644
--- a/tts_service/src/endpoints/mod.rs
+++ b/tts_service/src/endpoints/mod.rs
@@ -2,4 +2,6 @@ pub mod index;
 pub mod liveness;
 pub mod models;
 pub mod readiness;
-pub mod tts;
\ No newline at end of file
+pub mod speak;
+pub mod speakers;
+pub mod tts;
diff --git a/tts_service/src/endpoints/speak.rs b/tts_service/src/endpoints/speak.rs
new file mode 100644
index 0000000000..93e67f89e6
--- /dev/null
+++ b/tts_service/src/endpoints/speak.rs
@@ -0,0 +1,77 @@
+use actix_web::http::StatusCode;
+use actix_web::web::{
+  Data,
+  Json,
+};
+use actix_web::{
+  HttpRequest,
+  HttpResponse,
+  get,
+};
+
+use std::sync::Arc;
+use crate::AppState;
+use arpabet::Arpabet;
+use crate::text::text_to_arpabet_encoding;
+use crate::old_model::TacoMelModel;
+use crate::config::{Speaker, ModelPipeline};
+
+#[derive(Deserialize)]
+pub struct SpeakRequest {
+  /// Slug for the speaker
+  speaker: String,
+  /// Raw text to be spoken
+  text: String,
+}
+
+pub async fn post_speak(_request: HttpRequest,
+  query: Json<SpeakRequest>,
+  app_state: Data<Arc<AppState>>
+) -> std::io::Result<HttpResponse> {
+  println!("POST /tts");
+
+  let mut app_state = app_state.into_inner();
+
+  let speaker = match app_state.model_configs.find_speaker_by_slug(&query.speaker) {
+    Some(speaker) => speaker,
+    None => {
+      return Ok(HttpResponse::build(StatusCode::NOT_FOUND)
+          .content_type("text/plain")
+          .body("Speaker not found"));
+    },
+  };
+
+  match speaker.model_pipeline {
+    ModelPipeline::ArpabetTacotronMelgan => {
+      let tacotron_model = speaker.tacotron
+          .as_ref()
+          .map(|s| s.clone())
+          .expect("TODO ERROR HANDLING");
+
+      let melgan_model = speaker.melgan
+          .as_ref()
+          .map(|s| s.clone())
+          .expect("TODO ERROR HANDLING");
+
+      let text = query.text.to_string();
+      println!("Tacotron Model: {}", tacotron_model);
+      println!("Melgan Model: {}", melgan_model);
+      println!("Text: {}", text);
+
+      let arpabet = Arpabet::load_cmudict();
+      let encoded = text_to_arpabet_encoding(arpabet, &text);
+
+      println!("Encoded Text: {:?}", encoded);
+
+      let tacotron = app_state.model_cache.get_or_load_arbabet_tacotron(&tacotron_model).unwrap();
+      let melgan = app_state.model_cache.get_or_load_melgan(&melgan_model).unwrap();
+
+      let wav_data = TacoMelModel::new().run_tts_encoded(&tacotron, &melgan, &encoded);
+
+      return Ok(HttpResponse::build(StatusCode::OK)
+          .content_type("audio/wav")
+          .body(wav_data));
+    },
+    ModelPipeline::RawTextTacotronMelgan => unimplemented!(),
+  }
+}
diff --git a/tts_service/src/endpoints/speakers.rs b/tts_service/src/endpoints/speakers.rs
new file mode 100644
index 0000000000..f81d601636
--- /dev/null
+++ b/tts_service/src/endpoints/speakers.rs
@@ -0,0 +1,30 @@
+use actix_web::web::{
+  Data,
+  Json,
+};
+use actix_web::{
+  HttpRequest,
+  get,
+};
+
+use std::sync::Arc;
+use crate::AppState;
+use crate::config::{ModelConfigs, Speaker};
+
+#[derive(Serialize, Debug, Clone)]
+pub struct SpeakersResult {
+  speakers: Vec<Speaker>,
+}
+
+#[get("/speakers")]
+pub async fn get_speakers(
+  _request: HttpRequest,
+  app_state: Data<Arc<AppState>>
+) -> std::io::Result<Json<SpeakersResult>> {
+  println!("GET /models");
+  let app_state = app_state.into_inner();
+
+  Ok(Json(SpeakersResult {
+    speakers: app_state.model_configs.speakers.clone(),
+  }))
+}
diff --git a/tts_service/src/main.rs b/tts_service/src/main.rs
index f945de98c8..02829bda56 100644
--- a/tts_service/src/main.rs
+++ b/tts_service/src/main.rs
@@ -38,6 +38,8 @@ use crate::endpoints::index::get_root;
 use crate::endpoints::liveness::get_liveness;
 use crate::endpoints::models::get_models;
 use crate::endpoints::readiness::get_readiness;
+use crate::endpoints::speak::post_speak;
+use crate::endpoints::speakers::get_speakers;
 use crate::endpoints::tts::post_tts;
 
 const BIND_ADDRESS : &'static str = "BIND_ADDRESS";
@@ -103,10 +105,16 @@ async fn main() -> std::io::Result<()> {
             .route(web::post().to(post_tts))
             .route(web::head().to(|| HttpResponse::Ok()))
       )
+      .service(
+        web::resource("/speak")
+            .route(web::post().to(post_speak))
+            .route(web::head().to(|| HttpResponse::Ok()))
+      )
       .service(get_root)
       .service(get_readiness)
       .service(get_liveness)
       .service(get_models)
+      .service(get_speakers)
       .app_data(arc.clone())
     )
     .bind(bind_address)?