New endpoint to return base64 spectrogram + audio

Cargo.lock

@@ -408,6 +408,11 @@ name = "base64"
 version = "0.11.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 
+[[package]]
+name = "base64"
+version = "0.12.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+
 [[package]]
 name = "bindgen"
 version = "0.51.1"
@@ -2419,6 +2424,7 @@ dependencies = [
  "actix-web 2.0.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "anyhow 1.0.30 (registry+https://github.com/rust-lang/crates.io-index)",
  "arpabet 1.0.0 (registry+https://github.com/rust-lang/crates.io-index)",
+ "base64 0.12.2 (registry+https://github.com/rust-lang/crates.io-index)",
  "chrono 0.4.11 (registry+https://github.com/rust-lang/crates.io-index)",
  "diesel 1.4.4 (registry+https://github.com/rust-lang/crates.io-index)",
  "env_logger 0.7.1 (registry+https://github.com/rust-lang/crates.io-index)",
@@ -2742,6 +2748,7 @@ dependencies = [
 "checksum backtrace 0.3.40 (registry+https://github.com/rust-lang/crates.io-index)" = "924c76597f0d9ca25d762c25a4d369d51267536465dc5064bdf0eb073ed477ea"
 "checksum backtrace-sys 0.1.32 (registry+https://github.com/rust-lang/crates.io-index)" = "5d6575f128516de27e3ce99689419835fce9643a9b215a14d2b5b685be018491"
 "checksum base64 0.11.0 (registry+https://github.com/rust-lang/crates.io-index)" = "b41b7ea54a0c9d92199de89e20e58d49f02f8e699814ef3fdf266f6f748d15c7"
+"checksum base64 0.12.2 (registry+https://github.com/rust-lang/crates.io-index)" = "e223af0dc48c96d4f8342ec01a4974f139df863896b316681efd36742f22cc67"
 "checksum bindgen 0.51.1 (registry+https://github.com/rust-lang/crates.io-index)" = "ebd71393f1ec0509b553aa012b9b58e81dadbdff7130bd3b8cba576e69b32f75"
 "checksum bitflags 1.2.1 (registry+https://github.com/rust-lang/crates.io-index)" = "cf1de2fe8c75bc145a2f577add951f8134889b4795d47466a54a5c846d691693"
 "checksum brotli-sys 0.3.2 (registry+https://github.com/rust-lang/crates.io-index)" = "4445dea95f4c2b41cde57cc9fee236ae4dbae88d8fcbdb4750fc1bb5d86aaecd"

tts_service/Cargo.toml

@@ -19,6 +19,7 @@ actix-rt = "1.0.0"
 actix-utils = "1.0.3"
 actix-web = "2.0.0"
 anyhow = "1.0.30"
+base64 = "0.12.2"
 chrono = { version = "0.4.11", features = ["serde"] }
 diesel = { version = "1.4.4", features = ["mysql", "r2d2", "chrono"] }
 env_logger = "0.7.1"

tts_service/src/endpoints/mod.rs

@@ -5,5 +5,6 @@ pub mod models;
 pub mod readiness;
 pub mod sentences;
 pub mod speak;
+pub mod speak_with_spectrogram;
 pub mod speakers;
 pub mod tts;

tts_service/src/endpoints/speak_with_spectrogram.rs

@@ -0,0 +1,199 @@
+use actix_web::http::{StatusCode, header};
+use actix_web::web::{
+  Data,
+  Json,
+};
+use actix_web::{
+  Either,
+  HttpRequest,
+  HttpResponse,
+};
+
+use std::sync::Arc;
+use crate::AppState;
+use arpabet::Arpabet;
+use crate::text::arpabet::text_to_arpabet_encoding;
+use crate::model::model_config::ModelPipeline;
+use crate::model::old_model::TacoMelModel;
+use crate::text::cleaners::clean_text;
+use crate::model::pipelines::{arpabet_glow_tts_melgan_pipeline, arpabet_glow_tts_multi_speaker_melgan_pipeline, arpabet_glow_tts_melgan_pipeline_with_spectrogram, arpabet_glow_tts_multi_speaker_melgan_pipeline_with_spectrogram};
+use crate::database::model::NewSentence;
+
+#[derive(Deserialize)]
+pub struct SpeakRequest {
+  /// Slug for the speaker
+  speaker: String,
+  /// Raw text to be spoken
+  text: String,
+}
+
+#[derive(Serialize, Default)]
+pub struct Spectrogram {
+  pub bytes_base64: String,
+  pub width: i64,
+  pub height: i64,
+}
+
+#[derive(Serialize)]
+pub struct SpeakSpectrogramResponse {
+  pub audio_base64: String,
+  pub spectrogram: Spectrogram,
+}
+
+pub async fn post_speak_with_spectrogram(request: HttpRequest,
+  query: Json<SpeakRequest>,
+  app_state: Data<Arc<AppState>>)
+  -> Either<Json<SpeakSpectrogramResponse>, std::io::Result<HttpResponse>>
+{
+  let app_state = app_state.into_inner();
+
+  let speaker_slug = query.speaker.to_string();
+
+  let speaker = match app_state.model_configs.find_speaker_by_slug(&speaker_slug) {
+    Some(speaker) => speaker,
+    None => {
+      return Either::B(Ok(HttpResponse::build(StatusCode::NOT_FOUND)
+          .content_type("text/plain")
+          .body("Speaker not found")));
+    },
+  };
+
+  let text = query.text.to_string();
+
+  if text.is_empty() {
+    return Either::B(Ok(HttpResponse::build(StatusCode::BAD_REQUEST)
+        .content_type("text/plain")
+        .body("Request has empty text.")));
+  }
+
+  // NB: Actually, we want the X-Forwarded-For IP address, since otherwise
+  // we get the load balancer.
+  let ip_address = request.connection_info()
+      .remote()
+      .unwrap_or("")
+      .to_string();
+
+  let sentence_record = NewSentence {
+    sentence: text.clone(),
+    speaker: speaker_slug.clone(),
+    ip_address: ip_address,
+  };
+
+  match sentence_record.insert(&app_state.database_connector) {
+    Err(_) => error!("Could not insert sentence record for: {:?}", sentence_record),
+    Ok(_) => {},
+  }
+
+  let cleaned_text = clean_text(&text);
+
+  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");
+
+      info!("Tacotron Model: {}", tacotron_model);
+      info!("Melgan Model: {}", melgan_model);
+      info!("Text: {}", text);
+
+      let arpabet = Arpabet::load_cmudict();
+      let encoded = text_to_arpabet_encoding(arpabet, &cleaned_text);
+
+      let tacotron = app_state.model_cache.get_or_load_arbabet_tacotron(&tacotron_model)
+          .expect(&format!("Couldn't load tacotron model: {}", tacotron_model));
+
+      let melgan = app_state.model_cache.get_or_load_melgan(&melgan_model)
+          .expect(&format!("Couldn't load melgan model: {}", melgan_model));
+
+      match TacoMelModel::new().run_tts_encoded(&tacotron, &melgan, &encoded) {
+        None => {
+          Either::B(Ok(HttpResponse::build(StatusCode::TOO_MANY_REQUESTS)
+              .content_type("text/plain")
+              .body("The service is receiving too many requests. Although there are many worker \
+                     containers, model access is serialized on a per-container basis until the \
+                     segfaults are fixed.")))
+        },
+        Some(wav_data) => {
+          Either::A(Json(SpeakSpectrogramResponse {
+            audio_base64: base64::encode(wav_data),
+            spectrogram: Default::default()
+          }))
+        },
+      }
+    },
+    ModelPipeline::ArpabetGlowTtsMelgan => {
+      let glow_tts_model = speaker.glow_tts
+          .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");
+
+      info!("Glow-TTS Model: {}", glow_tts_model);
+      info!("Melgan Model: {}", melgan_model);
+      info!("Text: {}", text);
+
+      let glow_tts = app_state.model_cache.get_or_load_arbabet_glow_tts(&glow_tts_model)
+          .expect(&format!("Couldn't load glow-tts model: {}", glow_tts_model));
+
+      let melgan = app_state.model_cache.get_or_load_melgan(&melgan_model)
+          .expect(&format!("Couldn't load melgan model: {}", melgan_model));
+
+      let (mel_data, wav_data) = arpabet_glow_tts_melgan_pipeline_with_spectrogram(&cleaned_text, &glow_tts, &melgan);
+
+      println!("Mel Data: {}x{}", mel_data.width, mel_data.height);
+
+      Either::A(Json(SpeakSpectrogramResponse {
+        audio_base64: base64::encode(wav_data),
+        spectrogram: mel_data,
+      }))
+    },
+    ModelPipeline::ArpabetGlowTtsMultiSpeakerMelgan=> {
+      let glow_tts_multi_speaker_model = speaker.glow_tts_multi_speaker
+          .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");
+
+      info!("Glow-TTS Multi-Speaker Model: {}", glow_tts_multi_speaker_model);
+      info!("Melgan Model: {}", melgan_model);
+      info!("Text: {}", text);
+
+      let speaker_id = speaker.speaker_id.expect("Should have speaker_id");
+
+      let glow_tts_multi_speaker = app_state
+          .model_cache
+          .get_or_load_arbabet_glow_tts_multi_speaker(&glow_tts_multi_speaker_model)
+          .expect(&format!("Couldn't load glow-tts multi-speaker model: {}", glow_tts_multi_speaker_model));
+
+      let melgan = app_state
+          .model_cache
+          .get_or_load_melgan(&melgan_model)
+          .expect(&format!("Couldn't load melgan model: {}", melgan_model));
+
+      let (mel_data, wav_data) = arpabet_glow_tts_multi_speaker_melgan_pipeline_with_spectrogram(
+        &cleaned_text, speaker_id, &glow_tts_multi_speaker, &melgan);
+
+      println!("Mel Data: {}x{}", mel_data.width, mel_data.height);
+
+      Either::A(Json(SpeakSpectrogramResponse {
+        audio_base64: base64::encode(wav_data),
+        spectrogram: mel_data,
+      }))
+    },
+    ModelPipeline::RawTextTacotronMelgan => unimplemented!(),
+  }
+}

tts_service/src/main.rs

@@ -41,6 +41,7 @@ use crate::endpoints::tts::post_tts;
 use crate::model::model_cache::ModelCache;
 use crate::model::model_config::ModelConfigs;
 use crate::text::checker::TextChecker;
+use crate::endpoints::speak_with_spectrogram::post_speak_with_spectrogram;
 
 const ENV_ASSET_DIRECTORY: &'static str = "ASSET_DIRECTORY";
 const ENV_BIND_ADDRESS: &'static str = "BIND_ADDRESS";
@@ -233,6 +234,11 @@ async fn run_server(app_state: AppState, server_args: ServerArgs) -> std::io::Re
             .route(web::get().to(legacy_get_speak))
             .route(web::head().to(|| HttpResponse::Ok()))
       )
