intermediate plumbing for new vocoders

.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

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')

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<String>,
   ) -> 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,

src/database/enums/mod.rs

@@ -1 +1,2 @@
-pub mod record_visibility;
+pub mod record_visibility;
+pub mod vocoder_type;

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<Self> {
+    match vocoder_type {
+      "hifigan-supersample" => Ok(VocoderType::HifiGanSuperSample),
+      "waveglow" => Ok(VocoderType::WaveGlow),
+      _ => Err(anyhow!("invalid value: {:?}", vocoder_type)),
+    }
+  }
+}

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,