Upload 13 files

app.py

@@ -0,0 +1,33 @@
+import gradio as gr
+from infer_onnx import TTS
+from ruaccent import RUAccent # https://huggingface.co/TeraTTS/accentuator
+
+models = ["TeraTTS/natasha-g2p-vits", "TeraTTS/glados2-g2p-vits"]
+
+models = {k:TTS(k) for k in models}
+
+accentizer = RUAccent(workdir="./model/ruaccent")
+accentizer.load(omograph_model_size='medium', dict_load_startup=True)
+
+
+def process_text(text: str) -> str:
+    text = accentizer.process_all(text)
+    return text
+
+def text_to_speech(model_name, text, prep_text):
+    if prep_text:
+        text = process_text(text)
+    audio = models[model_name](text)
+    models[model_name].save_wav(audio, 'temp.wav')
+
+    return 'temp.wav', f"Обработанный текст: '{text}'"
+
+model_choice = gr.Dropdown(choices=list(models.keys()), value="TeraTTS/natasha-g2p-vits", label="Выберите модель")
+input_text = gr.Textbox(label="Введите текст для синтеза речи")
+prep_text = gr.Checkbox(label="Предобработать", info="Хотите пред обработать текст?(Ударения, ё)", value=True)
+
+output_audio = gr.Audio(label="Аудио", type="numpy")
+output_text = gr.Textbox(label="Обработанный текст")
+
+iface = gr.Interface(fn=text_to_speech, inputs=[model_choice, input_text, prep_text], outputs=[output_audio, output_text])
+iface.launch()

infer_onnx.py

@@ -0,0 +1,79 @@
+import scipy.io.wavfile
+import os
+import onnxruntime
+import numpy as np
+from huggingface_hub import snapshot_download
+
+class TTS:
+    def __init__(self, model_name: str, save_path: str = "./model", add_time_to_end: float = 0.8) -> None:
+        if not os.path.exists(save_path):
+            os.mkdir(save_path)
+        
+        model_dir = os.path.join(save_path, model_name)
+        
+        if not os.path.exists(model_dir):
+            snapshot_download(repo_id=model_name, 
+                              allow_patterns=["*.txt", "*.onnx", "*.json"], 
+                              local_dir=model_dir,
+                              local_dir_use_symlinks=False
+                            )
+        
+        sess_options = onnxruntime.SessionOptions()
+        self.model = onnxruntime.InferenceSession(os.path.join(model_dir, "exported/model.onnx"), sess_options=sess_options)
+        
+        if os.path.exists(os.path.join(model_dir, "exported/dictionary.txt")):
+            from tokenizer import TokenizerG2P
+            print("Use g2p")
+            self.tokenizer = TokenizerG2P(os.path.join(model_dir, "exported"))
+            
+        else:
+            from tokenizer import TokenizerGRUUT
+            print("Use gruut")
+            self.tokenizer = TokenizerGRUUT(os.path.join(model_dir, "exported"))
+        
+        self.add_time_to_end = add_time_to_end
+
+    
+    def _add_silent(self, audio, silence_duration: float = 1.0, sample_rate: int = 22050):
+        num_samples_silence = int(sample_rate * silence_duration)
+        silence_array = np.zeros(num_samples_silence, dtype=np.float32)
+        audio_with_silence = np.concatenate((audio, silence_array), axis=0)
+        return audio_with_silence
+
+
+    def save_wav(self, audio, path:str):
+        '''save audio to wav'''
+        scipy.io.wavfile.write(path, 22050, audio)
+    
+    
+    def _intersperse(self, lst, item):
+        result = [item] * (len(lst) * 2 + 1)
+        result[1::2] = lst
+        return result
+    
+    
+    def _get_seq(self, text):
+        phoneme_ids = self.tokenizer._get_seq(text)
+        phoneme_ids_inter = self._intersperse(phoneme_ids, 0)
+        return phoneme_ids_inter
+        
+    
+    def __call__(self, text: str):
+        phoneme_ids = self._get_seq(text)
+        text = np.expand_dims(np.array(phoneme_ids, dtype=np.int64), 0)
+        text_lengths = np.array([text.shape[1]], dtype=np.int64)
+        scales = np.array(
+            [0.667, 1, 0.8],
+            dtype=np.float32,
+        )
+        audio = self.model.run(
+            None,
+            {
+                "input": text,
+                "input_lengths": text_lengths,
+                "scales": scales,
+                "sid": None,
+            },
+        )[0][0,0][0]
+        audio = self._add_silent(audio, silence_duration = self.add_time_to_end)
+        return audio