+      .service(
+        web::resource("/speak_spectrogram")
+            .route(web::post().to(post_speak_with_spectrogram))
+            .route(web::head().to(|| HttpResponse::Ok()))
+      )
       .service(get_root)
       .service(get_readiness)
       .service(get_liveness)

tts_service/src/model/pipelines.rs

@@ -10,6 +10,7 @@ use std::io::{Cursor, BufWriter};
 use tch::Tensor;
 use crate::model::arpabet_glow_tts_model::ArpabetGlowTtsModel;
 use crate::model::arpabet_glow_tts_multi_speaker_model::ArpabetGlowTtsMultiSpeakerModel;
+use crate::endpoints::speak_with_spectrogram::Spectrogram;
 
 // TODO: This might be useful to implement as a multi-stage
 //  state machine struct with functions, that way you can pull out
@@ -32,6 +33,61 @@ pub fn arpabet_glow_tts_multi_speaker_melgan_pipeline(
   audio_signal_to_wav_bytes(audio_signal)
 }
 
+pub fn arpabet_glow_tts_multi_speaker_melgan_pipeline_with_spectrogram(
+  cleaned_text: &str,
+  speaker_id: i64,
+  arpabet_glow_tts: &ArpabetGlowTtsMultiSpeakerModel,
+  melgan: &MelganModel) -> (Spectrogram, Vec<u8>) {
+
+  let arpabet = Arpabet::load_cmudict(); // TODO: Inefficient.
+  let arpabet_encodings = text_to_arpabet_encoding_glow_tts(arpabet, &cleaned_text);
+
+  let mel_tensor = arpabet_glow_tts.encoded_arpabet_to_mel(&arpabet_encodings, speaker_id);
+
+  let spectrogram = {
+    let dims = mel_tensor.size3().expect("This should work");
+
+    let mut spectrogram = Spectrogram::default();
+    spectrogram.height = dims.1;
+    spectrogram.width = dims.2;
+
+    let length = dims.0 * dims.1 * dims.2;
+    let mut data = [0.0f32].repeat(length as usize);
+    mel_tensor.copy_data(data.as_mut_slice(), length as usize);
+
+    // NB: The real maxima and minima are +/-inf, but I don't want to deal with that.
+    let mut max_value = -100000.0f32;
+    let mut min_value = 100000.0f32;
+
+    for sample in data.iter() {
+      if *sample < min_value {
+        min_value = *sample;
+      } else if *sample > max_value {
+        max_value = *sample;
+      }
+    }
+
+    let mut scaled_data = Vec::with_capacity(data.len());
+    for sample in data {
+      let scaled_sample = ( (255.0f32 - 0.0f32) * (sample - min_value) ) / ( max_value - min_value );
+      scaled_data.push(scaled_sample);
+    }
+
+    let bytes : Vec<u8> = scaled_data.iter().map(|s| *s as u8).collect();
+
+    spectrogram.bytes_base64 = base64::encode(bytes);
+    spectrogram
+  };
+
+  let audio_tensor = melgan.tacotron_mel_to_audio(&mel_tensor);
+  let audio_signal = mel_audio_tensor_to_audio_signal(audio_tensor);
+
+  (
+    spectrogram,
+    audio_signal_to_wav_bytes(audio_signal),
+  )
+}
+
 pub fn arpabet_glow_tts_melgan_pipeline(
   cleaned_text: &str,
   arpabet_glow_tts: &ArpabetGlowTtsModel,
@@ -48,6 +104,60 @@ pub fn arpabet_glow_tts_melgan_pipeline(
   audio_signal_to_wav_bytes(audio_signal)
 }
 
+pub fn arpabet_glow_tts_melgan_pipeline_with_spectrogram(
+  cleaned_text: &str,
+  arpabet_glow_tts: &ArpabetGlowTtsModel,
+  melgan: &MelganModel) -> (Spectrogram, Vec<u8>) {
+
+  let arpabet = Arpabet::load_cmudict(); // TODO: Inefficient.
+  let arpabet_encodings = text_to_arpabet_encoding_glow_tts(arpabet, &cleaned_text);
+
+  let mel_tensor = arpabet_glow_tts.encoded_arpabet_to_mel(&arpabet_encodings);
+
+  let spectrogram = {
+    let dims = mel_tensor.size3().expect("This should work");
+
+    let mut spectrogram = Spectrogram::default();
+    spectrogram.height = dims.1;
+    spectrogram.width = dims.2;
+
+    let length = dims.0 * dims.1 * dims.2;
+    let mut data = [0.0f32].repeat(length as usize);
+    mel_tensor.copy_data(data.as_mut_slice(), length as usize);
+
+    // NB: The real maxima and minima are +/-inf, but I don't want to deal with that.
+    let mut max_value = -100000.0f32;
+    let mut min_value = 100000.0f32;
+
+    for sample in data.iter() {
+      if *sample < min_value {
+        min_value = *sample;
+      } else if *sample > max_value {
+        max_value = *sample;
+      }
+    }
+
+    let mut scaled_data = Vec::with_capacity(data.len());
+    for sample in data {
+      let scaled_sample = ( (255.0f32 - 0.0f32) * (sample - min_value) ) / ( max_value - min_value );
+      scaled_data.push(scaled_sample);
+    }
+
+    let bytes : Vec<u8> = scaled_data.iter().map(|s| *s as u8).collect();
+
+    spectrogram.bytes_base64 = base64::encode(bytes);
+    spectrogram
+  };
+
+  let audio_tensor = melgan.tacotron_mel_to_audio(&mel_tensor);
+  let audio_signal = mel_audio_tensor_to_audio_signal(audio_tensor);
+
+  (
+    spectrogram,
+    audio_signal_to_wav_bytes(audio_signal),
+  )
+}
+
 pub fn arpabet_tacotron_melgan_pipeline(
   cleaned_text: &str,
   arpabet_tacotron: &ArpabetTacotronModel,