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.tacotron2_config import Tacotron2Config from TTS.tts.datasets import load_tts_samples from TTS.tts.models.tacotron2 import Tacotron2 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=22050, do_trim_silence=True, trim_db=60.0, signal_norm=False, mel_fmin=0.0, mel_fmax=11025, spec_gain=1.0, log_func="np.log", ref_level_db=20, preemphasis=0.0, ) # Using the standard Capacitron config capacitron_config = CapacitronVAEConfig(capacitron_VAE_loss_alpha=1.0, capacitron_capacity=50) config = Tacotron2Config( run_name="Capacitron-Tacotron2", audio=audio_config, capacitron_vae=capacitron_config, use_capacitron_vae=True, batch_size=128, # Tune this to your gpu max_audio_len=8 * 22050, # Tune this to your gpu min_audio_len=1 * 22050, eval_batch_size=16, num_loader_workers=8, num_eval_loader_workers=8, precompute_num_workers=24, run_eval=True, test_delay_epochs=25, ga_alpha=0.0, 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="dynamic_convolution", grad_clip=0.0, # Important! We overwrite the standard grad_clip with capacitron_grad_clip double_decoder_consistency=False, 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=25, print_eval=True, mixed_precision=False, seq_len_norm=True, 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], [4e5, 3e-4], [6e4, 1e-4], [8e4, 5e-5], ] }, scheduler_after_epoch=False, # scheduler doesn't work without this flag # Need to experiment with these below for capacitron loss_masking=False, decoder_loss_alpha=1.0, postnet_loss_alpha=1.0, postnet_diff_spec_alpha=0.0, decoder_diff_spec_alpha=0.0, decoder_ssim_alpha=0.0, postnet_ssim_alpha=0.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 = Tacotron2(config, ap, tokenizer, speaker_manager=None) trainer = Trainer( TrainerArgs(), config, output_path, model=model, train_samples=train_samples, eval_samples=eval_samples, training_assets={"audio_processor": ap}, ) trainer.fit()