From f87592bfc3eebd1e6577fa93491ca57105e7818e Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Wed, 17 Jun 2020 01:35:34 -0400 Subject: [PATCH] New endpoint to return base64 spectrogram + audio --- Cargo.lock | 7 + tts_service/Cargo.toml | 1 + tts_service/src/endpoints/mod.rs | 1 + .../src/endpoints/speak_with_spectrogram.rs | 199 ++++++++++++++++++ tts_service/src/main.rs | 6 + tts_service/src/model/pipelines.rs | 110 ++++++++++ 6 files changed, 324 insertions(+) create mode 100644 tts_service/src/endpoints/speak_with_spectrogram.rs diff --git a/Cargo.lock b/Cargo.lock index e8727ec58e..cbc6329ea3 100644 --- a/Cargo.lock +++ b/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" diff --git a/tts_service/Cargo.toml b/tts_service/Cargo.toml index 38ccba089b..3ee71c03a4 100644 --- a/tts_service/Cargo.toml +++ b/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" diff --git a/tts_service/src/endpoints/mod.rs b/tts_service/src/endpoints/mod.rs index a4d0119002..f499c9fdc5 100644 --- a/tts_service/src/endpoints/mod.rs +++ b/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; \ No newline at end of file diff --git a/tts_service/src/endpoints/speak_with_spectrogram.rs b/tts_service/src/endpoints/speak_with_spectrogram.rs new file mode 100644 index 0000000000..2208dd7003 --- /dev/null +++ b/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, + app_state: Data>) + -> Either, std::io::Result> +{ + 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!(), + } +} diff --git a/tts_service/src/main.rs b/tts_service/src/main.rs index 30219bf5d8..0c1c21d410 100644 --- a/tts_service/src/main.rs +++ b/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) diff --git a/tts_service/src/model/pipelines.rs b/tts_service/src/model/pipelines.rs index 94535781e8..85fa75915a 100644 --- a/tts_service/src/model/pipelines.rs +++ b/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) { + + 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 = 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) { + + 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 = 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,