File size: 1,629 Bytes
83b7d03 |
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 |
import torch
from datasets import load_dataset
from transformers import Pipeline, SpeechT5Processor, SpeechT5HifiGan
class TTSPipeline(Pipeline):
def __init__(self, *args, vocoder=None, processor=None, **kwargs):
super().__init__(*args, **kwargs)
if vocoder is None:
raise ValueError("Must pass a vocoder to the TTSPipeline.")
if processor is None:
raise ValueError("Must pass a processor to the TTSPipeline.")
if isinstance(vocoder, str):
vocoder = SpeechT5HifiGan.from_pretrained(vocoder)
if isinstance(processor, str):
processor = SpeechT5Processor.from_pretrained(processor)
self.processor = processor
self.vocoder = vocoder
def preprocess(self, text, speaker_embeddings=None):
inputs = self.processor(text=text, return_tensors='pt')
if speaker_embeddings is None:
embeddings_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
speaker_embeddings = torch.tensor(embeddings_dataset[7305]["xvector"]).unsqueeze(0)
return {'inputs': inputs, 'speaker_embeddings': speaker_embeddings}
def _forward(self, model_inputs):
inputs = model_inputs['inputs']
speaker_embeddings = model_inputs['speaker_embeddings']
with torch.no_grad():
speech = self.model.generate_speech(inputs['input_ids'], speaker_embeddings, vocoder=self.vocoder)
return speech
def _sanitize_parameters(self, **pipeline_parameters):
return {}, {}, {}
def postprocess(self, speech):
return speech
|