From 9e2a63b17d9dabe60df6003ce6d5bfa071de2c12 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Sun, 21 Jun 2020 07:38:07 -0400 Subject: [PATCH] pipeline integrated into /spectrogram_speak endpoint --- .../src/endpoints/speak_with_spectrogram.rs | 96 +++++++------------ tts_service/src/inference/mod.rs | 8 +- tts_service/src/main.rs | 3 +- 3 files changed, 40 insertions(+), 67 deletions(-) diff --git a/tts_service/src/endpoints/speak_with_spectrogram.rs b/tts_service/src/endpoints/speak_with_spectrogram.rs index 2208dd7003..ef77ec06a4 100644 --- a/tts_service/src/endpoints/speak_with_spectrogram.rs +++ b/tts_service/src/endpoints/speak_with_spectrogram.rs @@ -9,15 +9,19 @@ use actix_web::{ HttpResponse, }; -use std::sync::Arc; -use crate::AppState; use arpabet::Arpabet; -use crate::text::arpabet::text_to_arpabet_encoding; +use crate::AppState; +use crate::database::model::NewSentence; +use crate::inference::inference::InferencePipelineStart; +use crate::inference::pipelines::glowtts_melgan::GlowTtsMelganPipeline; +use crate::inference::pipelines::glowtts_multispeaker_melgan::{GlowTtsMultiSpeakerMelganPipeline, GlowTtsMultiSpeakerMelganPipelineMelDone}; +use crate::inference::spectrogram::Base64MelSpectrogram; 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; +use crate::text::arpabet::text_to_arpabet_encoding; +use crate::text::cleaners::clean_text; +use std::sync::Arc; #[derive(Deserialize)] pub struct SpeakRequest { @@ -37,7 +41,7 @@ pub struct Spectrogram { #[derive(Serialize)] pub struct SpeakSpectrogramResponse { pub audio_base64: String, - pub spectrogram: Spectrogram, + pub spectrogram: Base64MelSpectrogram, } pub async fn post_speak_with_spectrogram(request: HttpRequest, @@ -86,47 +90,7 @@ pub async fn post_speak_with_spectrogram(request: HttpRequest, 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() - })) - }, - } - }, + let (base64_image, base64_audio) = match speaker.model_pipeline { ModelPipeline::ArpabetGlowTtsMelgan => { let glow_tts_model = speaker.glow_tts .as_ref() @@ -148,14 +112,16 @@ pub async fn post_speak_with_spectrogram(request: HttpRequest, 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); + let pipeline = GlowTtsMelganPipeline::new(&glow_tts, &melgan) + .infer_mel(&cleaned_text, 0) + .unwrap() + .infer_audio() + .unwrap(); - println!("Mel Data: {}x{}", mel_data.width, mel_data.height); + let base64_image = pipeline.get_base64_mel_spectrogram().unwrap(); + let base64_audio = pipeline.get_base64_audio().unwrap(); - Either::A(Json(SpeakSpectrogramResponse { - audio_base64: base64::encode(wav_data), - spectrogram: mel_data, - })) + (base64_image, base64_audio) }, ModelPipeline::ArpabetGlowTtsMultiSpeakerMelgan=> { let glow_tts_multi_speaker_model = speaker.glow_tts_multi_speaker @@ -184,16 +150,24 @@ pub async fn post_speak_with_spectrogram(request: HttpRequest, .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); + let pipeline = GlowTtsMultiSpeakerMelganPipeline::new(&glow_tts_multi_speaker, &melgan) + .infer_mel(&cleaned_text, speaker_id) + .unwrap() + .infer_audio() + .unwrap(); - Either::A(Json(SpeakSpectrogramResponse { - audio_base64: base64::encode(wav_data), - spectrogram: mel_data, - })) + let base64_image = pipeline.get_base64_mel_spectrogram().unwrap(); + let base64_audio = pipeline.get_base64_audio().unwrap(); + + (base64_image, base64_audio) }, ModelPipeline::RawTextTacotronMelgan => unimplemented!(), - } + ModelPipeline::ArpabetTacotronMelgan => unimplemented!(), + }; + + Either::A(Json(SpeakSpectrogramResponse { + audio_base64: base64_audio.bytes_base64, + spectrogram: base64_image, + })) } diff --git a/tts_service/src/inference/mod.rs b/tts_service/src/inference/mod.rs index 9b4925d98e..254b9bbfab 100644 --- a/tts_service/src/inference/mod.rs +++ b/tts_service/src/inference/mod.rs @@ -1,4 +1,4 @@ -mod audio; -mod spectrogram; -mod inference; -mod pipelines; \ No newline at end of file +pub mod audio; +pub mod spectrogram; +pub mod inference; +pub mod pipelines; \ No newline at end of file diff --git a/tts_service/src/main.rs b/tts_service/src/main.rs index 73bdc6b0aa..8baa210872 100644 --- a/tts_service/src/main.rs +++ b/tts_service/src/main.rs @@ -1,5 +1,3 @@ -mod inference; - #[macro_use] extern crate anyhow; #[macro_use] extern crate diesel; #[macro_use] extern crate log; @@ -14,6 +12,7 @@ extern crate tch; pub mod database; pub mod endpoints; +pub mod inference; pub mod model; pub mod schema; pub mod text;