beef up module config: shared base dir, speakers

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

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

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"
 

tts_service/src/config.rs

@@ -4,10 +4,50 @@ use std::path::Path;
 
 #[derive(Serialize, Deserialize, Debug, Clone)]
 pub struct ModelConfigs {
+  speakers: Vec<Speaker>,
+  pub model_locations: Vec<ModelLocation>,
   tacotron: Vec<ModelDetails>,
   melgan: Vec<ModelDetails>,
 }
 
+#[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<String>,
+  /// Melgan model, if used.
+  pub melgan: Option<String>,
+}
+
+/// 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<String>,
+}
+
+/// 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,

tts_service/src/endpoints/index.rs

@@ -10,5 +10,5 @@ pub async fn get_root(_request: HttpRequest) -> std::io::Result<HttpResponse> {
   println!("GET /");
   Ok(HttpResponse::build(StatusCode::OK)
       .content_type("text/html; charset=utf-8")
-      .body("Hello World"))
+      .body("tts engine"))
 }

tts_service/src/endpoints/tts.rs

@@ -25,7 +25,7 @@ pub struct TtsRequest {
   melgan_model: Option<String>,
 }
 
-pub async fn post_tts(request: HttpRequest,
+pub async fn post_tts(_request: HttpRequest,
   query: Json<TtsRequest>,
   app_state: Data<Arc<AppState>>
 ) -> std::io::Result<HttpResponse> {

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)

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<Self> {
+  pub fn load(filename: &Path) -> Result<Self> {
     let model_container = ModelContainer::load(filename)?;
     Ok(Self {
       model_container,

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<Self> {
+  pub fn load(filename: &Path) -> Result<Self> {
     let model_container = ModelContainer::load(filename)?;
     Ok(Self {
       model_container,

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<ModelType, PathBuf>,
+
   arpabet_tacotron_models: Arc<Mutex<HashMap<String, Arc<ArpabetTacotronModel>>>>,
   melgan_models: Arc<Mutex<HashMap<String, Arc<MelganModel>>>>,
 }
 
 impl ModelCache {
-  pub fn new() -> Self {
+  pub fn new(model_locations: &Vec<ModelLocation>) -> 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<Arc<ArpabetTacotronModel>> {
+    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<Arc<MelganModel>> {
+    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)
+  }
 }

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<Self> {
+  pub fn load(filename: &Path) -> Result<Self> {
     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
   }