repository·main·Indexed 19 days ago
https://github.com/lucidrains/e2-tts-pytorchA PyTorch implementation of E2-TTS (Embarrassingly Easy Fully Non-Autoregressive Zero-Shot TTS). It features a multi-stream transformer architecture for text and audio, a DurationPredictor for token duration estimation, and support for flow matching training with infilling and velocity consistency loss. The library includes utilities for UTF-8 and English phoneme tokenization, MelSpectrogram conversion, and a high-level E2Trainer for managing the training loop and EMA updates.
interpolated_text = True when initializing the E2TTS class.Install the package using pip:
$ pip install e2-tts-pytorchDuring the forward pass, E2TTS supports infilling by masking out random spans of the conditioning signal.
frac_lengths and mask_from_frac_lengths.cond is constructed such that only the unmasked parts of the input x1 are visible to the model. The masked parts are set to zero.rand_span_mask area, forcing the model to learn to predict the missing segments based on the surrounding context and text.The Transformer class is a multi-stream architecture used for both the main E2TTS model and the duration predictor. It supports:
TextAudioCrossCondition.has_freq_axis is set, the model can handle 4D tensors [b f n d] (batch, frequency, sequence, dimension).AdaptiveRMSNorm and AdaLNZero layers, similar to Diffusion Transformers (DiT).HyperConnections to manage multiple residual streams.registers to improve performance.RotaryEmbedding for positional encoding.RandomFourierEmbed and AdaLNZero.E2TTS implements Classifier-Free Guidance to improve sample quality. The cfg_transformer_with_pred_head method calculates the guidance update by comparing the prediction from the conditioned model against a prediction from a 'null' model (where text conditioning is dropped).
pred using the provided conditioning.null_pred is calculated using a cfg_null_model (which defaults to the current model itself) with drop_text_cond=True.cfg_update = pred - null_pred.remove_parallel_component is True, the update is projected to be orthogonal to the original prediction, controlled by keep_parallel_frac. This follows the approach in [arXiv:2410.02416].pred + cfg_update * cfg_strength.The project provides two primary ways to tokenize text for the models:
'char_utf8'): Converts text to a tensor of UTF-8 bytes. This is the default.'phoneme_en'): Uses g2p_en to convert text into phonemes, which are then mapped to indices. This includes support for extended characters like punctuation and ellipses.Callable[[list[str]], Int['b nt']]. If you do this, you must specify the text_num_embeds parameter in the parent module (E2TTS or DurationPredictor) so the embedding layer can be sized correctly.The DurationPredictor is used to predict audio durations from text. It takes a mel spectrogram and text as input and returns a loss value.
Parameters:
transformer: A dictionary defining the transformer architecture (e.g., dim and depth).Input Shapes:
mel: Tensor of shape (batch, time, channels).text: A list of strings.import torch
from e2_tts_pytorch import DurationPredictor
duration_predictor = DurationPredictor(
transformer = dict(
dim = 512,
depth = 8,
)
)
mel = torch.randn(2, 1024, 100)
text = ['Hello', 'Goodbye']
loss = duration_predictor(mel, text = text)
loss.backward()The E2TTS class implements the E2-TTS model. It can be used for both training (calculating loss) and sampling (generating audio).
Initialization:
duration_predictor: An instance of DurationPredictor.transformer: A dictionary defining the transformer architecture (e.g., dim and depth).interpolated_text (optional): Set to True to use an improvisation where text is interpolated to the length of the audio for conditioning.Methods:
__call__(mel, text): Performs a forward pass. Returns an object containing a .loss attribute for training.sample(mel, text): Generates audio samples based on the provided mel spectrogram and text.import torch
from e2_tts_pytorch import E2TTS, DurationPredictor
duration_predictor = DurationPredictor(
transformer = dict(
dim = 512,
depth = 8,
)
)
e2tts = E2TTS(
duration_predictor = duration_predictor,
transformer = dict(
dim = 512,
depth = 8
),
)
mel = torch.randn(2, 1024, 100)
text = ['Hello', 'Goodbye']
# Training usage
out = e2tts(mel, text = text)
out.loss.backward()
# Sampling usage
sampled = e2tts.sample(mel[:, :5], text = text)When initializing E2Trainer, you can tune the following parameters to control training behavior:
| Parameter | Type | Description |
|---|---|---|
optimizer | Optimizer | Custom optimizer instance. If None, Adopt is used. |
learning_rate | float | Initial learning rate (used if optimizer is None). |
num_warmup_steps | int | Number of steps for the linear warmup scheduler. |
grad_accumulation_steps | int | Number of steps to accumulate gradients before an optimizer step. |
max_grad_norm | float | Maximum gradient norm for clipping (default 1.0). |
accelerate_kwargs | dict | Arguments passed to the accelerate.Accelerator constructor. |
ema_kwargs | dict | Arguments passed to the EMA constructor. |
tensorboard_log_dir | str | Directory for TensorBoard logs. |
The DurationPredictor module is used to predict the duration of audio segments based on text and audio features. It uses a Transformer backbone and a HLGaussLayer for duration modeling.
Key parameters:
transformer: A Transformer instance or a dictionary of its hyperparameters.tokenizer: Supports 'char_utf8', 'phoneme_en', or a custom callable. If using a custom callable, you must provide text_num_embeds.use_regression: Boolean indicating if regression should be used.num_freq_tokens: Number of frequency tokens (if > 1, enables frequency axis processing).from e2_tts_pytorch.e2_tts import DurationPredictor
duration_predictor = DurationPredictor(
transformer = dict(
dim = 512,
depth = 8,
heads = 8
),
tokenizer = 'char_utf8'
)The MelSpec module is a wrapper around torchaudio.transforms.MelSpectrogram that also applies a log transformation to the output. It is used to convert raw audio waveforms into mel spectrograms.
Parameters:
filter_length, hop_length, win_length: Standard STFT parameters.n_mel_channels: Number of mel frequency bins.sampling_rate: The audio sampling rate.normalize: Whether to normalize the spectrogram.from e2_tts_pytorch.e2_tts import MelSpec
import torch
mels = MelSpec(n_mel_channels=100, sampling_rate=24000)
# input shape: [batch, waveform_length]
mel_spec = mels(torch.randn(1, 24000))The sample method is the primary high-level API for generating audio from a conditioning signal (like a mel spectrogram) and optional text. It handles the ODE integration (flow matching) and can optionally convert the resulting mel spectrogram into raw audio using a provided vocoder or the model's internal vocos decoder.
cond: The conditioning signal. Can be a mel spectrogram Float['b n d'] or raw waveform Float['b nw'].text: Optional text conditioning. Accepts Int['b nt'] (tokenized), list[str], or None.duration: Optional target duration. If None, the model uses its internal duration_predictor.steps: Number of ODE integration steps (default: 32).cfg_strength: Classifier-Free Guidance strength.cfg_null_model: An optional E2TTS instance for 'autoguidance'.vocoder: A callable that converts mel spectrograms to audio. If provided, the model's internal vocos is ignored.save_to_filename: If provided, saves the generated audio files to the specified path.Returns a tuple containing:
Float['b n d'].list[Float['_']] (if a vocoder was used or vocos is enabled).# Example sampling with text and a custom vocoder
output_mel, audio = model.sample(
cond=mel_spectrogram,
text=['Hello world'],
steps=50,
cfg_strength=1.5,
vocoder=my_custom_vocoder
)