Spaces:
Build error
Build error
import os | |
from trainer import Trainer, TrainerArgs | |
from TTS.config.shared_configs import BaseAudioConfig | |
from TTS.tts.configs.shared_configs import BaseDatasetConfig, CapacitronVAEConfig | |
from TTS.tts.configs.tacotron_config import TacotronConfig | |
from TTS.tts.datasets import load_tts_samples | |
from TTS.tts.models.tacotron import Tacotron | |
from TTS.tts.utils.text.tokenizer import TTSTokenizer | |
from TTS.utils.audio import AudioProcessor | |
output_path = os.path.dirname(os.path.abspath(__file__)) | |
data_path = "/srv/data/" | |
# Using LJSpeech like dataset processing for the blizzard dataset | |
dataset_config = BaseDatasetConfig(formatter="ljspeech", meta_file_train="metadata.csv", path=data_path) | |
audio_config = BaseAudioConfig( | |
sample_rate=24000, | |
do_trim_silence=True, | |
trim_db=60.0, | |
signal_norm=True, | |
mel_fmin=80.0, | |
mel_fmax=12000, | |
spec_gain=20.0, | |
log_func="np.log10", | |
ref_level_db=20, | |
preemphasis=0.0, | |
min_level_db=-100, | |
) | |
# Using the standard Capacitron config | |
capacitron_config = CapacitronVAEConfig(capacitron_VAE_loss_alpha=1.0) | |
config = TacotronConfig( | |
run_name="Blizzard-Capacitron-T1", | |
audio=audio_config, | |
capacitron_vae=capacitron_config, | |
use_capacitron_vae=True, | |
batch_size=128, # Tune this to your gpu | |
max_audio_len=6 * 24000, # Tune this to your gpu | |
min_audio_len=0.5 * 24000, | |
eval_batch_size=16, | |
num_loader_workers=12, | |
num_eval_loader_workers=8, | |
precompute_num_workers=24, | |
run_eval=True, | |
test_delay_epochs=5, | |
r=2, | |
optimizer="CapacitronOptimizer", | |
optimizer_params={"RAdam": {"betas": [0.9, 0.998], "weight_decay": 1e-6}, "SGD": {"lr": 1e-5, "momentum": 0.9}}, | |
attention_type="graves", | |
attention_heads=5, | |
epochs=1000, | |
text_cleaner="phoneme_cleaners", | |
use_phonemes=True, | |
phoneme_language="en-us", | |
phonemizer="espeak", | |
phoneme_cache_path=os.path.join(data_path, "phoneme_cache"), | |
stopnet_pos_weight=15, | |
print_step=50, | |
print_eval=True, | |
mixed_precision=False, | |
output_path=output_path, | |
datasets=[dataset_config], | |
lr=1e-3, | |
lr_scheduler="StepwiseGradualLR", | |
lr_scheduler_params={"gradual_learning_rates": [[0, 1e-3], [2e4, 5e-4], [4e4, 3e-4], [6e4, 1e-4], [8e4, 5e-5]]}, | |
scheduler_after_epoch=False, # scheduler doesn't work without this flag | |
loss_masking=False, | |
decoder_loss_alpha=1.0, | |
postnet_loss_alpha=1.0, | |
postnet_diff_spec_alpha=1.0, | |
decoder_diff_spec_alpha=1.0, | |
decoder_ssim_alpha=1.0, | |
postnet_ssim_alpha=1.0, | |
) | |
ap = AudioProcessor(**config.audio.to_dict()) | |
tokenizer, config = TTSTokenizer.init_from_config(config) | |
train_samples, eval_samples = load_tts_samples(dataset_config, eval_split=True) | |
model = Tacotron(config, ap, tokenizer, speaker_manager=None) | |
trainer = Trainer( | |
TrainerArgs(), | |
config, | |
output_path, | |
model=model, | |
train_samples=train_samples, | |
eval_samples=eval_samples, | |
) | |
# 🚀 | |
trainer.fit() | |