ruaccent/__init__.py

@@ -0,0 +1 @@
+from .ruaccent import RUAccent

ruaccent/accent_model.py

@@ -0,0 +1,27 @@
+import torch
+from .char_tokenizer import CharTokenizer
+from transformers import AutoModelForTokenClassification
+
+class AccentModel:
+    def __init__(self) -> None:
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    def load(self, path):
+        self.model = AutoModelForTokenClassification.from_pretrained(path).to(self.device)
+        self.tokenizer = CharTokenizer.from_pretrained(path)
+    
+    def render_stress(self, word, token_classes):
+        if 'STRESS' in token_classes:
+            index = token_classes.index('STRESS')
+            word = list(word)
+            word[index-1] = '+' + word[index-1]
+            return ''.join(word)
+        else:
+            return word
+    
+    def put_accent(self, word):
+        inputs = self.tokenizer(word, return_tensors="pt").to(self.device)
+        with torch.no_grad():
+            logits = self.model(**inputs).logits
+            predictions = torch.argmax(logits, dim=2)
+            predicted_token_class = [self.model.config.id2label[t.item()] for t in predictions[0]]
+        return self.render_stress(word, predicted_token_class)

ruaccent/char_tokenizer.py

@@ -0,0 +1,112 @@
+import os
+from typing import Optional, Tuple, List
+from collections import OrderedDict
+
+from transformers import PreTrainedTokenizer
+
+
+def load_vocab(vocab_file):
+    vocab = OrderedDict()
+    with open(vocab_file, "r", encoding="utf-8") as reader:
+        tokens = reader.readlines()
+    for index, token in enumerate(tokens):
+        token = token.rstrip("\n")
+        vocab[token] = index
+    return vocab
+
+
+class CharTokenizer(PreTrainedTokenizer):
+    vocab_files_names = {"vocab_file": "vocab.txt"}
+
+    def __init__(
+        self,
+        vocab_file=None,
+        pad_token="[pad]",
+        unk_token="[unk]",
+        bos_token="[bos]",
+        eos_token="[eos]",
+        do_lower_case=False,
+        *args,
+        **kwargs
+    ):
+        super().__init__(
+            pad_token=pad_token,
+            unk_token=unk_token,
+            bos_token=bos_token,
+            eos_token=eos_token,
+            do_lower_case=do_lower_case,
+            **kwargs
+        )
+        self.do_lower_case = do_lower_case
+
+        if not vocab_file or not os.path.isfile(vocab_file):
+            self.vocab = OrderedDict()
+            self.ids_to_tokens = OrderedDict()
+        else:
+            self.vocab = load_vocab(vocab_file)
+            self.ids_to_tokens = OrderedDict([(ids, tok) for tok, ids in self.vocab.items()])
+
+    @property
+    def vocab_size(self):
+        return len(self.vocab)
+
+    def get_vocab(self):
+        return self.vocab
+
+    def _convert_token_to_id(self, token):
+        if self.do_lower_case:
+            token = token.lower()
+        return self.vocab.get(token, self.vocab[self.unk_token])
+
+    def _convert_id_to_token(self, index):
+        return self.ids_to_tokens[index]
+
+    def _tokenize(self, text):
+        if self.do_lower_case:
+            text = text.lower()
+        return list(text)
+
+    def convert_tokens_to_string(self, tokens):
+        return "".join(tokens)
+
+    def build_inputs_with_special_tokens(
+        self,
+        token_ids_0: List[int],
+        token_ids_1: Optional[List[int]] = None
+    ) -> List[int]:
+        bos = [self.bos_token_id]
+        eos = [self.eos_token_id]
+        return bos + token_ids_0 + eos
+
+    def get_special_tokens_mask(
+         self,
+         token_ids_0: List[int],
+         token_ids_1: Optional[List[int]] = None
+    ) -> List[int]:
+        return [1] + ([0] * len(token_ids_0)) + [1]
+
+    def create_token_type_ids_from_sequences(
+        self,
+        token_ids_0: List[int],
+        token_ids_1: Optional[List[int]] = None
+    ) -> List[int]:
+        return (len(token_ids_0) + 2) * [0]
+
+    def save_vocabulary(
+        self,
+        save_directory: str,
+        filename_prefix: Optional[str] = None
+    ) -> Tuple[str]:
+        assert os.path.isdir(save_directory)
+        vocab_file = os.path.join(
+            save_directory,
+            (filename_prefix + "-" if filename_prefix else "") +
+            self.vocab_files_names["vocab_file"]
+        )
+        index = 0
+        with open(vocab_file, "w", encoding="utf-8") as writer:
+            for token, token_index in sorted(self.vocab.items(), key=lambda kv: kv[1]):
+                assert index == token_index
+                writer.write(token + "\n")
+                index += 1
+        return (vocab_file,)

