soundstorm-pytorch

If you already have pre-encoded codebook IDs (e.g., from a SoundStream model), you can train the SoundStorm model directly on these codes.

1. Initialize a ConformerWrapper with the appropriate codebook_size and num_quantizers.
2. Initialize SoundStorm with the conformer and specify the number of steps (e.g., 18) and a schedule (e.g., 'cosine').
3. Pass the codes of shape (batch, seq, num_residual_vq) to the model to compute loss.

import torch
from soundstorm_pytorch import SoundStorm, ConformerWrapper

conformer = ConformerWrapper(
    codebook_size = 1024,
    num_quantizers = 12,
    conformer = dict(
        dim = 512,
        depth = 2
    ),
)

model = SoundStorm(
    conformer,
    steps = 18,          # 18 steps, as in original maskgit paper
    schedule = 'cosine'  # currently the best schedule is cosine
)

# codes shape: (batch, seq, num residual VQ)
codes = torch.randint(0, 1024, (2, 1024, 12))

loss, _ = model(codes)
loss.backward()

# Generation
generated = model.generate(1024, batch_size = 2) # (2, 1024)

To train directly on raw audio, pass a pretrained SoundStream instance into the SoundStorm constructor. This allows the model to handle the audio waveform directly.

1. Initialize ConformerWrapper.
2. Initialize SoundStream (from audiolm-pytorch) with matching codebook_size and rq_num_quantizers.
3. Initialize SoundStorm passing the soundstream instance.
4. Pass raw audio tensors to the model for training.
5. Use model.generate(seconds = ...) to generate audio based on a duration in seconds.

import torch
from soundstorm_pytorch import SoundStorm, ConformerWrapper, Conformer, SoundStream

conformer = ConformerWrapper(
    codebook_size = 1024,
    num_quantizers = 12,
    conformer = dict(
        dim = 512,
        depth = 2
    ),
)

soundstream = SoundStream(
    codebook_size = 1024,
    rq_num_quantizers = 12,
    attn_window_size = 128,
    attn_depth = 2
)

model = SoundStorm(
    conformer,
    soundstream = soundstream   # pass in the soundstream
)

# audio shape: (batch, samples)
audio = torch.randn(2, 10080)

loss, _ = model(audio)
loss.backward()

# Generate 30 seconds of audio
generated_audio = model.generate(seconds = 30, batch_size = 2)

For complete text-to-speech (TTS) capabilities, you can integrate a trained TextToSemantic encoder/decoder (from spear-tts-pytorch) into SoundStorm using the spear_tts_text_to_semantic argument.

1. Load your trained TextToSemantic model weights.
2. Initialize SoundStorm with the conformer, soundstream, and the spear_tts_text_to_semantic model.
3. Use model.generate(texts = [...]) to generate raw waveforms from text strings.

from spear_tts_pytorch import TextToSemantic

text_to_semantic = TextToSemantic(
    dim = 512,
    source_depth = 12,
    target_depth = 12,
    num_text_token_ids = 50000,
    num_semantic_token_ids = 20000,
    use_openai_tokenizer = True
)

# load the trained text-to-semantic transformer
text_to_semantic.load('/path/to/trained/model.pt')

# pass it into the soundstorm
model = SoundStorm(
    conformer,
    soundstream = soundstream,
    spear_tts_text_to_semantic = text_to_semantic
).cuda()

# generate raw waveform decoded from soundstream
generated_speech = model.generate(
    texts = [
        'the rain in spain stays mainly in the plain',
        'the quick brown fox jumps over the lazy dog'
    ]
) # (2, n)