diff --git a/deployment.yaml b/deployment.yaml index 8effe03fc3..59d1f4ad27 100644 --- a/deployment.yaml +++ b/deployment.yaml @@ -86,7 +86,7 @@ spec: # LSpeech model that sounds 100% amazing # image: docker.pkg.github.com/echelon/voder/tts-service:6b5c22567614 # - image: docker.pkg.github.com/ml-applications/voder/tts-service:0281e206db16 + image: docker.pkg.github.com/ml-applications/voder/tts-service:42f51ecd18e2 resources: limits: # 4 GiB (4096 MiB) diff --git a/tts_service/models.toml b/tts_service/models.toml index e89f587bbe..0cf9361757 100644 --- a/tts_service/models.toml +++ b/tts_service/models.toml @@ -25,6 +25,14 @@ model_pipeline = "ArpabetTacotronMelgan" tacotron = "tacotron2_ljs_high_pitch_arpabet_2020-03-28.jit" melgan = "melgan_ljs-high-pitch_container.pt" +[[speakers]] +name = "Arnold Schwarzenegger" +slug = "arnold" +model_pipeline = "ArpabetGlowTtsMultiSpeakerMelgan" +glow_tts_multi_speaker = "glow_multispeaker.torchjit" +melgan = "melgan_trump_txlearn_2020-05-30_steno_chkpt9800.torchjit" +speaker_id = 2 + # --------------- Model configs --------------- [[model_locations]] @@ -40,6 +48,10 @@ base_directory = "/home/bt/dev/jit-models/glow-tts-jaywalnut" model_type = "Melgan" base_directory = "/home/bt/dev/tacotron-melgan/" +[[model_locations]] +model_type = "ArpabetGlowTtsMultiSpeaker" +base_directory = "/home/bt/dev/jit-models/glow-tts-jaywalnut-multispeaker" + # --------------- Tacotron --------------- # TODO: Get a high-quality LJS Tacotron @@ -160,6 +172,12 @@ description = "First experiments (continued); modified with longer input to JIT" file_path = "glow_tts_ljs_txlearn_arpabet_2020-06-01_G_3614.torchjit" description = "First experiments (continued); modified with longer input to JIT" + +# ---------------- Glow TTS Multi-speaker ---------------- +[[glow_tts_multi_speaker]] +file_path = "glow_multispeaker.torchjit" +description = "Mutli-speaker garbage" + # ---------------- Melgan ---------------- # TODO: Get a high-quality LJS Melgan diff --git a/tts_service/src/endpoints/legacy_speak.rs b/tts_service/src/endpoints/legacy_speak.rs index b54515426f..6a3dc0cfa8 100644 --- a/tts_service/src/endpoints/legacy_speak.rs +++ b/tts_service/src/endpoints/legacy_speak.rs @@ -217,6 +217,7 @@ pub async fn legacy_get_speak(request: HttpRequest, .set_header(header::CONTENT_RANGE, content_range_value) .body(wav_data)) }, + ModelPipeline::ArpabetGlowTtsMultiSpeakerMelgan=> unimplemented!(), ModelPipeline::RawTextTacotronMelgan => unimplemented!(), } } diff --git a/tts_service/src/endpoints/speak.rs b/tts_service/src/endpoints/speak.rs index 67f094be06..6653c7ff0f 100644 --- a/tts_service/src/endpoints/speak.rs +++ b/tts_service/src/endpoints/speak.rs @@ -15,7 +15,7 @@ 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; +use crate::model::pipelines::{arpabet_glow_tts_melgan_pipeline, arpabet_glow_tts_multi_speaker_melgan_pipeline}; use crate::database::model::NewSentence; #[derive(Deserialize)] @@ -148,6 +148,45 @@ pub async fn post_speak(request: HttpRequest, .set_header(header::CONTENT_RANGE, content_range_value) .body(wav_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 wav_data = arpabet_glow_tts_multi_speaker_melgan_pipeline( + &cleaned_text, speaker_id, &glow_tts_multi_speaker, &melgan); + + // To make iOS Safari work, you need a Content-Range and Content-Length header: + // https://stackoverflow.com/a/17835399 + let content_range_value = format!("bytes 0-{}/{}", wav_data.len(), wav_data.len()); + Ok(HttpResponse::build(StatusCode::OK) + .content_type("audio/wav") + .set_header(header::CONTENT_DISPOSITION, "attachment; filename = \"generated.wav\"") + .set_header(header::CONTENT_RANGE, content_range_value) + .body(wav_data)) + }, ModelPipeline::RawTextTacotronMelgan => unimplemented!(), } } 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 new file mode 100644 index 0000000000..ead6373c77 --- /dev/null +++ b/tts_service/src/model/arpabet_glow_tts_multi_speaker_model.rs @@ -0,0 +1,33 @@ +use crate::model::model_container::ModelContainer; +use anyhow::{Result, Error}; +use std::path::Path; +use std::thread; +use std::time::Duration; +use tch::Tensor; +use crate::text::arpabet::text_to_arpabet_encoding_glow_tts; + +pub struct ArpabetGlowTtsMultiSpeakerModel { + model_container: ModelContainer, +} + +impl ArpabetGlowTtsMultiSpeakerModel { + pub fn load(filename: &Path) -> Result { + let model_container = ModelContainer::load(filename)?; + Ok(Self { + model_container, + }) + } + + pub fn encoded_arpabet_to_mel(&self, arpabet_encodings: &Vec, speaker_id: i64) -> Tensor { + let text_tensor = Tensor::of_slice(&arpabet_encodings.as_slice()); + let text_tensor_2d = text_tensor.unsqueeze(0); // Add another axis + + let lengths = [arpabet_encodings.len() as i64]; + let lengths_tensor = Tensor::of_slice(&lengths); + + let speaker = [speaker_id]; + let speaker_tensor = Tensor::of_slice(&speaker); + + self.model_container.forward3(&text_tensor_2d, &lengths_tensor, &speaker_tensor) + } +} diff --git a/tts_service/src/model/mod.rs b/tts_service/src/model/mod.rs index 9ad40ca459..594be453d3 100644 --- a/tts_service/src/model/mod.rs +++ b/tts_service/src/model/mod.rs @@ -1,4 +1,5 @@ pub mod arpabet_glow_tts_model; +pub mod arpabet_glow_tts_multi_speaker_model; pub mod arpabet_tacotron_model; pub mod melgan_model; pub mod model_cache; diff --git a/tts_service/src/model/model_cache.rs b/tts_service/src/model/model_cache.rs index d3fcd68247..5b1cc5b81a 100644 --- a/tts_service/src/model/model_cache.rs +++ b/tts_service/src/model/model_cache.rs @@ -6,6 +6,7 @@ use std::path::PathBuf; use std::sync::Arc; use std::sync::Mutex; use crate::model::arpabet_glow_tts_model::ArpabetGlowTtsModel; +use crate::model::arpabet_glow_tts_multi_speaker_model::ArpabetGlowTtsMultiSpeakerModel; pub struct ModelCache { /// Base directories for each model type(if configured) @@ -13,6 +14,7 @@ pub struct ModelCache { arpabet_tacotron_models: Arc>>>, arpabet_glow_tts_models: Arc>>>, + arpabet_glow_tts_multi_speaker_models: Arc>>>, melgan_models: Arc>>>, } @@ -31,6 +33,7 @@ impl ModelCache { base_directories, arpabet_tacotron_models: Arc::new(Mutex::new(HashMap::new())), arpabet_glow_tts_models: Arc::new(Mutex::new(HashMap::new())), + arpabet_glow_tts_multi_speaker_models: Arc::new(Mutex::new(HashMap::new())), melgan_models: Arc::new(Mutex::new(HashMap::new())), } } @@ -95,6 +98,36 @@ impl ModelCache { } } + pub fn get_or_load_arbabet_glow_tts_multi_speaker(&self, filename: &str) -> Option> { + let base_directory = self.get_base_directory(ModelType::ArpabetGlowTtsMultiSpeaker); + + let mut file_path = PathBuf::from(filename); + + if !file_path.is_absolute() { + file_path = base_directory.map(|base_dir| base_dir.join(&file_path)) + .unwrap_or(file_path); + } + + let mut lock = self.arpabet_glow_tts_multi_speaker_models.lock() + .expect("should unlock"); + + if let Some(model) = lock.get(filename) { + return Some(model.clone()); + } + + match ArpabetGlowTtsMultiSpeakerModel::load(&file_path) { + Ok(model) => { + let arc = Arc::new(model); + lock.insert(filename.to_string(), arc.clone()); + Some(arc) + }, + Err(e) => { + warn!("There was an error loading the model `{:?}`: {}", file_path, e); + None + }, + } + } + pub fn get_or_load_melgan(&self, filename: &str) -> Option> { let base_directory = self.get_base_directory(ModelType::Melgan); diff --git a/tts_service/src/model/model_config.rs b/tts_service/src/model/model_config.rs index 4c750775e1..ee3926ca6e 100644 --- a/tts_service/src/model/model_config.rs +++ b/tts_service/src/model/model_config.rs @@ -8,6 +8,7 @@ pub struct ModelConfigs { pub model_locations: Vec, tacotron: Vec, glow_tts: Vec, + glow_tts_multi_speaker: Vec, melgan: Vec, } @@ -17,12 +18,16 @@ pub struct Speaker { pub name: String, /// URL slug / key pub slug: String, + /// Speaker Id, if used + pub speaker_id: Option, /// Model pipeline to use pub model_pipeline: ModelPipeline, /// Tacotron model, if used. pub tacotron: Option, /// Glow-tts model, if used. pub glow_tts: Option, + /// Glow-tts multi-speaker model, if used. + pub glow_tts_multi_speaker: Option, /// Melgan model, if used. pub melgan: Option, } @@ -33,6 +38,7 @@ pub enum ModelType { RawTextTacotron, ArpabetTacotron, ArpabetGlowTts, // NB: This uses a different arpabet preprocessor + ArpabetGlowTtsMultiSpeaker, // NB: This uses a different arpabet preprocessor Melgan, } @@ -51,6 +57,7 @@ pub enum ModelPipeline { ArpabetTacotronMelgan, RawTextTacotronMelgan, ArpabetGlowTtsMelgan, + ArpabetGlowTtsMultiSpeakerMelgan, } #[derive(Serialize, Deserialize, Debug, Clone, PartialEq)] diff --git a/tts_service/src/model/model_container.rs b/tts_service/src/model/model_container.rs index 4f9ff0ba5a..b2d7f77c13 100644 --- a/tts_service/src/model/model_container.rs +++ b/tts_service/src/model/model_container.rs @@ -40,4 +40,14 @@ impl ModelContainer { result.expect("SHOULD WORK") } + + pub fn forward3(&self, arg1: &Tensor, arg2: &Tensor, arg3: &Tensor) -> Tensor { + let result = self.jit_model.forward_ts(&[arg1, arg2, arg3]); + + if let Err(err) = result.as_ref() { + println!("error: {:?}", err); + } + + result.expect("SHOULD WORK") + } } diff --git a/tts_service/src/model/pipelines.rs b/tts_service/src/model/pipelines.rs index 772466d936..94535781e8 100644 --- a/tts_service/src/model/pipelines.rs +++ b/tts_service/src/model/pipelines.rs @@ -9,11 +9,29 @@ use hound::WavWriter; 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; // TODO: This might be useful to implement as a multi-stage // state machine struct with functions, that way you can pull out // intermediate pieces (eg. text to send to the db, or reject the request) +pub fn arpabet_glow_tts_multi_speaker_melgan_pipeline( + cleaned_text: &str, + speaker_id: i64, + arpabet_glow_tts: &ArpabetGlowTtsMultiSpeakerModel, + melgan: &MelganModel) -> 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 audio_tensor = melgan.tacotron_mel_to_audio(&mel_tensor); + let audio_signal = mel_audio_tensor_to_audio_signal(audio_tensor); + + audio_signal_to_wav_bytes(audio_signal) +} + pub fn arpabet_glow_tts_melgan_pipeline( cleaned_text: &str, arpabet_glow_tts: &ArpabetGlowTtsModel,