From c559fb74aea4d4a683ca91e10add4a387ccaf3d6 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Mon, 22 Jun 2020 01:53:39 -0400 Subject: [PATCH] moved mountains, but now there's Sized to deal with --- .../src/endpoints/speak_with_spectrogram.rs | 11 ++- tts_service/src/inference/inference.rs | 14 ++-- tts_service/src/inference/mod.rs | 6 +- .../src/inference/pipelines/glowtts_melgan.rs | 77 +++++++++---------- .../pipelines/glowtts_multispeaker_melgan.rs | 65 ++++++++-------- tts_service/src/inference/tts_model.rs | 7 ++ tts_service/src/inference/vocoder_model.rs | 5 ++ .../src/model/arpabet_glow_tts_model.rs | 11 +++ .../arpabet_glow_tts_multi_speaker_model.rs | 11 +++ tts_service/src/model/melgan_model.rs | 7 ++ 10 files changed, 132 insertions(+), 82 deletions(-) create mode 100644 tts_service/src/inference/tts_model.rs create mode 100644 tts_service/src/inference/vocoder_model.rs diff --git a/tts_service/src/endpoints/speak_with_spectrogram.rs b/tts_service/src/endpoints/speak_with_spectrogram.rs index b1599b8ba4..0f00bc77a4 100644 --- a/tts_service/src/endpoints/speak_with_spectrogram.rs +++ b/tts_service/src/endpoints/speak_with_spectrogram.rs @@ -22,6 +22,8 @@ use crate::model::pipelines::{arpabet_glow_tts_melgan_pipeline, arpabet_glow_tts use crate::text::arpabet::text_to_arpabet_encoding; use crate::text::cleaners::clean_text; use std::sync::Arc; +use crate::inference::vocoder_model::VocoderModelT; +use crate::inference::tts_model::TtsModelT; #[derive(Deserialize)] pub struct SpeakRequest { @@ -88,7 +90,8 @@ pub async fn post_speak_with_spectrogram(request: HttpRequest, Ok(_) => {}, } - let pipeline : Box = match speaker.model_pipeline { + //let pipeline : Box> = match speaker.model_pipeline { + let pipeline : Box, VocoderModel = Arc>> = match speaker.model_pipeline { ModelPipeline::ArpabetGlowTtsMelgan => { let glow_tts_model = speaker.glow_tts .as_ref() @@ -110,7 +113,8 @@ 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)); - Box::new(GlowTtsMelganPipeline::new(&glow_tts, &melgan)) + let inner = GlowTtsMelganPipeline::new(glow_tts.clone(), melgan.clone()); + Box::new(inner) }, ModelPipeline::ArpabetGlowTtsMultiSpeakerMelgan=> { let glow_tts_multi_speaker_model = speaker.glow_tts_multi_speaker @@ -137,7 +141,8 @@ pub async fn post_speak_with_spectrogram(request: HttpRequest, .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)); - Box::new(GlowTtsMultiSpeakerMelganPipeline::new(&glow_tts_multi_speaker, &melgan)) + let inner = GlowTtsMultiSpeakerMelganPipeline::new(glow_tts_multi_speaker.clone(), melgan.clone()); + Box::new(inner) }, ModelPipeline::RawTextTacotronMelgan => unimplemented!(), ModelPipeline::ArpabetTacotronMelgan => unimplemented!(), diff --git a/tts_service/src/inference/inference.rs b/tts_service/src/inference/inference.rs index 9759b792c9..6c019917b1 100644 --- a/tts_service/src/inference/inference.rs +++ b/tts_service/src/inference/inference.rs @@ -4,34 +4,34 @@ use crate::inference::audio::Base64WaveAudio; use crate::inference::pipelines::glowtts_multispeaker_melgan::GlowTtsMultiSpeakerMelganPipelineMelDone; /// 1) Stage of the pipeline before work is done. -pub trait InferencePipelineStart <'a> { +pub trait InferencePipelineStart { type TtsModel; type VocoderModel; fn clean_text(self, text: &str) - -> AnyhowResult + 'a>>; + -> AnyhowResult>>; } /// 2) Stage of the pipeline after text is cleaned. -pub trait InferencePipelineTextCleaningDone <'b> { +pub trait InferencePipelineTextCleaningDone { type TtsModel; type VocoderModel; fn infer_mel(self, speaker_id: i64) - -> AnyhowResult + 'b>>; + -> AnyhowResult>>; } /// 3) Stage of the pipeline after the mel is computed. -pub trait InferencePipelineMelDone <'c> { +pub trait InferencePipelineMelDone { type TtsModel; type VocoderModel; fn infer_audio(self: Box) - -> AnyhowResult + 'c>>; + -> AnyhowResult>>; } /// 4) Stage of the pipeline after the audio is computed. -pub trait InferencePipelineAudioDone <'d> { +pub trait InferencePipelineAudioDone { type TtsModel; type VocoderModel; diff --git a/tts_service/src/inference/mod.rs b/tts_service/src/inference/mod.rs index 254b9bbfab..2888709e31 100644 --- a/tts_service/src/inference/mod.rs +++ b/tts_service/src/inference/mod.rs @@ -1,4 +1,6 @@ pub mod audio; -pub mod spectrogram; pub mod inference; -pub mod pipelines; \ No newline at end of file +pub mod pipelines; +pub mod spectrogram; +pub mod tts_model; +pub mod vocoder_model; \ No newline at end of file diff --git a/tts_service/src/inference/pipelines/glowtts_melgan.rs b/tts_service/src/inference/pipelines/glowtts_melgan.rs index 515c9d81c2..deb76c763a 100644 --- a/tts_service/src/inference/pipelines/glowtts_melgan.rs +++ b/tts_service/src/inference/pipelines/glowtts_melgan.rs @@ -10,36 +10,39 @@ use crate::model::pipelines::{mel_audio_tensor_to_audio_signal, audio_signal_to_ use crate::text::arpabet::text_to_arpabet_encoding_glow_tts; use tch::Tensor; use crate::text::cleaners::clean_text; +use crate::inference::vocoder_model::VocoderModelT; +use crate::inference::tts_model::TtsModelT; +use std::sync::Arc; -pub struct GlowTtsMelganPipeline<'a> { - glow_tts: &'a ArpabetGlowTtsModel, - melgan: &'a MelganModel, +pub struct GlowTtsMelganPipeline { + glow_tts: Arc, + melgan: Arc, } -pub struct GlowTtsMelganPipelineTextCleaningDone<'a> { - glow_tts: &'a ArpabetGlowTtsModel, - melgan: &'a MelganModel, +pub struct GlowTtsMelganPipelineTextCleaningDone { + glow_tts: Arc, + melgan: Arc, cleaned_text: String, } -pub struct GlowTtsMelganPipelineMelDone<'a> { - glow_tts: &'a ArpabetGlowTtsModel, - melgan: &'a MelganModel, +pub struct GlowTtsMelganPipelineMelDone { + glow_tts: Arc, + melgan: Arc, mel_tensor: Tensor, mel_spectrogram: MelSpectrogram, } -pub struct GlowTtsMelganPipelineAudioDone<'a> { - glow_tts: &'a ArpabetGlowTtsModel, - melgan: &'a MelganModel, +pub struct GlowTtsMelganPipelineAudioDone { + glow_tts: Arc, + melgan: Arc, mel_tensor: Tensor, mel_spectrogram: MelSpectrogram, audio_tensor: Tensor, wav_audio_signal: Vec } -impl <'a> GlowTtsMelganPipeline <'a> { - pub fn new(glow_tts: &'a ArpabetGlowTtsModel, melgan: &'a MelganModel) -> Self { +impl GlowTtsMelganPipeline { + pub fn new(glow_tts: Arc, melgan: Arc) -> Self { Self { glow_tts, melgan, @@ -47,12 +50,12 @@ impl <'a> GlowTtsMelganPipeline <'a> { } } -impl <'a> InferencePipelineStart<'a> for GlowTtsMelganPipeline<'a> { - type TtsModel = &'a ArpabetGlowTtsModel; - type VocoderModel = &'a MelganModel; +impl InferencePipelineStart for GlowTtsMelganPipeline { + type TtsModel = Arc; + type VocoderModel = Arc; fn clean_text(self, text: &str) - -> AnyhowResult + 'a>> + -> AnyhowResult>> { let cleaned_text = clean_text(text); @@ -64,18 +67,18 @@ impl <'a> InferencePipelineStart<'a> for GlowTtsMelganPipeline<'a> { } } -impl <'a> InferencePipelineTextCleaningDone<'a> for GlowTtsMelganPipelineTextCleaningDone<'a> { - type TtsModel = &'a ArpabetGlowTtsModel; - type VocoderModel = &'a MelganModel; +impl InferencePipelineTextCleaningDone for GlowTtsMelganPipelineTextCleaningDone { + type TtsModel = Arc; + type VocoderModel = Arc; fn infer_mel(self, _speaker_id: i64) - -> AnyhowResult + 'a>> + -> AnyhowResult>> { // TODO: Creating arpabet instances every time is inefficient (even if lazy_static! under the hood). let arpabet = Arpabet::load_cmudict(); let arpabet_encodings = text_to_arpabet_encoding_glow_tts(arpabet, &self.cleaned_text); - let mel_tensor = self.glow_tts.encoded_arpabet_to_mel(&arpabet_encodings); + let mel_tensor = self.glow_tts.encoded_sequence_to_mel_single_speaker(&arpabet_encodings); let spectrogram = { let dims = mel_tensor.size3().expect("This should work"); @@ -112,45 +115,41 @@ impl <'a> InferencePipelineTextCleaningDone<'a> for GlowTtsMelganPipelineTextCle spectrogram }; - let boxed : Box> = Box::new(GlowTtsMelganPipelineMelDone { + Ok(Box::new(GlowTtsMelganPipelineMelDone { glow_tts: self.glow_tts, melgan: self.melgan, mel_tensor: mel_tensor, mel_spectrogram: spectrogram, - }); - - Ok(boxed) + })) } } -impl <'a> InferencePipelineMelDone<'a> for GlowTtsMelganPipelineMelDone<'a> { - type TtsModel = &'a ArpabetGlowTtsModel; - type VocoderModel = &'a MelganModel; +impl InferencePipelineMelDone for GlowTtsMelganPipelineMelDone { + type TtsModel = Arc; + type VocoderModel = Arc; fn infer_audio(self: Box) - -> AnyhowResult + 'a>> + -> AnyhowResult>> { - let audio_tensor = self.melgan.tacotron_mel_to_audio(&self.mel_tensor); + let audio_tensor = self.melgan.mel_to_audio(&self.mel_tensor); let raw_audio_signal = mel_audio_tensor_to_audio_signal(&audio_tensor); let wav_audio_signal = audio_signal_to_wav_bytes(raw_audio_signal); - let boxed : Box> = Box::new(GlowTtsMelganPipelineAudioDone { + Ok(Box::new(GlowTtsMelganPipelineAudioDone { glow_tts: self.glow_tts, melgan: self.melgan, mel_tensor: self.mel_tensor, mel_spectrogram: self.mel_spectrogram, audio_tensor: audio_tensor, wav_audio_signal: wav_audio_signal, - }); - - Ok(boxed) + })) } } -impl <'a> InferencePipelineAudioDone<'_> for GlowTtsMelganPipelineAudioDone <'a> { - type TtsModel = &'a ArpabetGlowTtsModel; - type VocoderModel = &'a MelganModel; +impl InferencePipelineAudioDone for GlowTtsMelganPipelineAudioDone { + type TtsModel = Arc; + type VocoderModel = Arc; fn get_base64_mel_spectrogram(&self) -> AnyhowResult { let base64_bytes = base64::encode(&self.mel_spectrogram.bytes); diff --git a/tts_service/src/inference/pipelines/glowtts_multispeaker_melgan.rs b/tts_service/src/inference/pipelines/glowtts_multispeaker_melgan.rs index 54d9436b65..54c8e6bb59 100644 --- a/tts_service/src/inference/pipelines/glowtts_multispeaker_melgan.rs +++ b/tts_service/src/inference/pipelines/glowtts_multispeaker_melgan.rs @@ -10,29 +10,32 @@ use crate::text::arpabet::text_to_arpabet_encoding_glow_tts; use tch::Tensor; use crate::model::pipelines::{mel_audio_tensor_to_audio_signal, audio_signal_to_wav_bytes}; use crate::text::cleaners::clean_text; +use crate::inference::vocoder_model::VocoderModelT; +use crate::inference::tts_model::TtsModelT; +use std::sync::Arc; -pub struct GlowTtsMultiSpeakerMelganPipeline<'a> { - glow_tts: &'a ArpabetGlowTtsMultiSpeakerModel, - melgan: &'a MelganModel, +pub struct GlowTtsMultiSpeakerMelganPipeline { + glow_tts: Arc, + melgan: Arc, } -pub struct GlowTtsMultiSpeakerMelganPipelineTextCleaningDone<'a> { - glow_tts: &'a ArpabetGlowTtsMultiSpeakerModel, - melgan: &'a MelganModel, +pub struct GlowTtsMultiSpeakerMelganPipelineTextCleaningDone { + glow_tts: Arc, + melgan: Arc, cleaned_text: String, } -pub struct GlowTtsMultiSpeakerMelganPipelineMelDone<'a> { - glow_tts: &'a ArpabetGlowTtsMultiSpeakerModel, - melgan: &'a MelganModel, +pub struct GlowTtsMultiSpeakerMelganPipelineMelDone { + glow_tts: Arc, + melgan: Arc, cleaned_text: String, mel_tensor: Tensor, mel_spectrogram: MelSpectrogram, } -pub struct GlowTtsMultiSpeakerMelganPipelineAudioDone<'a> { - glow_tts: &'a ArpabetGlowTtsMultiSpeakerModel, - melgan: &'a MelganModel, +pub struct GlowTtsMultiSpeakerMelganPipelineAudioDone { + glow_tts: Arc, + melgan: Arc, cleaned_text: String, mel_tensor: Tensor, mel_spectrogram: MelSpectrogram, @@ -40,8 +43,8 @@ pub struct GlowTtsMultiSpeakerMelganPipelineAudioDone<'a> { wav_audio_signal: Vec } -impl <'a> GlowTtsMultiSpeakerMelganPipeline <'a> { - pub fn new(glow_tts: &'a ArpabetGlowTtsMultiSpeakerModel, melgan: &'a MelganModel) -> Self { +impl GlowTtsMultiSpeakerMelganPipeline { + pub fn new(glow_tts: Arc, melgan: Arc) -> Self { Self { glow_tts, melgan, @@ -49,12 +52,12 @@ impl <'a> GlowTtsMultiSpeakerMelganPipeline <'a> { } } -impl <'a> InferencePipelineStart<'a> for GlowTtsMultiSpeakerMelganPipeline<'a> { - type TtsModel = &'a ArpabetGlowTtsMultiSpeakerModel; - type VocoderModel = &'a MelganModel; +impl InferencePipelineStart for GlowTtsMultiSpeakerMelganPipeline { + type TtsModel = Arc; + type VocoderModel = Arc; fn clean_text(self, text: &str) - -> AnyhowResult + 'a>> + -> AnyhowResult>> { let cleaned_text = clean_text(text); @@ -66,18 +69,18 @@ impl <'a> InferencePipelineStart<'a> for GlowTtsMultiSpeakerMelganPipeline<'a> { } } -impl <'a> InferencePipelineTextCleaningDone<'a> for GlowTtsMultiSpeakerMelganPipelineTextCleaningDone<'a> { - type TtsModel = &'a ArpabetGlowTtsMultiSpeakerModel; - type VocoderModel = &'a MelganModel; +impl InferencePipelineTextCleaningDone for GlowTtsMultiSpeakerMelganPipelineTextCleaningDone { + type TtsModel = Arc; + type VocoderModel = Arc; fn infer_mel(self, speaker_id: i64) - -> AnyhowResult + 'a>> + -> AnyhowResult>> { // TODO: Creating arpabet instances every time is inefficient (even if lazy_static! under the hood). let arpabet = Arpabet::load_cmudict(); let arpabet_encodings = text_to_arpabet_encoding_glow_tts(arpabet, &self.cleaned_text); - let mel_tensor = self.glow_tts.encoded_arpabet_to_mel(&arpabet_encodings, speaker_id); + let mel_tensor = self.glow_tts.encoded_sequence_to_mel(&arpabet_encodings, speaker_id); let spectrogram = { let dims = mel_tensor.size3().expect("This should work"); @@ -124,14 +127,14 @@ impl <'a> InferencePipelineTextCleaningDone<'a> for GlowTtsMultiSpeakerMelganPip } } -impl <'a> InferencePipelineMelDone<'a> for GlowTtsMultiSpeakerMelganPipelineMelDone<'a> { - type TtsModel = &'a ArpabetGlowTtsMultiSpeakerModel; - type VocoderModel = &'a MelganModel; +impl InferencePipelineMelDone for GlowTtsMultiSpeakerMelganPipelineMelDone { + type TtsModel = Arc; + type VocoderModel = Arc; fn infer_audio(self: Box) - -> AnyhowResult + 'a>> + -> AnyhowResult>> { - let audio_tensor = self.melgan.tacotron_mel_to_audio(&self.mel_tensor); + let audio_tensor = self.melgan.mel_to_audio(&self.mel_tensor); let raw_audio_signal = mel_audio_tensor_to_audio_signal(&audio_tensor); let wav_audio_signal = audio_signal_to_wav_bytes(raw_audio_signal); @@ -148,9 +151,9 @@ impl <'a> InferencePipelineMelDone<'a> for GlowTtsMultiSpeakerMelganPipelineMelD } } -impl <'a> InferencePipelineAudioDone<'_> for GlowTtsMultiSpeakerMelganPipelineAudioDone <'a> { - type TtsModel = &'a ArpabetGlowTtsMultiSpeakerModel; - type VocoderModel = &'a MelganModel; +impl InferencePipelineAudioDone for GlowTtsMultiSpeakerMelganPipelineAudioDone { + type TtsModel = Arc; + type VocoderModel = Arc; fn get_base64_mel_spectrogram(&self) -> AnyhowResult { let base64_bytes = base64::encode(&self.mel_spectrogram.bytes); diff --git a/tts_service/src/inference/tts_model.rs b/tts_service/src/inference/tts_model.rs new file mode 100644 index 0000000000..f2e5de0d82 --- /dev/null +++ b/tts_service/src/inference/tts_model.rs @@ -0,0 +1,7 @@ +use tch::Tensor; + +pub trait TtsModelT { + fn encoded_sequence_to_mel(&self, arpabet_encodings: &Vec, speaker_id: i64) -> Tensor; + + fn encoded_sequence_to_mel_single_speaker(&self, arpabet_encodings: &Vec) -> Tensor; +} diff --git a/tts_service/src/inference/vocoder_model.rs b/tts_service/src/inference/vocoder_model.rs new file mode 100644 index 0000000000..3d33d068cc --- /dev/null +++ b/tts_service/src/inference/vocoder_model.rs @@ -0,0 +1,5 @@ +use tch::Tensor; + +pub trait VocoderModelT { + fn mel_to_audio(&self, mel_tensor: &Tensor) -> Tensor; +} \ No newline at end of file diff --git a/tts_service/src/model/arpabet_glow_tts_model.rs b/tts_service/src/model/arpabet_glow_tts_model.rs index a0e37ec6a1..2d90c7a164 100644 --- a/tts_service/src/model/arpabet_glow_tts_model.rs +++ b/tts_service/src/model/arpabet_glow_tts_model.rs @@ -5,6 +5,7 @@ use std::thread; use std::time::Duration; use tch::Tensor; use crate::text::arpabet::text_to_arpabet_encoding_glow_tts; +use crate::inference::tts_model::TtsModelT; pub struct ArpabetGlowTtsModel { model_container: ModelContainer, @@ -28,3 +29,13 @@ impl ArpabetGlowTtsModel { self.model_container.forward2(&text_tensor_2d, &lengths_tensor) } } + +impl TtsModelT for ArpabetGlowTtsModel { + fn encoded_sequence_to_mel(&self, arpabet_encodings: &Vec, _speaker_id: i64) -> Tensor { + self.encoded_arpabet_to_mel(arpabet_encodings) + } + + fn encoded_sequence_to_mel_single_speaker(&self, arpabet_encodings: &Vec) -> Tensor { + self.encoded_arpabet_to_mel(arpabet_encodings) + } +} diff --git a/tts_service/src/model/arpabet_glow_tts_multi_speaker_model.rs b/tts_service/src/model/arpabet_glow_tts_multi_speaker_model.rs index 9d8f1318c7..3b1ae0a8d2 100644 --- a/tts_service/src/model/arpabet_glow_tts_multi_speaker_model.rs +++ b/tts_service/src/model/arpabet_glow_tts_multi_speaker_model.rs @@ -5,6 +5,7 @@ use std::thread; use std::time::Duration; use tch::Tensor; use crate::text::arpabet::text_to_arpabet_encoding_glow_tts; +use crate::inference::tts_model::TtsModelT; pub struct ArpabetGlowTtsMultiSpeakerModel { model_container: ModelContainer, @@ -34,3 +35,13 @@ impl ArpabetGlowTtsMultiSpeakerModel { self.model_container.forward3(&text_tensor_2d, &lengths_tensor, &speaker_tensor) } } + +impl TtsModelT for ArpabetGlowTtsMultiSpeakerModel { + fn encoded_sequence_to_mel(&self, arpabet_encodings: &Vec, speaker_id: i64) -> Tensor { + self.encoded_arpabet_to_mel(arpabet_encodings, speaker_id) + } + + fn encoded_sequence_to_mel_single_speaker(&self, arpabet_encodings: &Vec) -> Tensor { + self.encoded_arpabet_to_mel(arpabet_encodings, -1) + } +} diff --git a/tts_service/src/model/melgan_model.rs b/tts_service/src/model/melgan_model.rs index d7a039ddda..44f23812cc 100644 --- a/tts_service/src/model/melgan_model.rs +++ b/tts_service/src/model/melgan_model.rs @@ -2,6 +2,7 @@ use crate::model::model_container::ModelContainer; use anyhow::Result; use tch::Tensor; use std::path::Path; +use crate::inference::vocoder_model::VocoderModelT; pub struct MelganModel { model_container: ModelContainer, @@ -19,3 +20,9 @@ impl MelganModel { self.model_container.forward(&mel_tensor) } } + +impl VocoderModelT for MelganModel { + fn mel_to_audio(&self, mel_tensor: &Tensor) -> Tensor { + self.model_container.forward(&mel_tensor) + } +}