From a20a7259b5b5ab6f004e2faec6dc7ecff383b81d Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Tue, 26 May 2020 20:36:08 -0400 Subject: [PATCH] beef up module config: shared base dir, speakers --- deployment.yaml | 3 +- tacotron_melgan/wrap_tacotron2_model.py | 17 ++++ tts_service/models.toml | 99 +++++++++++++++---- tts_service/src/config.rs | 40 ++++++++ tts_service/src/endpoints/index.rs | 2 +- tts_service/src/endpoints/tts.rs | 2 +- tts_service/src/main.rs | 4 +- .../src/model/arpabet_tacotron_model.rs | 3 +- tts_service/src/model/melgan_model.rs | 3 +- tts_service/src/model/model_cache.rs | 44 ++++++++- tts_service/src/model/model_container.rs | 9 +- 11 files changed, 193 insertions(+), 33 deletions(-) diff --git a/deployment.yaml b/deployment.yaml index 1a79820c45..d8594a9feb 100644 --- a/deployment.yaml +++ b/deployment.yaml @@ -33,6 +33,7 @@ spec: annotations: bump: '1' spec: + terminationGracePeriodSeconds: 60 imagePullSecrets: - name: regcred volumes: @@ -87,7 +88,7 @@ spec: # The following image runs a rudimentary Arbapet engine, but doesn't # handle puncutation and tokenization #image: docker.pkg.github.com/ml-applications/voder/tts-service:6e766720ca7d - image: docker.pkg.github.com/ml-applications/voder/tts-service:dafe174dc364 + image: docker.pkg.github.com/ml-applications/voder/tts-service:805bad97c311 volumeMounts: - name: mntdatas3fs mountPath: /data:shared diff --git a/tacotron_melgan/wrap_tacotron2_model.py b/tacotron_melgan/wrap_tacotron2_model.py index fe65ee009b..bf136447a1 100755 --- a/tacotron_melgan/wrap_tacotron2_model.py +++ b/tacotron_melgan/wrap_tacotron2_model.py @@ -20,6 +20,14 @@ tacotron_model_file = '/home/bt/models/tacotron2-nvidia/tacotron2_trump_txlearn_ tacotron_model_file = '/home/bt/models/tacotron2-nvidia/tacotron2_trump_txlearn_arpabet_2020-05-14_chkpt27000' tacotron_model_file = '/home/bt/models/tacotron2-nvidia/tacotron2_trump_txlearn_arpabet_2020-05-14_chkpt27000-18000' tacotron_model_file = '/home/bt/models/tacotron2-nvidia/tacotron2_trump_txlearn_arpabet_2020-05-14_chkpt53500' +tacotron_model_file = '/home/bt/models/tacotron2-nvidia/tacotron2_trump_txlearn_arpabet_2020-05-14_chkpt112500' +tacotron_model_file = '/home/bt/models/tacotron2-nvidia/tacotron2_trump_txlearn_arpabet_2020-05-23_chkpt8000' +tacotron_model_file = '/home/bt/models/tacotron2-nvidia/tacotron2_trump_txlearn_arpabet_2020-05-23_chkpt13000' +tacotron_model_file = '/home/bt/models/tacotron2-nvidia/tacotron2_trump_txlearn_arpabet_2020-05-24_chkpt2000' +tacotron_model_file = '/home/bt/models/tacotron2-nvidia/tacotron2_trump_txlearn_arpabet_2020-05-24_chkpt16000' +tacotron_model_file = '/home/bt/models/tacotron2-nvidia/tacotron2_trump_txlearn_arpabet_2020-05-25_chkpt7000' +tacotron_model_file = '/home/bt/models/tacotron2-nvidia/tacotron2_trump_txlearn_arpabet_2020-05-25_chkpt17000' +tacotron_model_file = '/home/bt/models/tacotron2-nvidia/tacotron2_trump_txlearn_arpabet_2020-05-25_chkpt26500' #output_filename = 'tacotron_arpabet_1.pt' @@ -30,6 +38,15 @@ output_filename = 'tacotron2_trump_txlearn_ljs_arpabet_2020-05-14_ckpt17000.jit' output_filename = 'tacotron2_trump_txlearn_ljs_arpabet_2020-05-14_ckpt27000.jit' output_filename = 'tacotron2_trump_txlearn_ljs_arpabet_2020-05-14_ckpt27000-18000.jit' output_filename = 'tacotron2_trump_txlearn_ljs_arpabet_2020-05-14_ckpt53500.jit' +output_filename = 'tacotron2_trump_txlearn_ljs_arpabet_2020-05-14_chkpt112500.jit' +output_filename = 'tacotron2_trump_txlearn_ljs_arpabet_2020-05-23_chkpt8000.jit' +output_filename = 'tacotron2_trump_txlearn_ljs_arpabet_2020-05-23_chkpt13000.jit' +output_filename = 'tacotron2_trump_txlearn_ljs_arpabet_2020-05-24_chkpt2000.jit' +output_filename = 'tacotron2_trump_txlearn_ljs_arpabet_2020-05-24_chkpt16000.jit' +output_filename = 'tacotron2_trump_txlearn_ljs_arpabet_2020-05-25_chkpt7000.jit' +output_filename = 'tacotron2_trump_txlearn_ljs_arpabet_2020-05-25_chkpt17000.jit' +output_filename = 'tacotron2_trump_txlearn_ljs_arpabet_2020-05-25_chkpt26500.jit' + # NB: Tacotron's original shape is [148, 512]. Arpabet encoding is [256, 512]. NUM_SYMBOLS_ENCODED = 256 # 148 by default, 256 with my custom Arpabet encoding diff --git a/tts_service/models.toml b/tts_service/models.toml index e605f585e1..c2c55095df 100644 --- a/tts_service/models.toml +++ b/tts_service/models.toml @@ -1,48 +1,107 @@ +# --------------- Speakers --------------- +[[speakers]] +name = "Donald Trump" +slug = "trump" +model_pipeline = "ArpabetTacotronMelgan" +tacotron = "tacotron2_trump_txlearn_ljs_arpabet_2020-05-25_chkpt26500.jit" +melgan = "melgan_trump-txlearn-2020.05.05_13675.jit" + +[[speakers]] +name = "High-pitched Linda Johnson" +slug = "ljsh" +model_pipeline = "ArpabetTacotronMelgan" +tacotron = "" +melgan = "" + +# --------------- Model configs --------------- + +[[model_locations]] +model_type = "ArpabetTacotron" +base_directory = "/home/bt/dev/voder/tacotron_melgan" + +[[model_locations]] +model_type = "Melgan" +base_directory = "/home/bt/dev/tacotron-melgan/" + # --------------- Tacotron --------------- # TODO: Get a high-quality LJS Tacotron [[tacotron]] -file_path = "/home/bt/dev/voder/tacotron_melgan/tacotron2_trump_txlearn_ljs_arpabet_2020-02-06_ckpt30000.jit" +file_path = "tacotron2_trump_txlearn_ljs_arpabet_2020-02-06_ckpt30000.jit" description = "First Trump txlearn with Arpabet. Trained 30k iters and sounds great!" [[tacotron]] -file_path = "/home/bt/dev/voder/tacotron_melgan/tacotron2_trump_txlearn_ljs_arpabet_2020-05-14_ckpt7000.jit" +file_path = "tacotron2_trump_txlearn_ljs_arpabet_2020-05-14_ckpt7000.jit" description = "Trump txlearn with Arpabet 2020-05-14" [[tacotron]] -file_path = "/home/bt/dev/voder/tacotron_melgan/tacotron2_trump_txlearn_ljs_arpabet_2020-05-14_ckpt11500.jit" -description = "Trump txlearn with Arpabet 2020-05-14" +file_path = "tacotron2_trump_txlearn_ljs_arpabet_2020-05-14_ckpt11500.jit" +description = "(see previous)" [[tacotron]] -file_path = "/home/bt/dev/voder/tacotron_melgan/tacotron2_trump_txlearn_ljs_arpabet_2020-05-14_ckpt17000.jit" -description = "Trump txlearn with Arpabet 2020-05-14" +file_path = "tacotron2_trump_txlearn_ljs_arpabet_2020-05-14_ckpt17000.jit" +description = "(see previous)" [[tacotron]] -file_path = "/home/bt/dev/voder/tacotron_melgan/tacotron2_trump_txlearn_ljs_arpabet_2020-05-14_ckpt27000.jit" -description = "Trump txlearn with Arpabet 2020-05-14" +file_path = "tacotron2_trump_txlearn_ljs_arpabet_2020-05-14_ckpt27000.jit" +description = "(see previous)" [[tacotron]] -file_path = "/home/bt/dev/voder/tacotron_melgan/tacotron2_trump_txlearn_ljs_arpabet_2020-05-14_ckpt27000-18000.jit" -description = "Trump txlearn with Arpabet 2020-05-14" +file_path = "tacotron2_trump_txlearn_ljs_arpabet_2020-05-14_ckpt27000-18000.jit" +description = "(see previous)" [[tacotron]] -file_path = "/home/bt/dev/voder/tacotron_melgan/tacotron2_trump_txlearn_ljs_arpabet_2020-05-14_ckpt53500.jit" -description = "Trump txlearn with Arpabet 2020-05-14" +file_path = "tacotron2_trump_txlearn_ljs_arpabet_2020-05-14_ckpt53500.jit" +description = "(see previous)" + +[[tacotron]] +file_path = "tacotron2_trump_txlearn_ljs_arpabet_2020-05-14_chkpt112500.jit" +description = "(see previous)" + +[[tacotron]] +file_path = "tacotron2_trump_txlearn_ljs_arpabet_2020-05-23_chkpt8000.jit" +description = "(see previous)" + +[[tacotron]] +file_path = "tacotron2_trump_txlearn_ljs_arpabet_2020-05-23_chkpt13000.jit" +description = "(see previous)" + +[[tacotron]] +file_path = "tacotron2_trump_txlearn_ljs_arpabet_2020-05-24_chkpt2000.jit" +description = "Trump txlearn with Arpabet 2020-05-24; adds 400 new training files from steno" + +[[tacotron]] +file_path = "tacotron2_trump_txlearn_ljs_arpabet_2020-05-24_chkpt16000.jit" +description = "(see previous)" + +[[tacotron]] +file_path = "acotron2_trump_txlearn_ljs_arpabet_2020-05-25_chkpt7000.jit" +description = "Trump txlearn with Arpabet 2020-05-25; adds lots from steno, disabled one bad quality source" + +[[tacotron]] +file_path = "tacotron2_trump_txlearn_ljs_arpabet_2020-05-25_chkpt17000.jit" +description = "(see previous)" + +[[tacotron]] +file_path = "tacotron2_trump_txlearn_ljs_arpabet_2020-05-25_chkpt26500.jit" +description = "(see previous)" + +# ======== OLDER RUNS ====== [[tacotron]] # OLD NAME, perhaps I can get the epoch: file_path = "/home/bt/dev/voder/tacotron_melgan/tacotron_arpabet_1.pt" -file_path = "/home/bt/dev/voder/tacotron_melgan/tacotron2_ljs_high_pitch_arpabet_2020-03-28.jit" +file_path = "tacotron2_ljs_high_pitch_arpabet_2020-03-28.jit" description = "High-quality high-pitched LJS; arpabet encodings." [[tacotron]] # TODO: Remove. This blows up. -file_path = "/home/bt/dev/voder/tacotron_melgan/tacotron_container.pt" +file_path = "tacotron_container.pt" description = "Unknown (broken)" [[tacotron]] # TODO: Remove. This hangs forever. -file_path = "/home/bt/dev/voder/tacotron_melgan/tacotron_jit_model_voder_c0cac635.pt" +file_path = "acotron_jit_model_voder_c0cac635.pt" description = "Unknown (broken)" @@ -51,22 +110,22 @@ description = "Unknown (broken)" # TODO: Get a high-quality LJS Melgan [[melgan]] -file_path = "/home/bt/dev/tacotron-melgan/melgan_trump-txlearn-2020.05.05_7525.jit" +file_path = "melgan_trump-txlearn-2020.05.05_7525.jit" description = "Latest LJS transfer learned Trump." [[melgan]] -file_path = "/home/bt/dev/tacotron-melgan/melgan_trump-txlearn-2020.05.05_13675.jit" +file_path = "melgan_trump-txlearn-2020.05.05_13675.jit" description = "Latest LJS transfer learned Trump." [[melgan]] -file_path = "/home/bt/dev/tacotron-melgan/melgan_ljs-high-pitch_container.pt" +file_path = "melgan_ljs-high-pitch_container.pt" description = "High-pitched LJS" [[melgan]] -file_path = "/home/bt/dev/tacotron-melgan/shared_melgan_container-2020.01.04.pt" +file_path = "shared_melgan_container-2020.01.04.pt" description = "Unknown" [[melgan]] -file_path = "/home/bt/dev/tacotron-melgan/shared_melgan_container-2020.01.05.pt" +file_path = "shared_melgan_container-2020.01.05.pt" description = "Unknown" diff --git a/tts_service/src/config.rs b/tts_service/src/config.rs index b4b3a18217..76f4ef3ca9 100644 --- a/tts_service/src/config.rs +++ b/tts_service/src/config.rs @@ -4,10 +4,50 @@ use std::path::Path; #[derive(Serialize, Deserialize, Debug, Clone)] pub struct ModelConfigs { + speakers: Vec, + pub model_locations: Vec, tacotron: Vec, melgan: Vec, } +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)] +pub struct Speaker { + /// Name of the speaker + pub name: String, + /// URL slug / key + pub slug: String, + /// Model pipeline to use + pub model_pipeline: ModelPipeline, + /// Tacotron model, if used. + pub tacotron: Option, + /// Melgan model, if used. + pub melgan: Option, +} + +/// The types of models supported by our system. +#[derive(Serialize, Deserialize, Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum ModelType { + RawTextTacotron, + ArpabetTacotron, + Melgan, +} + +/// Where the various models are stored. +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)] +pub struct ModelLocation { + pub model_type: ModelType, + pub base_directory: Option, +} + +/// The valid model pipelines in use by our system. +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)] +pub enum ModelPipeline { + //ArpabetTacotronGriffinLim, + //ArpabetTacotronWorld, + ArpabetTacotronMelgan, + RawTextTacotronMelgan, +} + #[derive(Serialize, Deserialize, Debug, Clone, PartialEq)] pub struct ModelDetails { pub file_path: String, diff --git a/tts_service/src/endpoints/index.rs b/tts_service/src/endpoints/index.rs index 875736977f..950a078c5a 100644 --- a/tts_service/src/endpoints/index.rs +++ b/tts_service/src/endpoints/index.rs @@ -10,5 +10,5 @@ pub async fn get_root(_request: HttpRequest) -> std::io::Result { println!("GET /"); Ok(HttpResponse::build(StatusCode::OK) .content_type("text/html; charset=utf-8") - .body("Hello World")) + .body("tts engine")) } diff --git a/tts_service/src/endpoints/tts.rs b/tts_service/src/endpoints/tts.rs index d60576dcd8..dad07f70c8 100644 --- a/tts_service/src/endpoints/tts.rs +++ b/tts_service/src/endpoints/tts.rs @@ -25,7 +25,7 @@ pub struct TtsRequest { melgan_model: Option, } -pub async fn post_tts(request: HttpRequest, +pub async fn post_tts(_request: HttpRequest, query: Json, app_state: Data> ) -> std::io::Result { diff --git a/tts_service/src/main.rs b/tts_service/src/main.rs index 99a1244200..f945de98c8 100644 --- a/tts_service/src/main.rs +++ b/tts_service/src/main.rs @@ -75,9 +75,11 @@ async fn main() -> std::io::Result<()> { let model_configs = ModelConfigs::load_from_file(&model_config_file); println!("Model configs: {:?}", model_configs); + let model_cache = ModelCache::new(&model_configs.model_locations); + let app_state = AppState { model_configs, - model_cache: ModelCache::new(), + model_cache, }; let bind_address = env::var(BIND_ADDRESS) diff --git a/tts_service/src/model/arpabet_tacotron_model.rs b/tts_service/src/model/arpabet_tacotron_model.rs index 52dbeacd68..12a1385253 100644 --- a/tts_service/src/model/arpabet_tacotron_model.rs +++ b/tts_service/src/model/arpabet_tacotron_model.rs @@ -1,13 +1,14 @@ use crate::model::model_container::ModelContainer; use anyhow::Result; use tch::Tensor; +use std::path::Path; pub struct ArpabetTacotronModel { model_container: ModelContainer, } impl ArpabetTacotronModel { - pub fn load(filename: &str) -> Result { + pub fn load(filename: &Path) -> Result { let model_container = ModelContainer::load(filename)?; Ok(Self { model_container, diff --git a/tts_service/src/model/melgan_model.rs b/tts_service/src/model/melgan_model.rs index 3047bf4622..d7a039ddda 100644 --- a/tts_service/src/model/melgan_model.rs +++ b/tts_service/src/model/melgan_model.rs @@ -1,13 +1,14 @@ use crate::model::model_container::ModelContainer; use anyhow::Result; use tch::Tensor; +use std::path::Path; pub struct MelganModel { model_container: ModelContainer, } impl MelganModel { - pub fn load(filename: &str) -> Result { + pub fn load(filename: &Path) -> Result { let model_container = ModelContainer::load(filename)?; Ok(Self { model_container, diff --git a/tts_service/src/model/model_cache.rs b/tts_service/src/model/model_cache.rs index a4fe4f2dce..0ef15c26a1 100644 --- a/tts_service/src/model/model_cache.rs +++ b/tts_service/src/model/model_cache.rs @@ -4,21 +4,46 @@ use crate::model::melgan_model::MelganModel; use std::collections::HashMap; use std::sync::Arc; use std::sync::Mutex; +use crate::config::{ModelType, ModelLocation}; +use crate::config::ModelType::ArpabetTacotron; +use std::path::{PathBuf, Path}; pub struct ModelCache { + /// Base directories for each model type(if configured) + base_directories: HashMap, + arpabet_tacotron_models: Arc>>>, melgan_models: Arc>>>, } impl ModelCache { - pub fn new() -> Self { + pub fn new(model_locations: &Vec) -> Self { + let mut base_directories = HashMap::new(); + + for model_location in model_locations.iter() { + if let Some(base_directory) = model_location.base_directory.as_ref() { + let path = PathBuf::from(base_directory); + base_directories.insert(model_location.model_type.clone(), path); + } + } + Self { + base_directories, arpabet_tacotron_models: Arc::new(Mutex::new(HashMap::new())), melgan_models: Arc::new(Mutex::new(HashMap::new())), } } pub fn get_or_load_arbabet_tacotron(&self, filename: &str) -> Option> { + let base_directory = self.get_base_directory(ModelType::ArpabetTacotron); + + 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_tacotron_models.lock() .expect("should unlock"); @@ -26,7 +51,7 @@ impl ModelCache { return Some(model.clone()); } - match ArpabetTacotronModel::load(&filename) { + match ArpabetTacotronModel::load(&file_path) { Ok(model) => { let arc = Arc::new(model); lock.insert(filename.to_string(), arc.clone()); @@ -37,6 +62,15 @@ impl ModelCache { } pub fn get_or_load_melgan(&self, filename: &str) -> Option> { + let base_directory = self.get_base_directory(ModelType::Melgan); + + 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.melgan_models.lock() .expect("should unlock"); @@ -44,7 +78,7 @@ impl ModelCache { return Some(model.clone()); } - match MelganModel::load(&filename) { + match MelganModel::load(&file_path) { Ok(model) => { let arc = Arc::new(model); lock.insert(filename.to_string(), arc.clone()); @@ -53,4 +87,8 @@ impl ModelCache { Err(_) => None, } } + + fn get_base_directory(&self, model_type: ModelType) -> Option<&PathBuf> { + self.base_directories.get(&model_type) + } } diff --git a/tts_service/src/model/model_container.rs b/tts_service/src/model/model_container.rs index 54fc6ccdae..f2849f19de 100644 --- a/tts_service/src/model/model_container.rs +++ b/tts_service/src/model/model_container.rs @@ -3,25 +3,26 @@ use tch::CModule; use tch::Tensor; use tch::nn::Module; use tch::nn::ModuleT; +use std::path::{Path, PathBuf}; /// Holds the loaded pytorch JIT model pub struct ModelContainer { - filename: String, + filename: PathBuf, jit_model: CModule, } impl ModelContainer { - pub fn load(filename: &str) -> Result { + pub fn load(filename: &Path) -> Result { let model = CModule::load(filename) .map_err(|e| anyhow!("jit model failed to load: {}", e))?; Ok(Self { - filename: filename.to_string(), + filename: filename.to_path_buf(), jit_model: model, }) } - pub fn get_filename(&self) -> &str { + pub fn get_filename(&self) -> &Path { &self.filename }