Fix enhancer usage in ASR-TTS examples (NVIDIA#6116) (NVIDIA#6117)

Signed-off-by: Vladimir Bataev <[email protected]> Co-authored-by: Vladimir Bataev <[email protected]> Signed-off-by: hsiehjackson <[email protected]>
hsiehjackson · Jun 2, 2023 · 1b835ba · 1b835ba
1 parent 1f7721c
commit 1b835ba
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/examples/asr/asr_with_tts/speech_to_text_bpe_with_text.py b/examples/asr/asr_with_tts/speech_to_text_bpe_with_text.py
@@ -21,6 +21,7 @@
     # (Optional: --config-path=<path to dir of configs> --config-name=<name of config without .yaml>) \
     ++asr_model_type=<rnnt_bpe or ctc_bpe> \
     ++tts_model_path=<path to compatible tts model> \
+    ++enhancer_model_path=<optional path to enhancer model> \
     model.tokenizer.dir=<path to tokenizer> \
     model.tokenizer.type="bpe" \
     model.train_ds.manifest_filepath=<path(s) to manifest with audio-text pairs or null> \
@@ -70,7 +71,11 @@ def main(cfg):
     exp_manager(trainer, cfg.get("exp_manager", None))
 
     asr_model = ASRWithTTSModel.from_asr_config(
-        asr_cfg=cfg.model, asr_model_type=cfg.asr_model_type, tts_model_path=cfg.tts_model_path, trainer=trainer
+        asr_cfg=cfg.model,
+        asr_model_type=cfg.asr_model_type,
+        tts_model_path=cfg.tts_model_path,
+        enhancer_model_path=cfg.get("enhancer_model_path", None),
+        trainer=trainer,
     )
 
     # Initialize the weights of the model from another model, if provided via config

diff --git a/examples/asr/asr_with_tts/speech_to_text_bpe_with_text_finetune.py b/examples/asr/asr_with_tts/speech_to_text_bpe_with_text_finetune.py
@@ -19,6 +19,7 @@
     # (Optional: --config-path=<path to dir of configs> --config-name=<name of config without .yaml>) \
     model.asr_model_path=<path to ASR model> \
     model.tts_model_path=<path to compatible TTS model> \
+    model.enhancer_model_path=<optional path to enhancer model> \
     model.asr_model_fuse_bn=<true recommended if ConformerEncoder with BatchNorm, false otherwise> \
     model.train_ds.manifest_filepath=<path to manifest with audio-text pairs or null> \
     model.train_ds.text_data.manifest_filepath=<path(s) to manifest with train text> \