ruaccent/omograph_model.py

@@ -0,0 +1,21 @@
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+import torch
+
+class OmographModel:
+    def __init__(self) -> None:
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        
+    def load(self, path):
+        self.nli_model = AutoModelForSequenceClassification.from_pretrained(path, torch_dtype=torch.bfloat16).to(self.device)
+        self.tokenizer = AutoTokenizer.from_pretrained(path)
+        
+    def classify(self, text, hypotheses):
+        encodings = self.tokenizer.batch_encode_plus([(text, hyp) for hyp in hypotheses], return_tensors='pt', padding=True)
+        input_ids = encodings['input_ids'].to(self.device)
+        with torch.no_grad():
+            logits = self.nli_model(input_ids)[0]
+            entail_contradiction_logits = logits[:,[0,2]]
+            probs = entail_contradiction_logits.softmax(dim=1)
+            prob_label_is_true = [float(p[1]) for p in probs]
+
+        return hypotheses[prob_label_is_true.index(max(prob_label_is_true))]

ruaccent/ruaccent.py

@@ -0,0 +1,153 @@
+import json
+import pathlib
+from huggingface_hub import snapshot_download
+import os
+from os.path import join as join_path
+from .omograph_model import OmographModel
+from .accent_model import AccentModel
+import re
+
+
+class RUAccent:
+    def __init__(self, workdir=None):
+        self.omograph_model = OmographModel()
+        self.accent_model = AccentModel()
+        if not workdir:
+            self.workdir = str(pathlib.Path(__file__).resolve().parent)
+        else:
+            self.workdir = workdir
+
+    def load(
+        self,
+        omograph_model_size="medium",
+        dict_load_startup=False,
+        disable_accent_dict=False,
+        repo="TeraTTS/accentuator",
+    ):
+        if not os.path.exists(
+            join_path(self.workdir, "dictionary")
+        ) or not os.path.exists(join_path(self.workdir, "nn")):
+            snapshot_download(
+                repo_id=repo,
+                ignore_patterns=["*.md", "*.gitattributes"],
+                local_dir=self.workdir,
+                local_dir_use_symlinks=False,
+            )
+        self.omographs = json.load(
+            open(join_path(self.workdir, "dictionary/omographs.json"), encoding='utf-8')
+        )
+        self.yo_words = json.load(
+            open(join_path(self.workdir, "dictionary/yo_words.json"), encoding='utf-8')
+        )
+        self.dict_load_startup = dict_load_startup
+
+        if dict_load_startup:
+            self.accents = json.load(
+                open(join_path(self.workdir, "dictionary/accents.json"), encoding='utf-8')
+            )
+        if disable_accent_dict:
+            self.accents = {}
+            self.disable_accent_dict = True
+        else:
+            self.disable_accent_dict = False
+
+        if omograph_model_size not in ["small", "medium"]:
+            raise NotImplementedError
+
+        self.omograph_model.load(
+            join_path(self.workdir, f"nn/nn_omograph/{omograph_model_size}/")
+        )
+        self.accent_model.load(join_path(self.workdir, "nn/nn_accent/"))
+
+
+    def split_by_words(self, string):
+        result = re.findall(r"\w*(?:\+\w+)*|[^\w\s]+", string.lower())
+        return [res for res in result if res]
+
+    def extract_initial_letters(self, text):
+        words = text
+        initial_letters = []
+        for word in words:
+            if len(word) > 2 and '+' not in word and not bool(re.search('[a-zA-Z]', word)):
+                initial_letters.append(word[0])
+        return initial_letters
+
+    def load_dict(self, text):
+        chars = self.extract_initial_letters(text)
+        out_dict = {}
+        for char in chars:
+            out_dict.update(
+                json.load(
+                    open(
+                        join_path(self.workdir, f"dictionary/letter_accent/{char}.json"),
+                        encoding='utf-8'
+                    )
+                )
+            )
+        return out_dict
+
+    def count_vowels(self, text):
+        vowels = "аеёиоуыэюяАЕЁИОУЫЭЮЯ"
+        return sum(1 for char in text if char in vowels)
+
+    def has_punctuation(self, text):
+        for char in text:
+            if char in "!\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~":
+                return True
+        return False
+
+    def delete_spaces_before_punc(self, text):
+        punc = "!\"#$%&'()*,./:;<=>?@[\\]^_`{|}~"
+        for char in punc:
+            text = text.replace(" " + char, char)
+        return text
+
+    def process_yo(self, text):
+        splitted_text = text
+
+        for i, word in enumerate(splitted_text):
+            splitted_text[i] = self.yo_words.get(word, word)
+        return splitted_text
+
+    def process_omographs(self, text):
+        splitted_text = text
+
+        founded_omographs = []
+        for i, word in enumerate(splitted_text):
+            variants = self.omographs.get(word)
+            if variants:
+                founded_omographs.append(
+                    {"word": word, "variants": variants, "position": i}
+                )
+        for omograph in founded_omographs:
+            splitted_text[
+                omograph["position"]
+            ] = f"<w>{splitted_text[omograph['position']]}</w>"
+            cls = self.omograph_model.classify(
+                " ".join(splitted_text), omograph["variants"]
+            )
+            splitted_text[omograph["position"]] = cls
+        return splitted_text
+
+    def process_accent(self, text):
+        if not self.dict_load_startup and not self.disable_accent_dict:
+            self.accents = self.load_dict(text)
+
+        splitted_text = text
+
+        for i, word in enumerate(splitted_text):
+            stressed_word = self.accents.get(word, word)
+            if stressed_word == word and not self.has_punctuation(word) and self.count_vowels(word) > 1:
+                splitted_text[i] = self.accent_model.put_accent(word)
+            else:
+                splitted_text[i] = stressed_word
+        return splitted_text
+
+    def process_all(self, text):
+        text = self.split_by_words(text)
+        processed_text = self.process_yo(text)
+        processed_text = self.process_omographs(processed_text)
+        processed_text = self.process_accent(processed_text)
+        processed_text = " ".join(processed_text)
+        processed_text = self.delete_spaces_before_punc(processed_text)
+        return processed_text

