add dev and prod model paths

tts_service/models.toml

@@ -61,13 +61,15 @@ speaker_id = 3
 
 [[model_locations]]
 model_type = "ArpabetGlowTts"
-base_directory = "/home/bt/dev/jit-models/glow-tts-jaywalnut/OUTPUT"
-
-[[model_locations]]
-model_type = "Melgan"
-base_directory = "/home/bt/dev/jit-models/melgan-swpark/OUTPUT"
+dev_base_directory = "/home/bt/dev/jit-models/glow-tts-jaywalnut/OUTPUT"
+prod_base_directory = "/home/bt/models/jit/glow-tts-jaywalnut"
 
 [[model_locations]]
 model_type = "ArpabetGlowTtsMultiSpeaker"
-base_directory = "/home/bt/dev/jit-models/glow-tts-jaywalnut-multispeaker/OUTPUT"
+dev_base_directory = "/home/bt/dev/jit-models/glow-tts-jaywalnut-multispeaker/OUTPUT"
+prod_base_directory = "/home/bt/models/jit/glow-tts-jaywalnut-multispeaker"
 
+[[model_locations]]
+model_type = "Melgan"
+dev_base_directory = "/home/bt/dev/jit-models/melgan-swpark/OUTPUT"
+prod_base_directory = "/home/bt/models/jit/melgan-swpark"

tts_service/src/model/model_cache.rs

