VTTS-speechT5 / src /model.py
linh-truong's picture
update remove_special_characters
45bafcb
raw
history blame
5.54 kB
import re
import torch
import requests
import torchaudio
import numpy as np
# from src.reduce_noise import smooth_and_reduce_noise, model_remove_noise, model, df_state
import io
from transformers import SpeechT5Processor, SpeechT5ForTextToSpeech, SpeechT5HifiGan
from pydub import AudioSegment
import re
from uroman import uroman
# from src.pynote_speaker_embedding import create_speaker_embedding
from src.speechbrain_speaker_embedding import create_speaker_embedding
from datasets import load_dataset
dataset = load_dataset("truong-xuan-linh/vi-xvector-speechbrain",
download_mode="force_redownload",
verification_mode="no_checks",
cache_dir="temp/",
revision="5ea5e4345258333cbc6d1dd2544f6c658e66a634")
dataset = dataset["train"].to_list()
dataset_dict = {}
for rc in dataset:
dataset_dict[rc["speaker_id"]] = rc["embedding"]
vocoder = SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan")
def remove_special_characters(sentence):
# Use regular expression to keep only letters, periods, and commas
sentence_after_removal = re.sub(r'[^a-zA-Z\s,.\u00C0-\u1EF9]', ' ,', sentence)
return sentence_after_removal
from scipy.signal import butter, lfilter
def butter_bandpass(lowcut, highcut, fs, order=5):
nyq = 0.5 * fs
low = lowcut / nyq
high = highcut / nyq
b, a = butter(order, [low, high], btype='band')
return b, a
def butter_bandpass_filter(data, lowcut, highcut, fs, order=5):
b, a = butter_bandpass(lowcut, highcut, fs, order=order)
y = lfilter(b, a, data)
return y
def korean_splitter(string):
pattern = re.compile('[가-힣]+')
matches = pattern.findall(string)
return matches
def uroman_normalization(string):
korean_inputs = korean_splitter(string)
for korean_input in korean_inputs:
korean_roman = uroman(korean_input)
string = string.replace(korean_input, korean_roman)
return string
class Model():
def __init__(self, model_name, speaker_url=""):
self.model_name = model_name
self.processor = SpeechT5Processor.from_pretrained(model_name)
self.model = SpeechT5ForTextToSpeech.from_pretrained(model_name)
# self.model.generate = partial(self.model.generate, use_cache=True)
self.model.eval()
self.speaker_url = speaker_url
if speaker_url:
print(f"download speaker_url")
response = requests.get(speaker_url)
audio_stream = io.BytesIO(response.content)
audio_segment = AudioSegment.from_file(audio_stream, format="wav")
audio_segment = audio_segment.set_channels(1)
audio_segment = audio_segment.set_frame_rate(16000)
audio_segment = audio_segment.set_sample_width(2)
wavform, _ = torchaudio.load(audio_segment.export())
self.speaker_embeddings = create_speaker_embedding(wavform)[0]
else:
self.speaker_embeddings = None
if model_name == "truong-xuan-linh/speecht5-vietnamese-commonvoice" or model_name == "truong-xuan-linh/speecht5-irmvivoice":
self.speaker_embeddings = torch.zeros((1, 512)) # or load xvectors from a file
def inference(self, text, speaker_id=None):
# if self.model_name == "truong-xuan-linh/speecht5-vietnamese-voiceclone-v2":
# # self.speaker_embeddings = torch.tensor(dataset_dict_v2[speaker_id])
# wavform, _ = torchaudio.load(speaker_id)
# self.speaker_embeddings = create_speaker_embedding(wavform)[0]
if "voiceclone" in self.model_name:
if not self.speaker_url:
self.speaker_embeddings = torch.tensor(dataset_dict[speaker_id])
# self.speaker_embeddings = create_speaker_embedding(speaker_id)[0]
# wavform, _ = torchaudio.load("voices/kcbn1.wav")
# self.speaker_embeddings = create_speaker_embedding(wavform)[0]
# wavform, _ = torchaudio.load(wav_file)
# self.speaker_embeddings = create_speaker_embedding(wavform)[0]
with torch.no_grad():
full_speech = []
separators = r";|\.|!|\?|\n"
text = uroman_normalization(text)
text = remove_special_characters(text)
text = text.replace(" ", "▁")
split_texts = re.split(separators, text)
for split_text in split_texts:
if split_text != "▁":
split_text = split_text.lower() + "▁"
print(split_text)
inputs = self.processor.tokenizer(text=split_text, return_tensors="pt")
speech = self.model.generate_speech(inputs["input_ids"], threshold=0.5, speaker_embeddings=self.speaker_embeddings, vocoder=vocoder)
full_speech.append(speech.numpy())
# full_speech.append(butter_bandpass_filter(speech.numpy(), lowcut=10, highcut=5000, fs=16000, order=2))
# out_audio = model_remove_noise(model, df_state, np.concatenate(full_speech))
return np.concatenate(full_speech)
@staticmethod
def moving_average(data, window_size):
return np.convolve(data, np.ones(window_size)/window_size, mode='same')
# woman: VIVOSSPK26, VIVOSSPK02, VIVOSSPK40
# man: VIVOSSPK28, VIVOSSPK36, VIVOSDEV09, VIVOSSPK33, VIVOSSPK23