tokenizer/__init__.py

@@ -0,0 +1,2 @@
+from .gruut import Tokenizer as TokenizerGRUUT
+from .g2p import Tokenizer as TokenizerG2P

tokenizer/g2p/__init__.py

@@ -0,0 +1 @@
+from .tokenizer import Tokenizer

tokenizer/g2p/g2p.py

@@ -0,0 +1,94 @@
+
+softletters=set(u"яёюиье")
+startsyl=set(u"#ъьаяоёуюэеиы-")
+others = set(["#", "+", "-", u"ь", u"ъ"])
+
+softhard_cons = {
+    u"б" : u"b",
+    u"в" : u"v",
+    u"г" : u"g",
+    u"Г" : u"g",
+    u"д" : u"d",
+    u"з" : u"z",
+    u"к" : u"k",
+    u"л" : u"l",
+    u"м" : u"m",
+    u"н" : u"n",
+    u"п" : u"p",
+    u"р" : u"r",
+    u"с" : u"s",
+    u"т" : u"t",
+    u"ф" : u"f",
+    u"х" : u"h"
+}
+
+other_cons = {
+    u"ж" : u"zh",
+    u"ц" : u"c",
+    u"ч" : u"ch",
+    u"ш" : u"sh",
+    u"щ" : u"sch",
+    u"й" : u"j"
+}
+
+vowels = {
+    u"а" : u"a",
+    u"я" : u"a",
+    u"у" : u"u",
+    u"ю" : u"u",
+    u"о" : u"o",
+    u"ё" : u"o",
+    u"э" : u"e",
+    u"е" : u"e",
+    u"и" : u"i",
+    u"ы" : u"y",
+}                                
+
+def pallatize(phones):
+    for i, phone in enumerate(phones[:-1]):
+        if phone[0] in softhard_cons:
+            if phones[i+1][0] in softletters:
+                phones[i] = (softhard_cons[phone[0]] + "j", 0)
+            else:
+                phones[i] = (softhard_cons[phone[0]], 0)
+        if phone[0] in other_cons:
+            phones[i] = (other_cons[phone[0]], 0)
+
+def convert_vowels(phones):
+    new_phones = []
+    prev = ""
+    for phone in phones:
+        if prev in startsyl:
+            if phone[0] in set(u"яюеё"):
+                new_phones.append("j")
+        if phone[0] in vowels:
+            new_phones.append(vowels[phone[0]] + str(phone[1]))
+        else:
+            new_phones.append(phone[0])
+        prev = phone[0]
+
+    return new_phones
+
+def convert(stressword):
+    phones = ("#" + stressword + "#")
+
+
+    # Assign stress marks
+    stress_phones = []
+    stress = 0
+    for phone in phones:
+        if phone == "+":
+            stress = 1
+        else:
+            stress_phones.append((phone, stress))
+            stress = 0
+
+    # Pallatize
+    pallatize(stress_phones)
+
+    # Assign stress
+    phones = convert_vowels(stress_phones)
+
+    # Filter
+    phones = [x for x in phones if x not in others]
+    return " ".join(phones)

