|
import os |
|
|
|
from trainer import Trainer, TrainerArgs |
|
|
|
from TTS.tts.configs.shared_configs import BaseDatasetConfig, CharactersConfig |
|
from TTS.tts.configs.vits_config import VitsConfig |
|
from TTS.tts.datasets import load_tts_samples |
|
from TTS.tts.models.vits import Vits, VitsArgs, VitsAudioConfig |
|
from TTS.tts.utils.speakers import SpeakerManager |
|
from TTS.tts.utils.text.tokenizer import TTSTokenizer |
|
from TTS.utils.audio import AudioProcessor |
|
|
|
output_path = os.path.dirname(os.path.abspath(__file__)) |
|
|
|
|
|
|
|
|
|
CONTINUE_PATH=None |
|
RESTORE_PATH=None |
|
START_WITH_EVAL=True |
|
GRAD_ACUMM_STEPS=1 |
|
|
|
meta_file = '/home/ubuntu/nctb-cropped/metadata.txt' |
|
root_path = '/home/ubuntu/nctb-cropped/' |
|
|
|
def formatter(root_path, meta_file, **kwargs): |
|
"""Normalizes the LJSpeech meta data file to TTS format |
|
https://keithito.com/LJ-Speech-Dataset/""" |
|
txt_file = meta_file |
|
items = [] |
|
with open(txt_file, "r", encoding="utf-8") as ttf: |
|
for line in ttf: |
|
cols = line.split("|") |
|
wav_file = os.path.join(root_path,'audio', cols[0]) |
|
speaker_name = cols[0].split('_')[-1].split('.')[0] |
|
try: |
|
text = cols[1] |
|
except: |
|
print("not found") |
|
|
|
items.append({"text": text, "audio_file": wav_file, "speaker_name": speaker_name, "root_path": root_path}) |
|
return items |
|
|
|
|
|
dataset_config = BaseDatasetConfig( |
|
meta_file_train=meta_file, path=os.path.join(root_path, ""), language="bn" |
|
) |
|
|
|
|
|
characters_config = CharactersConfig( |
|
pad = '<PAD>', |
|
eos = '<EOS>', |
|
bos = '<BOS>', |
|
blank = '<BLNK>', |
|
phonemes = None, |
|
characters = "abcdefghijklmnopqrstuvwxyz0123456789+=/*√তট৫ভিঐঋখঊড়ইজমএেঘঙসীঢ়হঞ‘ঈকণ৬ঁৗশঢঠ\u200c১্২৮দৃঔগও—ছউংবৈঝাযফ\u200dচরষঅৌৎথড়৪ধ০ুূ৩আঃপয়’'”নলো_…ৰ", |
|
punctuations = "-–:;!,|.?॥। “", |
|
) |
|
|
|
|
|
audio_config = VitsAudioConfig( |
|
sample_rate=16000, win_length=1024, hop_length=256, num_mels=80, mel_fmin=0, mel_fmax=None |
|
) |
|
|
|
vitsArgs = VitsArgs( |
|
use_speaker_embedding=True, |
|
) |
|
|
|
config = VitsConfig( |
|
model_args=vitsArgs, |
|
audio=audio_config, |
|
run_name="vits_nctb", |
|
batch_size=64, |
|
eval_batch_size=8, |
|
batch_group_size=5, |
|
num_loader_workers=4, |
|
num_eval_loader_workers=4, |
|
run_eval=True, |
|
test_delay_epochs=-1, |
|
epochs=1000, |
|
|
|
text_cleaner='phoneme_cleaners', |
|
use_phonemes=True, |
|
phoneme_language="bn", |
|
phoneme_cache_path=os.path.join(output_path, "phoneme_cache"), |
|
compute_input_seq_cache=True, |
|
print_step=25, |
|
print_eval=False, |
|
mixed_precision=True, |
|
max_text_len=325, |
|
output_path=output_path, |
|
datasets=[dataset_config], |
|
characters=characters_config, |
|
save_step=5000, |
|
cudnn_benchmark=True, |
|
test_sentences = [ |
|
'আমার সোনার বাংলা, আমি তোমায় ভালোবাসি।', |
|
'চিরদিন তোমার আকাশ, তোমার বাতাস, আমার প্রাণে বাজায় বাঁশি', |
|
'ও মা, ফাগুনে তোর আমের বনে ঘ্রাণে পাগল করে,মরি হায়, হায় রে।' |
|
] |
|
) |
|
|
|
|
|
|
|
|
|
ap = AudioProcessor.init_from_config(config) |
|
|
|
|
|
|
|
|
|
tokenizer, config = TTSTokenizer.init_from_config(config) |
|
|
|
|
|
|
|
|
|
|
|
|
|
train_samples, eval_samples = load_tts_samples( |
|
dataset_config, |
|
formatter=formatter, |
|
eval_split=True, |
|
eval_split_max_size=config.eval_split_max_size, |
|
eval_split_size=config.eval_split_size, |
|
) |
|
|
|
|
|
|
|
speaker_manager = SpeakerManager() |
|
speaker_manager.set_ids_from_data(train_samples + eval_samples, parse_key="speaker_name") |
|
config.model_args.num_speakers = speaker_manager.num_speakers |
|
|
|
|
|
model = Vits(config, ap, tokenizer, speaker_manager) |
|
|
|
|
|
trainer = Trainer( |
|
TrainerArgs( |
|
continue_path=CONTINUE_PATH, |
|
restore_path=RESTORE_PATH, |
|
skip_train_epoch=False, |
|
start_with_eval=START_WITH_EVAL, |
|
grad_accum_steps=GRAD_ACUMM_STEPS, |
|
), |
|
config, |
|
output_path, |
|
model=model, |
|
train_samples=train_samples, |
|
eval_samples=eval_samples, |
|
) |
|
trainer.fit() |