usrbinenv python codingutf8 2023 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license usr bin env python coding utf 8 2023 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license
import torch

from ..models.speecht5 import SpeechT5ForTextToSpeech, SpeechT5HifiGan, SpeechT5Processor
from ..utils import is_datasets_available
from .base import PipelineTool


if is_datasets_available():
    from datasets import load_dataset


class TextToSpeechTool(PipelineTool):
    default_checkpoint = "microsoft/speecht5_tts"
    description = (
        "This is a tool that reads an English text out loud. It takes an input named `text` which should contain the "
        "text to read (in English) and returns a waveform object containing the sound."
    )
    name = "text_reader"
    pre_processor_class = SpeechT5Processor
    model_class = SpeechT5ForTextToSpeech
    post_processor_class = SpeechT5HifiGan

    inputs = ["text"]
    outputs = ["audio"]

    def setup(self):
        if self.post_processor is None:
            self.post_processor = "microsoft/speecht5_hifigan"
        super().setup()

    def encode(self, text, speaker_embeddings=None):
        inputs = self.pre_processor(text=text, return_tensors="pt", truncation=True)

        if speaker_embeddings is None:
            if not is_datasets_available():
                raise ImportError("Datasets needs to be installed if not passing speaker embeddings.")

            embeddings_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
            speaker_embeddings = torch.tensor(embeddings_dataset[7305]["xvector"]).unsqueeze(0)

        return {"input_ids": inputs["input_ids"], "speaker_embeddings": speaker_embeddings}

    def forward(self, inputs):
        with torch.no_grad():
            return self.model.generate_speech(**inputs)

    def decode(self, outputs):
        with torch.no_grad():
            return self.post_processor(outputs).cpu().detach()