tokenizer/g2p/tokenizer.py

@@ -0,0 +1,47 @@
+import re
+from .g2p import *
+import json
+import os
+
+class Tokenizer():
+    def __init__(self, data_path: str) -> None:
+        self.dic = {}
+        for line in open(os.path.join(data_path, "dictionary.txt")):
+           items = line.split()
+           self.dic[items[0]] = " ".join(items[1:])
+
+        self.config = json.load(open(os.path.join(data_path, "config.json")))
+    
+    def g2p(self, text):
+        text = re.sub("—", "-", text)
+        text = re.sub("([!'(),-.:;?])", r' \1 ', text)
+
+        phonemes = []
+        for word in text.split():
+            if re.match("[!'(),-.:;?]", word):
+                phonemes.append(word)
+                continue
+
+            word = word.lower()
+            if len(phonemes) > 0: phonemes.append(' ')
+
+            if word in self.dic:
+                phonemes.extend(self.dic[word].split())
+            else:
+                phonemes.extend(convert(word).split())
+
+        phoneme_id_map = self.config["phoneme_id_map"]
+        phoneme_ids = []
+        phoneme_ids.extend(phoneme_id_map["^"])
+        phoneme_ids.extend(phoneme_id_map["_"])
+        for p in phonemes:
+            if p in phoneme_id_map:
+                phoneme_ids.extend(phoneme_id_map[p])
+                phoneme_ids.extend(phoneme_id_map["_"])
+        phoneme_ids.extend(phoneme_id_map["$"])
+
+        return phoneme_ids, phonemes
+    
+    def _get_seq(self, text: str) -> list[int]:
+        seq = self.g2p(text)[0]
+        return seq

tokenizer/gruut/__init__.py

@@ -0,0 +1 @@
+from .tokenizer import Tokenizer

tokenizer/gruut/tokenizer.py

@@ -0,0 +1,36 @@
+from gruut import sentences
+import os
+import re
+
+class Tokenizer():
+    def __init__(self, path) -> None:
+        with open(os.path.join(path, "vocab.txt"), "r", encoding="utf-8") as vocab_file:
+            self.symbols = vocab_file.read().split("\n")
+            self.symbols = list(map(chr, list(map(int, self.symbols))))
+            
+        self.symbol_to_id = {s: i for i, s in enumerate(self.symbols)}
+        
+    def _ru_phonems(self, text: str) -> str:
+        text = text.lower()
+        phonemes = ""
+        for sent in sentences(text, lang="ru"):
+            for word in sent:
+                if word.phonemes:
+                    phonemes += "".join(word.phonemes)
+        phonemes = re.sub(re.compile(r'\s+'), ' ', phonemes).lstrip().rstrip()
+        return phonemes
+    
+    
+    def _text_to_sequence(self, text: str) -> list[int]:
+        '''convert text to seq'''
+        sequence = []
+        clean_text = self._ru_phonems(text)
+        for symbol in clean_text:
+            symbol_id = self.symbol_to_id[symbol]
+            sequence += [symbol_id]
+        return sequence
+    
+    
+    def _get_seq(self, text: str) -> list[int]:
+        seq = self._text_to_sequence(text)
+        return seq