import json import os from pathlib import Path from typing import List, Tuple import numpy as np import torch # from app.pipelines import Pipeline from app.pipelines.utils import ARG_OVERRIDES_MAP from fairseq import hub_utils from fairseq.checkpoint_utils import load_model_ensemble_and_task_from_hf_hub from fairseq.models.speech_to_speech.hub_interface import S2SHubInterface 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 class SpeechToSpeechPipeline(): def __init__(self, model_id: str): arg_overrides = ARG_OVERRIDES_MAP.get( model_id, {} ) # Model specific override. TODO: Update on checkpoint side in the future arg_overrides["config_yaml"] = "config.yaml" # common override models, cfg, task = load_model_ensemble_and_task_from_hf_hub( model_id, arg_overrides=arg_overrides, cache_dir=os.getenv("HUGGINGFACE_HUB_CACHE"), ) self.cfg = cfg self.model = models[0].cpu() self.model.eval() self.task = task self.sampling_rate = getattr(self.task, "sr", None) or 16_000 tgt_lang = self.task.data_cfg.hub.get("tgt_lang", None) pfx = f"{tgt_lang}_" if self.task.data_cfg.prepend_tgt_lang_tag else "" generation_args = self.task.data_cfg.hub.get(f"{pfx}generation_args", None) if generation_args is not None: for key in generation_args: setattr(cfg.generation, key, generation_args[key]) self.generator = task.build_generator([self.model], cfg.generation) tts_model_id = self.task.data_cfg.hub.get(f"{pfx}tts_model_id", None) self.unit_vocoder = self.task.data_cfg.hub.get(f"{pfx}unit_vocoder", None) self.tts_model, self.tts_task, self.tts_generator = None, None, None if tts_model_id is not None: _id = tts_model_id.split(":")[-1] cache_dir = os.getenv("HUGGINGFACE_HUB_CACHE") if self.unit_vocoder is not None: library_name = "fairseq" cache_dir = ( cache_dir or (Path.home() / ".cache" / library_name).as_posix() ) cache_dir = snapshot_download( f"facebook/{_id}", cache_dir=cache_dir, library_name=library_name ) x = hub_utils.from_pretrained( cache_dir, "model.pt", ".", archive_map=CodeHiFiGANVocoder.hub_models(), config_yaml="config.json", fp16=False, is_vocoder=True, ) 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) self.tts_model = VocoderHubInterface(vocoder_cfg, vocoder) else: ( tts_models, tts_cfg, self.tts_task, ) = load_model_ensemble_and_task_from_hf_hub( f"facebook/{_id}", arg_overrides={"vocoder": "griffin_lim", "fp16": False}, cache_dir=cache_dir, ) self.tts_model = tts_models[0].cpu() self.tts_model.eval() tts_cfg["task"].cpu = True TTSHubInterface.update_cfg_with_data_cfg( tts_cfg, self.tts_task.data_cfg ) self.tts_generator = self.tts_task.build_generator( [self.tts_model], tts_cfg ) def __call__(self, inputs: np.array) -> Tuple[np.array, int, List[str]]: """ Args: inputs (:obj:`np.array`): The raw waveform of audio received. By default sampled at `self.sampling_rate`. The shape of this array is `T`, where `T` is the time axis Return: A :obj:`tuple` containing: - :obj:`np.array`: The return shape of the array must be `C'`x`T'` - a :obj:`int`: the sampling rate as an int in Hz. - a :obj:`List[str]`: the annotation for each out channel. This can be the name of the instruments for audio source separation or some annotation for speech enhancement. The length must be `C'`. """ _inputs = torch.from_numpy(inputs).unsqueeze(0) sample, text = None, None if self.cfg.task._name in ["speech_to_text", "speech_to_text_sharded"]: sample = S2THubInterface.get_model_input(self.task, _inputs) text = S2THubInterface.get_prediction( self.task, self.model, self.generator, sample ) elif self.cfg.task._name in ["speech_to_speech"]: s2shubinerface = S2SHubInterface(self.cfg, self.task, self.model) sample = s2shubinerface.get_model_input(self.task, _inputs) text = S2SHubInterface.get_prediction( self.task, self.model, self.generator, sample ) wav, sr = np.zeros((0,)), self.sampling_rate if self.unit_vocoder is not None: tts_sample = self.tts_model.get_model_input(text) wav, sr = self.tts_model.get_prediction(tts_sample) text = "" else: tts_sample = TTSHubInterface.get_model_input(self.tts_task, text) wav, sr = TTSHubInterface.get_prediction( self.tts_task, self.tts_model, self.tts_generator, tts_sample ) return wav.unsqueeze(0).numpy(), sr, [text]