@@ -10,7 +10,8 @@ use crate::model::arpabet_glow_tts_multi_speaker_model::ArpabetGlowTtsMultiSpeak
 
 pub struct ModelCache {
   /// Base directories for each model type(if configured)
-  base_directories: HashMap<ModelType, PathBuf>,
+  dev_base_directories: HashMap<ModelType, PathBuf>,
+  prod_base_directories: HashMap<ModelType, PathBuf>,
 
   arpabet_tacotron_models: Arc<Mutex<HashMap<String, Arc<ArpabetTacotronModel>>>>,
   arpabet_glow_tts_models: Arc<Mutex<HashMap<String, Arc<ArpabetGlowTtsModel>>>>,
@@ -20,17 +21,23 @@ pub struct ModelCache {
 
 impl ModelCache {
   pub fn new(model_locations: &Vec<ModelLocation>) -> Self {
-    let mut base_directories = HashMap::new();
+    let mut dev_base_directories = HashMap::new();
+    let mut prod_base_directories = HashMap::new();
 
     for model_location in model_locations.iter() {
-      if let Some(base_directory) = model_location.base_directory.as_ref() {
+      if let Some(base_directory) = model_location.dev_base_directory.as_ref() {
         let path = PathBuf::from(base_directory);
-        base_directories.insert(model_location.model_type.clone(), path);
+        dev_base_directories.insert(model_location.model_type.clone(), path);
+      }
+      if let Some(base_directory) = model_location.prod_base_directory.as_ref() {
+        let path = PathBuf::from(base_directory);
+        prod_base_directories.insert(model_location.model_type.clone(), path);
       }
     }
 
     Self {
-      base_directories,
+      dev_base_directories,
+      prod_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())),
@@ -39,15 +46,6 @@ impl ModelCache {
   }
 
   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");
 
@@ -55,29 +53,36 @@ impl ModelCache {
       return Some(model.clone());
     }
 
-    match ArpabetTacotronModel::load(&file_path) {
-      Ok(model) => {
+    let mut file_path = PathBuf::from(filename);
+
+    let mut maybe_loaded = None;
+
+    if file_path.is_absolute() {
+      maybe_loaded = self.load_arpabet_tacotron_model(&file_path);
+    } else {
+      let check_path = self.resolve_prod_path(ModelType::ArpabetTacotron, &file_path);
+      maybe_loaded = self.load_arpabet_tacotron_model(&check_path);
+
+      if maybe_loaded.is_none() {
+        let check_path = self.resolve_dev_path(ModelType::ArpabetTacotron, &file_path);
+        maybe_loaded = self.load_arpabet_tacotron_model(&check_path);
+      }
+    }
+
+    match maybe_loaded {
+      Some(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 => {
+        warn!("There was an error loading the model `{:?}`", file_path);
         None
-      },
+      }
     }
   }
 
   pub fn get_or_load_arbabet_glow_tts(&self, filename: &str) -> Option<Arc<ArpabetGlowTtsModel>> {
-    let base_directory = self.get_base_directory(ModelType::ArpabetGlowTts);
-
-    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_models.lock()
         .expect("should unlock");
 
@@ -85,29 +90,36 @@ impl ModelCache {
       return Some(model.clone());
     }
 
-    match ArpabetGlowTtsModel::load(&file_path) {
-      Ok(model) => {
+    let mut file_path = PathBuf::from(filename);
+
+    let mut maybe_loaded = None;
+
+    if file_path.is_absolute() {
+      maybe_loaded = self.load_arpabet_glow_tts_model(&file_path);
+    } else {
+      let check_path = self.resolve_prod_path(ModelType::ArpabetGlowTts, &file_path);
+      maybe_loaded = self.load_arpabet_glow_tts_model(&check_path);
+
+      if maybe_loaded.is_none() {
+        let check_path = self.resolve_dev_path(ModelType::ArpabetGlowTts, &file_path);
+        maybe_loaded = self.load_arpabet_glow_tts_model(&check_path);
+      }
+    }
+
+    match maybe_loaded {
+      Some(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 => {
+        warn!("There was an error loading the model `{:?}`", file_path);
         None
-      },
+      }
     }
   }
 
   pub fn get_or_load_arbabet_glow_tts_multi_speaker(&self, filename: &str) -> Option<Arc<ArpabetGlowTtsMultiSpeakerModel>> {
-    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");
 
@@ -115,29 +127,36 @@ impl ModelCache {
       return Some(model.clone());
     }
 
-    match ArpabetGlowTtsMultiSpeakerModel::load(&file_path) {
-      Ok(model) => {
+    let mut file_path = PathBuf::from(filename);
+
+    let mut maybe_loaded = None;
+
+    if file_path.is_absolute() {
+      maybe_loaded = self.load_arpabet_glow_tts_multi_speaker_model(&file_path);
+    } else {
+      let check_path = self.resolve_prod_path(ModelType::ArpabetGlowTtsMultiSpeaker, &file_path);
+      maybe_loaded = self.load_arpabet_glow_tts_multi_speaker_model(&check_path);
+
+      if maybe_loaded.is_none() {
+        let check_path = self.resolve_dev_path(ModelType::ArpabetGlowTtsMultiSpeaker, &file_path);
+        maybe_loaded = self.load_arpabet_glow_tts_multi_speaker_model(&check_path);
+      }
+    }
+
+    match maybe_loaded {
+      Some(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 => {
+        warn!("There was an error loading the model `{:?}`", file_path);
         None
-      },
+      }
     }
   }
 
   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");
 
@@ -145,20 +164,68 @@ impl ModelCache {
       return Some(model.clone());
     }
 
-    match MelganModel::load(&file_path) {
-      Ok(model) => {
+    let mut file_path = PathBuf::from(filename);
+
+    let mut maybe_loaded = None;
+
+    if file_path.is_absolute() {
+      maybe_loaded = self.load_melgan_model(&file_path);
+    } else {
+      let check_path = self.resolve_prod_path(ModelType::Melgan, &file_path);
+      maybe_loaded = self.load_melgan_model(&check_path);
+
+      if maybe_loaded.is_none() {
+        let check_path = self.resolve_dev_path(ModelType::Melgan, &file_path);
+        maybe_loaded = self.load_melgan_model(&check_path);
+      }
+    }
+
+    match maybe_loaded {
+      Some(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 => {
+        warn!("There was an error loading the model `{:?}`", file_path);
         None
-      },
+      }
     }
   }
 
-  fn get_base_directory(&self, model_type: ModelType) -> Option<&PathBuf> {
-    self.base_directories.get(&model_type)
+  fn load_melgan_model(&self, file_path: &PathBuf) -> Option<MelganModel> {
+    MelganModel::load(&file_path).ok()
+  }
+
+  fn load_arpabet_tacotron_model(&self, file_path: &PathBuf) -> Option<ArpabetTacotronModel> {
+    ArpabetTacotronModel::load(&file_path).ok()
+  }
+
+  fn load_arpabet_glow_tts_model(&self, file_path: &PathBuf) -> Option<ArpabetGlowTtsModel> {
+    ArpabetGlowTtsModel::load(&file_path).ok()
+  }
+
+  fn load_arpabet_glow_tts_multi_speaker_model(&self, file_path: &PathBuf) -> Option<ArpabetGlowTtsMultiSpeakerModel> {
+    ArpabetGlowTtsMultiSpeakerModel::load(&file_path).ok()
+  }
+
+  fn resolve_prod_path(&self, model_type: ModelType, file_name: &PathBuf) -> PathBuf {
+    let base_directory = self.get_prod_base_directory(model_type);
+    base_directory.map(|base_dir| base_dir.join(&file_name))
+        .unwrap_or(file_name.clone())
+  }
+
+  fn resolve_dev_path(&self, model_type: ModelType, file_name: &PathBuf) -> PathBuf {
+    let base_directory = self.get_dev_base_directory(model_type);
+    base_directory.map(|base_dir| base_dir.join(&file_name))
+        .unwrap_or(file_name.clone())
+  }
+
+  fn get_prod_base_directory(&self, model_type: ModelType) -> Option<&PathBuf> {
+    self.prod_base_directories.get(&model_type)
+  }
+
+  fn get_dev_base_directory(&self, model_type: ModelType) -> Option<&PathBuf> {
+    self.dev_base_directories.get(&model_type)
   }
 }

tts_service/src/model/model_config.rs

@@ -26,6 +26,9 @@ pub struct Speaker {
   pub glow_tts_multi_speaker: Option<String>,
   /// Melgan model, if used.
   pub melgan: Option<String>,
+  /// Whether the model is "production" or not
+  /// Production models are backed up on GDrive and are synced to DigitalOcean
+  pub production: Option<bool>,
 }
 
 /// The types of models supported by our system.
@@ -42,7 +45,8 @@ pub enum ModelType {
 #[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
 pub struct ModelLocation {
   pub model_type: ModelType,
-  pub base_directory: Option<String>,
+  pub dev_base_directory: Option<String>,
+  pub prod_base_directory: Option<String>,
 }
 
 /// The valid model pipelines in use by our system.