Spaces:
Runtime error
Runtime error
File size: 6,662 Bytes
32e2145 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 |
import json
import os
from pathlib import Path
from typing import List, Tuple
import tempfile
import soundfile as sf
import gradio as gr
import numpy as np
import torch
import torchaudio
# from app.pipelines import Pipeline
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
ARG_OVERRIDES_MAP = {
"facebook/xm_transformer_s2ut_800m-es-en-st-asr-bt_h1_2022": {
"config_yaml": "config.yaml",
"task": "speech_to_text",
}
}
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: str) -> 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)
# print(f"input: {inputs}")
# _inputs = torchaudio.load(inputs)
_inputs = inputs
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
)
temp_file = ""
with tempfile.NamedTemporaryFile(suffix=".wav") as tmp_output_file:
sf.write(tmp_output_file, wav.detach().cpu().numpy(), sr)
tmp_output_file.seek(0)
temp_file = gr.Audio(tmp_output_file.name)
# return wav, sr, [text]
return temp_file |