diff --git a/.env b/.env index 995c16971e..cef4ab21e2 100644 --- a/.env +++ b/.env @@ -45,6 +45,6 @@ TTS_MODEL_CHECK_SCRIPT_NAME=vocodes_model_check.py TTS_INFERENCE_SIDECAR_HOSTNAME=localhost:8000 SIDECAR_MAX_SYNTHESIZER_MODELS=2 -#TTS_VOCODER_MODEL_FILENAME=waveglow_256channels_universal_v5.pt # THIS SOUNDS HORRID -#TTS_VOCODER_MODEL_FILENAME=waveglow_256channels_v4_uberduck.pt # THIS IS TOO SMALL! 300MB vs 1.2 GB! -TTS_VOCODER_MODEL_FILENAME=waveglow_256channels_v4.pt +#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 diff --git a/models/tacotron2/vocodes_server.py b/models/tacotron2/vocodes_server.py index 21852fa4c1..dc120d46e4 100755 --- a/models/tacotron2/vocodes_server.py +++ b/models/tacotron2/vocodes_server.py @@ -210,7 +210,7 @@ class ApiHandler(): # Request parameters synthesizer_checkpoint_path = raw_data.get('synthesizer_checkpoint_path') - vocoder_checkpoint_path = raw_data.get('vocoder_checkpoint_path') + waveglow_vocoder_checkpoint_path = raw_data.get('waveglow_vocoder_checkpoint_path') inference_text = raw_data.get('inference_text') output_audio_filename = raw_data.get('output_audio_filename') output_spectrogram_filename = raw_data.get('output_spectrogram_filename') @@ -221,7 +221,7 @@ class ApiHandler(): raw_data.get('maybe_clear_synthesizer_checkpoint_path') print('synthesizer_checkpoint_path: {}'.format(synthesizer_checkpoint_path)) - print('vocoder_checkpoint_path: {}'.format(vocoder_checkpoint_path)) + print('waveglow_vocoder_checkpoint_path: {}'.format(waveglow_vocoder_checkpoint_path)) print('inference_text: {}'.format(inference_text)) print('output_audio_filename: {}'.format(output_audio_filename)) print('output_spectrogram_filename: {}'.format(output_spectrogram_filename)) @@ -232,7 +232,7 @@ class ApiHandler(): if maybe_clear_synthesizer_checkpoint_path: pipeline.uncache_tacotron_model(maybe_clear_synthesizer_checkpoint_path) - pipeline.maybe_load_waveglow_model(vocoder_checkpoint_path) + pipeline.maybe_load_waveglow_model(waveglow_vocoder_checkpoint_path) pipeline.maybe_load_tacotron_model_to_cache(synthesizer_checkpoint_path) pipeline.maybe_load_hifigan('/home/bt/models/hifigan/universal_hifigan_g_02500000') diff --git a/src/clients/tts_inference_sidecar_client.rs b/src/clients/tts_inference_sidecar_client.rs index 2d762df964..4699c740a3 100644 --- a/src/clients/tts_inference_sidecar_client.rs +++ b/src/clients/tts_inference_sidecar_client.rs @@ -5,17 +5,27 @@ 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, - //client: Client, } #[derive(Serialize)] struct InferenceRequest { - pub vocoder_checkpoint_path : String, + // Vocoder information + pub vocoder_type: VocoderType, + pub waveglow_vocoder_checkpoint_path : String, + pub hifigan_vocoder_checkpoint_path: String, + pub hifigan_supersample_vocoder_checkpoint_path: String, + + // Synthesizer information pub synthesizer_checkpoint_path : String, + + // Text pub inference_text : String, + + // Output information pub output_audio_filename : String, pub output_spectrogram_filename : String, pub output_metadata_filename : String, @@ -39,18 +49,33 @@ impl TtsInferenceSidecarClient { &self, raw_text: &str, synthesizer_checkpoint_path: P, - vocoder_checkpoint_path: P, + vocoder_type: VocoderType, + hifigan_vocoder_checkpoint_path: P, + hifigan_supersample_vocoder_checkpoint_path: P, + waveglow_vocoder_checkpoint_path: P, output_audio_filename: P, output_spectrogram_filename: P, output_metadata_filename: P, maybe_unload_model_path: Option, ) -> AnyhowResult<()> { - let vocoder_checkpoint_path = vocoder_checkpoint_path + let waveglow_vocoder_checkpoint_path = waveglow_vocoder_checkpoint_path .as_ref() .to_str() .map(|s| s.to_string()) - .ok_or(anyhow!("bad vocoder path"))?; + .ok_or(anyhow!("bad waveglow vocoder path"))?; + + let hifigan_vocoder_checkpoint_path = hifigan_vocoder_checkpoint_path + .as_ref() + .to_str() + .map(|s| s.to_string()) + .ok_or(anyhow!("bad hifigan vocoder path"))?; + + let hifigan_supersample_vocoder_checkpoint_path = hifigan_supersample_vocoder_checkpoint_path + .as_ref() + .to_str() + .map(|s| s.to_string()) + .ok_or(anyhow!("bad hifigan supersample vocoder path"))?; let synthesizer_checkpoint_path = synthesizer_checkpoint_path .as_ref() @@ -78,7 +103,10 @@ impl TtsInferenceSidecarClient { let request = InferenceRequest { inference_text: raw_text.to_string(), - vocoder_checkpoint_path, + vocoder_type, + waveglow_vocoder_checkpoint_path, + hifigan_vocoder_checkpoint_path, + hifigan_supersample_vocoder_checkpoint_path, synthesizer_checkpoint_path, output_audio_filename, output_spectrogram_filename, diff --git a/src/database/enums/mod.rs b/src/database/enums/mod.rs index fa9a069908..7ffceb80cc 100644 --- a/src/database/enums/mod.rs +++ b/src/database/enums/mod.rs @@ -1 +1,2 @@ -pub mod record_visibility; \ No newline at end of file +pub mod record_visibility; +pub mod vocoder_type; \ No newline at end of file diff --git a/src/database/enums/vocoder_type.rs b/src/database/enums/vocoder_type.rs new file mode 100644 index 0000000000..cb1ea153ae --- /dev/null +++ b/src/database/enums/vocoder_type.rs @@ -0,0 +1,29 @@ +//! These denote the pretrained vocoders + +use anyhow::anyhow; +use crate::util::anyhow_result::AnyhowResult; + +#[derive(Clone, Copy, PartialEq, Eq, Debug, Serialize, Deserialize, sqlx::Type)] +#[serde(rename_all = "lowercase")] +#[sqlx(rename_all = "lowercase")] +pub enum VocoderType { + HifiGanSuperSample, + WaveGlow, +} + +impl VocoderType { + pub fn to_str(&self) -> &'static str { + match self { + VocoderType::HifiGanSuperSample => "hifigan-supersample", + VocoderType::WaveGlow => "waveglow", + } + } + + pub fn from_str(vocoder_type: &str) -> AnyhowResult { + match vocoder_type { + "hifigan-supersample" => Ok(VocoderType::HifiGanSuperSample), + "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 6bcfd18f4d..e157d02c2e 100644 --- a/src/tts_inference_job.rs +++ b/src/tts_inference_job.rs @@ -61,6 +61,7 @@ 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"; @@ -102,8 +103,14 @@ struct Inferencer { pub virtual_model_lfu: SyncVirtualLfuCache, pub cache_miss_strategizers: SyncMultiCacheMissStrategizer, - // Command to run - pub tts_vocoder_model_filename: String, + // Waveglow vocoder filename + pub waveglow_vocoder_model_filename: String, + + // Hifigan vocoder filename + pub hifigan_vocoder_model_filename: String, + + // Hifigan supersampling vocoder filename + pub hifigan_supersampling_vocoder_model_filename: String, // Sleep between batches pub job_batch_wait_millis: u64, @@ -209,8 +216,8 @@ async fn main() -> AnyhowResult<()> { semi_persistent_cache.create_tts_synthesizer_model_path()?; semi_persistent_cache.create_tts_vocoder_model_path()?; - let tts_vocoder_model_filename = easyenv::get_env_string_or_default( - "TTS_VOCODER_MODEL_FILENAME", "waveglow.pth"); + let waveglow_vocoder_model_filename = easyenv::get_env_string_or_default( + "TTS_WAVEGLOW_VOCODER_MODEL_FILENAME", "waveglow.pth"); let sidecar_max_synthesizer_models = easyenv::get_env_num( "SIDECAR_MAX_SYNTHESIZER_MODELS", 3)?; @@ -253,9 +260,6 @@ async fn main() -> AnyhowResult<()> { mysql_pool, public_bucket_client, private_bucket_client, - //ffmpeg_image_preview_generator: FfmpegGeneratePreviewImageCommand {}, - //ffmpeg_video_preview_generator: FfmpegGeneratePreviewVideoCommand {}, - //imagemagick_image_preview_generator: ImagemagickGeneratePreviewImageCommand {}, tts_inference_command, tts_inference_sidecar_client, virtual_model_lfu: virtual_lfu_cache, @@ -263,7 +267,9 @@ async fn main() -> AnyhowResult<()> { bucket_path_unifier: BucketPathUnifier::default_paths(), semi_persistent_cache, firehose_publisher, - tts_vocoder_model_filename, + waveglow_vocoder_model_filename, + hifigan_vocoder_model_filename: "".to_string(), // TODO + hifigan_supersampling_vocoder_model_filename: "".to_string(), // TODO 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, @@ -507,25 +513,25 @@ async fn process_job( return Ok(()) } - // ==================== CONFIRM OR DOWNLOAD TTS VOCODER MODEL ==================== // + // ==================== CONFIRM OR DOWNLOAD WAVEGLOW VOCODER MODEL ==================== // - let tts_vocoder_model_filename = inferencer.tts_vocoder_model_filename.clone(); - let tts_vocoder_model_fs_path = inferencer.semi_persistent_cache.tts_vocoder_model_path(&tts_vocoder_model_filename); + 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); - if !tts_vocoder_model_fs_path.exists() { - warn!("Vocoder model file does not exist: {:?}", &tts_vocoder_model_fs_path); + if !waveglow_vocoder_model_fs_path.exists() { + warn!("Waveglow vocoder model file does not exist: {:?}", &waveglow_vocoder_model_fs_path); - let tts_vocoder_model_object_path = inferencer.bucket_path_unifier - .tts_pretrained_vocoders_path(&tts_vocoder_model_filename); + let waveglow_vocoder_model_object_path = inferencer.bucket_path_unifier + .tts_pretrained_vocoders_path(&waveglow_vocoder_model_filename); - info!("Download vocoder from bucket path: {:?}", &tts_vocoder_model_object_path); + info!("Download waveglow vocoder from bucket path: {:?}", &waveglow_vocoder_model_object_path); inferencer.private_bucket_client.download_file_to_disk( - &tts_vocoder_model_object_path, - &tts_vocoder_model_fs_path + &waveglow_vocoder_model_object_path, + &waveglow_vocoder_model_fs_path ).await?; - info!("Downloaded tts vocoder model from bucket!"); + info!("Downloaded waveglow vocoder model from bucket!"); } // // ==================== LOOK UP TTS SYNTHESIZER RECORD (WHICH CONTAINS ITS BUCKET PATH) ==================== // @@ -611,7 +617,10 @@ async fn process_job( inferencer.tts_inference_sidecar_client.request_inference( &job.raw_inference_text, &tts_synthesizer_fs_path, - &tts_vocoder_model_fs_path, + VocoderType::WaveGlow, + &PathBuf::new(), // TODO + &PathBuf::new(), // TODO + &waveglow_vocoder_model_fs_path, &output_audio_fs_path, &output_spectrogram_fs_path, &output_metadata_fs_path,