voicebox-pytorch

Perform text-to-speech training and sampling with SpearTTS integration

main

To use Voicebox for text-guided speech generation, you need to integrate it with SpearTTS components. This involves using HubertWithKmeans for audio processing, TextToSemantic for converting text to semantic tokens, and ConditionalFlowMatcherWrapper to manage the training and sampling process.

Key components:

HubertWithKmeans: Loads a Hubert checkpoint and Kmeans weights.
TextToSemantic: Converts text to semantic tokens using a wav2vec model.
VoiceBox: The core model, which can be configured with an audio_enc_dec (like EncodecVoco()).
ConditionalFlowMatcherWrapper: Wraps the model and semantic converter to provide .sample() and training loss functionality.

import torch
from voicebox_pytorch import (
    VoiceBox,
    EncodecVoco,
    ConditionalFlowMatcherWrapper,
    HubertWithKmeans,
    TextToSemantic
)

# 1. Setup Hubert with Kmeans
wav2vec = HubertWithKmeans(
    checkpoint_path = '/path/to/hubert/checkpoint.pt',
    kmeans_path = '/path/to/hubert/kmeans.bin'
)

# 2. Setup Text to Semantic conversion
text_to_semantic = TextToSemantic(
    wav2vec = wav2vec,
    dim = 512,
    source_depth = 1,
    target_depth = 1,
    use_openai_tokenizer = True
)
text_to_semantic.load('/path/to/trained/spear-tts/model.pt')

# 3. Initialize VoiceBox model
model = VoiceBox(
    dim = 512,
    audio_enc_dec = EncodecVoco(),
    num_cond_tokens = 500,
    depth = 2,
    dim_head = 64,
    heads = 16
)

# 4. Wrap with ConditionalFlowMatcher
cfm_wrapper = ConditionalFlowMatcherWrapper(
    voicebox = model,
    text_to_semantic = text_to_semantic
)

# --- Training ---
audio = torch.randn(2, 12000)
loss = cfm_wrapper(audio)
loss.backward()

# --- Sampling ---
texts = [
    'the rain in spain falls mainly in the plains',
    'she sells sea shells by the seashore'
]
cond = torch.randn(2, 12000)
sampled = cfm_wrapper.sample(cond = cond, texts = texts) # (2, 1, <audio length>)

Install voicebox-pytorch
main
Install the package using pip:
```
$ pip install voicebox-pytorch
```
https://github.com/lucidrains/voicebox-pytorch/blob/main/README.md

Perform unconditional training and sampling

main

For unconditional training (where the model is not conditioned on text), set condition_on_text = False when initializing the VoiceBox model. In this mode, the ConditionalFlowMatcherWrapper does not require a text_to_semantic module.

import torch
from voicebox_pytorch import (
    VoiceBox,
    ConditionalFlowMatcherWrapper
)

# Initialize VoiceBox without text conditioning
model = VoiceBox(
    dim = 512,
    num_cond_tokens = 500,
    depth = 2,
    dim_head = 64,
    heads = 16,
    condition_on_text = False
)

cfm_wrapper = ConditionalFlowMatcherWrapper(
    voicebox = model
)

# --- Training ---
x = torch.randn(2, 1024, 512)
loss = cfm_wrapper(x)
loss.backward()

# --- Sampling ---
cond = torch.randn(2, 1024, 512)
sampled = cfm_wrapper.sample(cond = cond) # (2, 1024, 512)

What's inside voicebox-pytorch