diff --git a/tts_service/models.toml b/tts_service/models.toml index ceb2b820e5..c08ad81114 100644 --- a/tts_service/models.toml +++ b/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" diff --git a/tts_service/src/model/model_cache.rs b/tts_service/src/model/model_cache.rs index 5b1cc5b81a..d3a703f983 100644 --- a/tts_service/src/model/model_cache.rs +++ b/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, + dev_base_directories: HashMap, + prod_base_directories: HashMap, arpabet_tacotron_models: Arc>>>, arpabet_glow_tts_models: Arc>>>, @@ -20,17 +21,23 @@ pub struct ModelCache { impl ModelCache { pub fn new(model_locations: &Vec) -> 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> { - 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> { - 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> { - 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> { - 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::load(&file_path).ok() + } + + fn load_arpabet_tacotron_model(&self, file_path: &PathBuf) -> Option { + ArpabetTacotronModel::load(&file_path).ok() + } + + fn load_arpabet_glow_tts_model(&self, file_path: &PathBuf) -> Option { + ArpabetGlowTtsModel::load(&file_path).ok() + } + + fn load_arpabet_glow_tts_multi_speaker_model(&self, file_path: &PathBuf) -> Option { + 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) } } diff --git a/tts_service/src/model/model_config.rs b/tts_service/src/model/model_config.rs index cff886c4ff..8745f1f354 100644 --- a/tts_service/src/model/model_config.rs +++ b/tts_service/src/model/model_config.rs @@ -26,6 +26,9 @@ pub struct Speaker { pub glow_tts_multi_speaker: Option, /// Melgan model, if used. pub melgan: Option, + /// Whether the model is "production" or not + /// Production models are backed up on GDrive and are synced to DigitalOcean + pub production: Option, } /// 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, + pub dev_base_directory: Option, + pub prod_base_directory: Option, } /// The valid model pipelines in use by our system.