Back to Repositories

Testing VITS Speaker Embedding Training and Inference in Coqui-AI TTS

This test suite validates the training and inference functionality of VITS text-to-speech model with speaker embeddings in the Coqui-AI TTS framework. It covers model configuration, training initialization, checkpoint management, and inference capabilities with speaker identification.

Test Coverage Overview

The test suite provides comprehensive coverage of VITS speaker embedding functionality:

Model configuration validation with speaker embedding parameters
Initial training cycle verification
Checkpoint saving and loading
Inference with speaker identification
Training continuation from saved checkpoints

Implementation Analysis

The testing approach implements a systematic validation workflow:

Uses VitsConfig for model configuration with speaker embedding settings
Executes training through CLI commands with CUDA device management
Validates checkpoint integrity and config persistence
Tests inference pipeline with speaker identification

Technical Details

Key technical components include:

CUDA GPU acceleration support
JSON configuration management
LJSpeech dataset integration
Phoneme cache implementation
Speaker embedding dimension: 256
Checkpoint management utilities

Best Practices Demonstrated

The test implementation showcases several testing best practices:

Proper resource cleanup after test execution
Systematic configuration verification
Explicit CUDA device management
Comprehensive model lifecycle testing
Integration of multiple testing aspects in a single suite

coqui-ai/tts

tests/tts_tests/test_vits_speaker_emb_train.py

            
import glob
import json
import os
import shutil

from trainer import get_last_checkpoint

from tests import get_device_id, get_tests_output_path, run_cli
from TTS.tts.configs.vits_config import VitsConfig

config_path = os.path.join(get_tests_output_path(), "test_model_config.json")
output_path = os.path.join(get_tests_output_path(), "train_outputs")


config = VitsConfig(
    batch_size=2,
    eval_batch_size=2,
    num_loader_workers=0,
    num_eval_loader_workers=0,
    text_cleaner="english_cleaners",
    use_phonemes=True,
    phoneme_language="en-us",
    phoneme_cache_path="tests/data/ljspeech/phoneme_cache/",
    run_eval=True,
    test_delay_epochs=-1,
    epochs=1,
    print_step=1,
    print_eval=True,
    test_sentences=[
        ["Be a voice, not an echo.", "ljspeech-1"],
    ],
)
# set audio config
config.audio.do_trim_silence = True
config.audio.trim_db = 60

# active multispeaker d-vec mode
config.model_args.use_speaker_embedding = True
config.model_args.use_d_vector_file = False
config.model_args.d_vector_file = None
config.model_args.d_vector_dim = 256


config.save_json(config_path)

# train the model for one epoch
command_train = (
    f"CUDA_VISIBLE_DEVICES='{get_device_id()}' python TTS/bin/train_tts.py --config_path {config_path} "
    f"--coqpit.output_path {output_path} "
    "--coqpit.datasets.0.formatter ljspeech_test "
    "--coqpit.datasets.0.meta_file_train metadata.csv "
    "--coqpit.datasets.0.meta_file_val metadata.csv "
    "--coqpit.datasets.0.path tests/data/ljspeech "
    "--coqpit.datasets.0.meta_file_attn_mask tests/data/ljspeech/metadata_attn_mask.txt "
    "--coqpit.test_delay_epochs 0"
)
run_cli(command_train)

# Find latest folder
continue_path = max(glob.glob(os.path.join(output_path, "*/")), key=os.path.getmtime)

# Inference using TTS API
continue_config_path = os.path.join(continue_path, "config.json")
continue_restore_path, _ = get_last_checkpoint(continue_path)
out_wav_path = os.path.join(get_tests_output_path(), "output.wav")
speaker_id = "ljspeech-1"
continue_speakers_path = os.path.join(continue_path, "speakers.json")

# Check integrity of the config
with open(continue_config_path, "r", encoding="utf-8") as f:
    config_loaded = json.load(f)
assert config_loaded["characters"] is not None
assert config_loaded["output_path"] in continue_path
assert config_loaded["test_delay_epochs"] == 0

# Load the model and run inference
inference_command = f"CUDA_VISIBLE_DEVICES='{get_device_id()}' tts --text 'This is an example.' --speaker_idx {speaker_id} --speakers_file_path {continue_speakers_path} --config_path {continue_config_path} --model_path {continue_restore_path} --out_path {out_wav_path}"
run_cli(inference_command)

# restore the model and continue training for one more epoch
command_train = f"CUDA_VISIBLE_DEVICES='{get_device_id()}' python TTS/bin/train_tts.py --continue_path {continue_path} "
run_cli(command_train)
shutil.rmtree(continue_path)