Spear-TTS PyTorch

You can use SemanticToTextDatasetGenerator to create a new dataset by running a model over an existing dataset and saving the results to a folder. Once generated, you can load this new dataset using GeneratedAudioTextDataset.

Workflow:

1. Initialize SemanticToTextDatasetGenerator with your TextToSemantic model, an existing dataset, and a target folder.
2. Call the generator with a max_length parameter.
3. Load the results using GeneratedAudioTextDataset.

from spear_tts_pytorch import (
    TextToSemantic,
    SemanticToTextDatasetGenerator,
    GeneratedAudioTextDataset,
    MockDataset
)

# Assuming 'model' is an initialized TextToSemantic instance
ds = MockDataset(10)

dataset_generator = SemanticToTextDatasetGenerator(
    model = model,
    dataset = ds,
    folder = './output_folder'
)

# Generate the dataset
dataset_generator(max_length = 2)

# Load the generated dataset
generated_dataset = GeneratedAudioTextDataset(
    folder = './output_folder'
)

assert len(generated_dataset) == 10

The TextToSemantic model is a multi-speaker text-to-speech attention network. It requires a wav2vec component (typically from audiolm_pytorch) for semantic processing.

Key parameters for TextToSemantic:

• wav2vec: A HubertWithKmeans instance.
• dim: Dimension of the model.
• num_text_token_ids: Number of text token IDs.
• heads: Number of attention heads.
• target_kv_heads: Number of target KV heads (used for grouped query attention to enable memory-efficient decoding).
• source_depth: Depth of the source.
• target_depth: Depth of the target.

import torch
from audiolm_pytorch import HubertWithKmeans
from spear_tts_pytorch import TextToSemantic

wav2vec = HubertWithKmeans(
    checkpoint_path = './hubert_base_ls960.pt',
    kmeans_path = './hubert_base_ls960_L9_km500.bin'
)

model = TextToSemantic(
    wav2vec = wav2vec,
    dim = 512,
    num_text_token_ids = 256,
    heads = 8,
    target_kv_heads = 2, # grouped query attention, for memory efficient decoding
    source_depth = 1,
    target_depth = 1
)