diff --git a/.env b/.env index cef4ab21e2..9793be9269 100644 --- a/.env +++ b/.env @@ -48,3 +48,6 @@ SIDECAR_MAX_SYNTHESIZER_MODELS=2 #TTS_WAVEGLOW_VOCODER_MODEL_FILENAME=waveglow_256channels_universal_v5.pt # THIS SOUNDS HORRID #TTS_WAVEGLOW_VOCODER_MODEL_FILENAME=waveglow_256channels_v4_uberduck.pt # THIS IS TOO SMALL! 300MB vs 1.2 GB! TTS_WAVEGLOW_VOCODER_MODEL_FILENAME=waveglow_256channels_v4.pt + +TTS_HIFIGAN_VOCODER_MODEL_FILENAME=todo +TTS_HIFIGAN_SUPERRES_VOCODER_MODEL_FILENAME=todo diff --git a/models/tacotron2/vocodes_common.py b/models/tacotron2/vocodes_common.py index 7f79cb5c47..3fa33bd77e 100644 --- a/models/tacotron2/vocodes_common.py +++ b/models/tacotron2/vocodes_common.py @@ -259,7 +259,7 @@ class TacotronWaveglowPipeline: print('Generating metadata file...') generate_metadata_file(args['output_audio_filename'], args['output_metadata_filename']) - #return + return with torch.no_grad(): print('Running hifigan...') diff --git a/src/clients/tts_inference_sidecar_client.rs b/src/clients/tts_inference_sidecar_client.rs index 4699c740a3..354cdfd1b0 100644 --- a/src/clients/tts_inference_sidecar_client.rs +++ b/src/clients/tts_inference_sidecar_client.rs @@ -1,11 +1,11 @@ use actix_web::http::header; use anyhow::anyhow; +use crate::database::enums::vocoder_type::VocoderType; use crate::util::anyhow_result::AnyhowResult; use hyper::client::Client; use hyper::{Body, Request}; use log::info; use std::path::Path; -use crate::database::enums::vocoder_type::VocoderType; pub struct TtsInferenceSidecarClient { hostname: String, @@ -17,7 +17,7 @@ struct InferenceRequest { pub vocoder_type: VocoderType, pub waveglow_vocoder_checkpoint_path : String, pub hifigan_vocoder_checkpoint_path: String, - pub hifigan_supersample_vocoder_checkpoint_path: String, + pub hifigan_superres_vocoder_checkpoint_path: String, // Synthesizer information pub synthesizer_checkpoint_path : String, @@ -51,7 +51,7 @@ impl TtsInferenceSidecarClient { synthesizer_checkpoint_path: P, vocoder_type: VocoderType, hifigan_vocoder_checkpoint_path: P, - hifigan_supersample_vocoder_checkpoint_path: P, + hifigan_superres_vocoder_checkpoint_path: P, waveglow_vocoder_checkpoint_path: P, output_audio_filename: P, output_spectrogram_filename: P, @@ -71,11 +71,11 @@ impl TtsInferenceSidecarClient { .map(|s| s.to_string()) .ok_or(anyhow!("bad hifigan vocoder path"))?; - let hifigan_supersample_vocoder_checkpoint_path = hifigan_supersample_vocoder_checkpoint_path + let hifigan_superres_vocoder_checkpoint_path = hifigan_superres_vocoder_checkpoint_path .as_ref() .to_str() .map(|s| s.to_string()) - .ok_or(anyhow!("bad hifigan supersample vocoder path"))?; + .ok_or(anyhow!("bad hifigan super resolution vocoder path"))?; let synthesizer_checkpoint_path = synthesizer_checkpoint_path .as_ref() @@ -106,7 +106,7 @@ impl TtsInferenceSidecarClient { vocoder_type, waveglow_vocoder_checkpoint_path, hifigan_vocoder_checkpoint_path, - hifigan_supersample_vocoder_checkpoint_path, + hifigan_superres_vocoder_checkpoint_path, synthesizer_checkpoint_path, output_audio_filename, output_spectrogram_filename, diff --git a/src/database/enums/vocoder_type.rs b/src/database/enums/vocoder_type.rs index cb1ea153ae..e75d3209f8 100644 --- a/src/database/enums/vocoder_type.rs +++ b/src/database/enums/vocoder_type.rs @@ -7,21 +7,21 @@ use crate::util::anyhow_result::AnyhowResult; #[serde(rename_all = "lowercase")] #[sqlx(rename_all = "lowercase")] pub enum VocoderType { - HifiGanSuperSample, + HifiGanSuperResolution, WaveGlow, } impl VocoderType { pub fn to_str(&self) -> &'static str { match self { - VocoderType::HifiGanSuperSample => "hifigan-supersample", + VocoderType::HifiGanSuperResolution => "hifigan-superres", VocoderType::WaveGlow => "waveglow", } } pub fn from_str(vocoder_type: &str) -> AnyhowResult { match vocoder_type { - "hifigan-supersample" => Ok(VocoderType::HifiGanSuperSample), + "hifigan-superres" => Ok(VocoderType::HifiGanSuperResolution), "waveglow" => Ok(VocoderType::WaveGlow), _ => Err(anyhow!("invalid value: {:?}", vocoder_type)), } diff --git a/src/tts_inference_job.rs b/src/tts_inference_job.rs index e157d02c2e..9bfc2b387c 100644 --- a/src/tts_inference_job.rs +++ b/src/tts_inference_job.rs @@ -22,6 +22,7 @@ use anyhow::{anyhow, Error}; use chrono::{Utc, DateTime, TimeZone}; use crate::clients::tts_inference_sidecar_client::TtsInferenceSidecarClient; use crate::common_env::CommonEnv; +use crate::database::enums::vocoder_type::VocoderType; use crate::database::mediators::firehose_publisher::FirehosePublisher; use crate::job_queries::tts_inference_job_queries::get_tts_model_by_token; use crate::job_queries::tts_inference_job_queries::grab_job_lock_and_mark_pending; @@ -61,7 +62,6 @@ use std::process::Command; use std::thread; use std::time::Duration; use tempdir::TempDir; -use crate::database::enums::vocoder_type::VocoderType; // Buckets (shared config) const ENV_ACCESS_KEY : &'static str = "ACCESS_KEY"; @@ -109,8 +109,8 @@ struct Inferencer { // Hifigan vocoder filename pub hifigan_vocoder_model_filename: String, - // Hifigan supersampling vocoder filename - pub hifigan_supersampling_vocoder_model_filename: String, + // Hifigan super resolution vocoder filename + pub hifigan_superres_vocoder_model_filename: String, // Sleep between batches pub job_batch_wait_millis: u64, @@ -214,11 +214,17 @@ async fn main() -> AnyhowResult<()> { info!("Creating pod semi-persistent cache dirs..."); semi_persistent_cache.create_tts_synthesizer_model_path()?; - semi_persistent_cache.create_tts_vocoder_model_path()?; + semi_persistent_cache.create_tts_pretrained_vocoder_model_path()?; let waveglow_vocoder_model_filename = easyenv::get_env_string_or_default( "TTS_WAVEGLOW_VOCODER_MODEL_FILENAME", "waveglow.pth"); + let hifigan_vocoder_model_filename = easyenv::get_env_string_or_default( + "TTS_HIFIGAN_VOCODER_MODEL_FILENAME", "hifigan.pth"); + + let hifigan_superres_vocoder_model_filename = easyenv::get_env_string_or_default( + "TTS_HIFIGAN_SUPERRES_VOCODER_MODEL_FILENAME", "hifigan_superres.pth"); + let sidecar_max_synthesizer_models = easyenv::get_env_num( "SIDECAR_MAX_SYNTHESIZER_MODELS", 3)?; @@ -268,8 +274,8 @@ async fn main() -> AnyhowResult<()> { semi_persistent_cache, firehose_publisher, waveglow_vocoder_model_filename, - hifigan_vocoder_model_filename: "".to_string(), // TODO - hifigan_supersampling_vocoder_model_filename: "".to_string(), // TODO + hifigan_vocoder_model_filename, + hifigan_superres_vocoder_model_filename, job_batch_wait_millis: common_env.job_batch_wait_millis, job_max_attempts: common_env.job_max_attempts as i32, job_batch_size: common_env.job_batch_size, @@ -516,7 +522,7 @@ async fn process_job( // ==================== CONFIRM OR DOWNLOAD WAVEGLOW VOCODER MODEL ==================== // let waveglow_vocoder_model_filename = inferencer.waveglow_vocoder_model_filename.clone(); - let waveglow_vocoder_model_fs_path = inferencer.semi_persistent_cache.tts_vocoder_model_path(&waveglow_vocoder_model_filename); + let waveglow_vocoder_model_fs_path = inferencer.semi_persistent_cache.tts_pretrained_vocoder_model_path(&waveglow_vocoder_model_filename); if !waveglow_vocoder_model_fs_path.exists() { warn!("Waveglow vocoder model file does not exist: {:?}", &waveglow_vocoder_model_fs_path); diff --git a/src/util/semi_persistent_cache_dir.rs b/src/util/semi_persistent_cache_dir.rs index 5cd05d51f5..93a34ad372 100644 --- a/src/util/semi_persistent_cache_dir.rs +++ b/src/util/semi_persistent_cache_dir.rs @@ -8,7 +8,7 @@ pub struct SemiPersistentCacheDir { cache_root: PathBuf, tts_synthesizer_model_root: PathBuf, - tts_vocoder_model_root: PathBuf, + tts_pretrained_vocoder_model_root: PathBuf, w2l_model_root: PathBuf, w2l_end_bump_root: PathBuf, @@ -32,7 +32,7 @@ impl SemiPersistentCacheDir { cache_root: cache_root.clone(), tts_synthesizer_model_root: cache_root.join("tts/synthesizer_models/"), - tts_vocoder_model_root: cache_root.join("tts/vocoder_models/"), + tts_pretrained_vocoder_model_root: cache_root.join("tts/vocoder_models_pretrained/"), w2l_model_root: cache_root.join("w2l/models/"), w2l_end_bump_root: cache_root.join("w2l/end_bumps/"), @@ -59,19 +59,19 @@ impl SemiPersistentCacheDir { Ok(()) } - // ==================== TTS VOCODER MODELS ==================== + // ==================== TTS PRETRAINED VOCODER MODELS (WAVEGLOW, HIFIGAN, ETC) ==================== /// We'll start with just a few vocoders, but this may grow. - pub fn tts_vocoder_model_path(&self, model_filename: &str) -> PathBuf { - self.tts_vocoder_model_root.join(model_filename) + pub fn tts_pretrained_vocoder_model_path(&self, model_filename: &str) -> PathBuf { + self.tts_pretrained_vocoder_model_root.join(model_filename) } - pub fn tts_vocoder_model_directory(&self) -> &Path { - &self.tts_vocoder_model_root + pub fn tts_pretrained_vocoder_model_directory(&self) -> &Path { + &self.tts_pretrained_vocoder_model_root } - pub fn create_tts_vocoder_model_path(&self) -> AnyhowResult<()> { - let _ = fs::create_dir_all(self.tts_vocoder_model_directory())?; + pub fn create_tts_pretrained_vocoder_model_path(&self) -> AnyhowResult<()> { + let _ = fs::create_dir_all(self.tts_pretrained_vocoder_model_directory())?; Ok(()) }