intermediate plumbing for new vocoders (2)

.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

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

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,

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

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

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(())
   }