audiolm-pytorch

To perform text-conditioned audio synthesis (similar to VALL-E), use the SemanticTransformer and SemanticTransformerTrainer.

1. Initialize HubertWithKmeans: Load a pre-trained Hubert model and its corresponding K-means quantization file.
2. Configure SemanticTransformer: Set has_condition = True. You can use cond_as_self_attn_prefix = True to condition via self-attention prefix instead of cross-attention.
3. Prepare Dataset: Extend torch.utils.data.Dataset to return a tuple containing a string (the text caption) and an audio tensor. The framework automatically routes these to the transformer.
4. Train: Use SemanticTransformerTrainer to manage the training loop.
5. Generate: Use trainer.generate(text=[...]) to synthesize audio from text prompts.

import torch
from audiolm_pytorch import HubertWithKmeans, SemanticTransformer, SemanticTransformerTrainer
from torch.utils.data import Dataset

# 1. Setup Wav2Vec/Hubert
wav2vec = HubertWithKmeans(
    checkpoint_path = './hubert/hubert_base_ls960.pt',
    kmeans_path = './hubert/hubert_base_ls960_L9_km500.bin'
)

# 2. Setup Semantic Transformer
semantic_transformer = SemanticTransformer(
    num_semantic_tokens = 500,
    dim = 1024,
    depth = 6,
    has_condition = True,
    cond_as_self_attn_prefix = True
).cuda()

# 3. Define Dataset (must return caption and audio)
class MockTextAudioDataset(Dataset):
    def __init__(self, length = 100, audio_length = 320 * 32):
        super().__init__()
        self.audio_length = audio_length
        self.len = length

    def __len__(self):
        return self.len

    def __getitem__(self, idx):
        mock_audio = torch.randn(self.audio_length)
        mock_caption = 'audio caption'
        return mock_caption, mock_audio

dataset = MockTextAudioDataset()

# 4. Train
trainer = SemanticTransformerTrainer(
    transformer = semantic_transformer,
    wav2vec = wav2vec,
    dataset = dataset,
    batch_size = 4,
    grad_accum_every = 8,
    data_max_length = 320 * 32,
    num_train_steps = 1_000_000
)

trainer.train()

# 5. Generate
sample = trainer.generate(text = ['sound of rain drops on the rooftops'], batch_size = 1, max_length = 2)

You can train a SoundStream model using SoundStreamTrainer. Once trained, you can use it for audio autoencoding, tokenization, and decoding.

Key features:

• SoundStream.init_and_load_from(path): Loads a model from a checkpoint without needing to manually re-specify configurations.
• soundstream.tokenize(audio): Converts audio to codebook indices.
• soundstream.decode_from_codebook_indices(codes): Reconstructs audio from indices.
• soundstream(audio, return_recons_only = True): Returns only the reconstructed audio.

from audiolm_pytorch import SoundStream, SoundStreamTrainer

soundstream = SoundStream(
    codebook_size = 4096,
    rq_num_quantizers = 8,
    rq_groups = 2,
    use_lookup_free_quantizer = True,
    use_finite_scalar_quantizer = False,
    attn_window_size = 128,
    attn_depth = 2
)

trainer = SoundStreamTrainer(
    soundstream,
    folder = '/path/to/audio/files',
    batch_size = 4,
    grad_accum_every = 8,
    data_max_length_seconds = 2,
    num_train_steps = 1_000_000
).cuda()

trainer.train()

# Testing autoencoding
soundstream.eval()
audio = torch.randn(10080).cuda()
recons = soundstream(audio, return_recons_only = True)

# Tokenization and decoding
codes = soundstream.tokenize(audio)
recon_audio_from_codes = soundstream.decode_from_codebook_indices(codes)

Train the CoarseTransformer using CoarseTransformerTrainer. This component bridges semantic tokens and acoustic tokens, requiring both a HubertWithKmeans instance and a trained SoundStream codec.

wav2vec = HubertWithKmeans(
    checkpoint_path = f'./{hubert_ckpt}',
    kmeans_path = f'./{hubert_quantizer}'
)

soundstream = SoundStream(
    codebook_size = 1024,
    rq_num_quantizers = 8,
)
soundstream.load(f"./{soundstream_ckpt}")

coarse_transformer = CoarseTransformer(
    num_semantic_tokens = wav2vec.codebook_size,
    codebook_size = 1024,
    num_coarse_quantizers = 3,
    dim = 512,
    depth = 6
)

trainer = CoarseTransformerTrainer(
    transformer = coarse_transformer,
    codec = soundstream,
    wav2vec = wav2vec,
    folder = dataset_folder,
    batch_size = 1,
    data_max_length = 320 * 32,
    save_results_every = 2,
    save_model_every = 4,
    num_train_steps = 9
)

trainer.train()