From 90e1db42c481884152f9fb9f2d0e7e24dce24fde Mon Sep 17 00:00:00 2001
From: Brandon Thomas <echelon@gmail.com>
Date: Wed, 25 Aug 2021 02:21:59 -0400
Subject: [PATCH] Sidecar supports both vocoders and uses the one in the HTTP
 request

---
 models/tacotron2/vocodes_common.py | 13 ++++++++-----
 models/tacotron2/vocodes_server.py |  3 +++
 src/database/enums/vocoder_type.rs |  4 ++++
 src/tts_inference_job.rs           |  2 +-
 4 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/models/tacotron2/vocodes_common.py b/models/tacotron2/vocodes_common.py
index 3029392c17..423d5fe02b 100644
--- a/models/tacotron2/vocodes_common.py
+++ b/models/tacotron2/vocodes_common.py
@@ -234,6 +234,7 @@ class TacotronWaveglowPipeline:
 
     def infer(self, args):
         assert('synthesizer_checkpoint_path' in args)
+        assert('vocoder_type' in args)
         assert('raw_text' in args)
         assert('output_audio_filename' in args)
         assert('output_spectrogram_filename' in args)
@@ -246,13 +247,15 @@ class TacotronWaveglowPipeline:
         with torch.no_grad():
             mel_outputs, mel_outputs_postnet, _, alignments = tacotron.inference(sequence)
 
-        print('Saving spectrogram as JSON...')
-        save_spectorgram_json_file(mel_outputs_postnet, args['output_spectrogram_filename'])
-
-        if False:
+        if args['vocoder_type'] == 'hifigan-superres':
+            self.vocode_hifigan_superres(args, mel_outputs_postnet)
+        elif args['vocoder_type'] == 'waveglow':
             self.vocode_waveglow(args, mel_outputs_postnet)
         else:
-            self.vocode_hifigan_superres(args, mel_outputs_postnet)
+            raise Exception('Wrong vocoder type: {}'.format(args['vocoder_type']))
+
+        print('Saving spectrogram as JSON...')
+        save_spectorgram_json_file(mel_outputs_postnet, args['output_spectrogram_filename'])
 
         print('Generating metadata file...')
         generate_metadata_file(args['output_audio_filename'], args['output_metadata_filename'])
diff --git a/models/tacotron2/vocodes_server.py b/models/tacotron2/vocodes_server.py
index a15a57f983..e940b9b240 100755
--- a/models/tacotron2/vocodes_server.py
+++ b/models/tacotron2/vocodes_server.py
@@ -210,6 +210,7 @@ class ApiHandler():
 
         # Request parameters
         synthesizer_checkpoint_path = raw_data.get('synthesizer_checkpoint_path')
+        vocoder_type = raw_data.get('vocoder_type')
         waveglow_vocoder_checkpoint_path = raw_data.get('waveglow_vocoder_checkpoint_path')
         hifigan_vocoder_checkpoint_path = raw_data.get('hifigan_vocoder_checkpoint_path')
         hifigan_superres_vocoder_checkpoint_path = raw_data.get('hifigan_superres_vocoder_checkpoint_path')
@@ -223,6 +224,7 @@ class ApiHandler():
             raw_data.get('maybe_clear_synthesizer_checkpoint_path')
 
         print('synthesizer_checkpoint_path: {}'.format(synthesizer_checkpoint_path))
+        print('vocoder_type: {}'.format(vocoder_type))
         print('waveglow_vocoder_checkpoint_path: {}'.format(waveglow_vocoder_checkpoint_path))
         print('hifigan_vocoder_checkpoint_path: {}'.format(hifigan_vocoder_checkpoint_path))
         print('hifigan_superres_vocoder_checkpoint_path: {}'.format(hifigan_superres_vocoder_checkpoint_path))
@@ -245,6 +247,7 @@ class ApiHandler():
         inference_args = {
             'raw_text': inference_text,
             'synthesizer_checkpoint_path': synthesizer_checkpoint_path,
+            'vocoder_type': vocoder_type,
             'output_audio_filename': output_audio_filename,
             'output_spectrogram_filename': output_spectrogram_filename,
             'output_metadata_filename': output_metadata_filename,
diff --git a/src/database/enums/vocoder_type.rs b/src/database/enums/vocoder_type.rs
index e75d3209f8..8ebcb54278 100644
--- a/src/database/enums/vocoder_type.rs
+++ b/src/database/enums/vocoder_type.rs
@@ -7,7 +7,11 @@ use crate::util::anyhow_result::AnyhowResult;
 #[serde(rename_all = "lowercase")]
 #[sqlx(rename_all = "lowercase")]
 pub enum VocoderType {
+  /// Current state of the art vocoder
+  #[serde(rename = "hifigan-superres")]
   HifiGanSuperResolution,
+
+  /// Older WaveGlow vocoder
   WaveGlow,
 }
 
diff --git a/src/tts_inference_job.rs b/src/tts_inference_job.rs
index b59ce4608b..9a1291ecfa 100644
--- a/src/tts_inference_job.rs
+++ b/src/tts_inference_job.rs
@@ -665,7 +665,7 @@ async fn process_job(
   inferencer.tts_inference_sidecar_client.request_inference(
     &job.raw_inference_text,
     &tts_synthesizer_fs_path,
-    VocoderType::WaveGlow,
+    VocoderType::HifiGanSuperResolution,
     &hifigan_vocoder_model_fs_path,
     &hifigan_superres_vocoder_model_fs_path,
     &waveglow_vocoder_model_fs_path,