Update README.md

README.md

@@ -17,17 +17,34 @@ widget:
   src: https://huggingface.co/facebook/xm_transformer_600m-es_en-multi_domain/resolve/main/common_voice_es_19966634.flac
 ---
 
+## xm_transformer_s2ut_800m-es-en-st-asr-bt_h1_2022
+speech-to-speech translation model from fairseq S2UT (paper/code):
+
+-Spanish-English
+-Trained on 
+-Speech synthesis with facebook/unit_hifigan_mhubert_vp_en_es_fr_it3_400k_layer11_km1000_lj_dur
+
 ## Usage
 ```python
-from fairseq.checkpoint_utils import load_model_ensemble_and_task_from_hf_hub
-from fairseq.models.text_to_speech.hub_interface import S2THubInterface
-from fairseq.models.text_to_speech.hub_interface import TTSHubInterface
+import json
+import os
+from pathlib import Path
+
 import IPython.display as ipd
+from fairseq import hub_utils
+from fairseq.checkpoint_utils import load_model_ensemble_and_task_from_hf_hub
+from fairseq.models.speech_to_text.hub_interface import S2THubInterface
+from fairseq.models.text_to_speech import CodeHiFiGANVocoder
+from fairseq.models.text_to_speech.hub_interface import (
+    TTSHubInterface,
+    VocoderHubInterface,
+)
+from huggingface_hub import snapshot_download
 import torchaudio
 
 
 models, cfg, task = load_model_ensemble_and_task_from_hf_hub(
-    "facebook/xm_transformer_600m-es_en-multi_domain",
+    "facebook/xm_transformer_s2ut_800m-es-en-st-asr-bt_h1_2022",
     arg_overrides={"config_yaml": "config.yaml"},
 )
 model = models[0]
@@ -38,22 +55,41 @@ generator = task.build_generator(model, cfg)
 audio, _ = torchaudio.load("/path/to/an/audio/file")
 
 sample = S2THubInterface.get_model_input(task, audio)
-text = S2THubInterface.get_prediction(task, model, generator, sample)
+unit = S2THubInterface.get_prediction(task, model, generator, sample)
 
 # speech synthesis
-tts_models, tts_cfg, tts_task = load_model_ensemble_and_task_from_hf_hub(
-  f"facebook/fastspeech2-en-ljspeech",
-  arg_overrides={"vocoder": "griffin_lim", "fp16": False},
+cache_dir = os.getenv("HUGGINGFACE_HUB_CACHE")
+            
+library_name = "fairseq"
+cache_dir = (
+    cache_dir or (Path.home() / ".cache" / library_name).as_posix()
+)
+cache_dir = snapshot_download(
+    f"facebook/unit_hifigan_mhubert_vp_en_es_fr_it3_400k_layer11_km1000_lj_dur", cache_dir=cache_dir, library_name=library_name
 )
-tts_model = tts_models[0]
-TTSHubInterface.update_cfg_with_data_cfg(tts_cfg, tts_task.data_cfg)
-tts_generator = tts_task.build_generator([tts_model], tts_cfg)
 
-tts_sample = TTSHubInterface.get_model_input(tts_task, text)
-wav, sr = TTSHubInterface.get_prediction(
-    tts_task, tts_model, tts_generator, tts_sample
+x = hub_utils.from_pretrained(
+    cache_dir,
+    "model.pt",
+    ".",
+    archive_map=CodeHiFiGANVocoder.hub_models(),
+    config_yaml="config.json",
+    fp16=False,
+    is_vocoder=True,
 )
 
-ipd.Audio(wav, rate=rate)
+with open(f"{x['args']['data']}/config.json") as f:
+    vocoder_cfg = json.load(f)
+assert (
+    len(x["args"]["model_path"]) == 1
+), "Too many vocoder models in the input"
+
+vocoder = CodeHiFiGANVocoder(x["args"]["model_path"][0], vocoder_cfg)
+tts_model = VocoderHubInterface(vocoder_cfg, vocoder)
+
+tts_sample = tts_model.get_model_input(unit)
+wav, sr = tts_model.get_prediction(tts_sample)
+
+ipd.Audio(wav, rate=sr)
 ```