espnet3.systems.f5tts.src.f5tts.F5TTS

source

class espnet3.systems.f5tts.src.f5tts.F5TTS(token_list: str | Sequence[str], feats_extract_config: Dict[str, Any] | None = None, hidden_size: int = 1024, depth: int = 22, attention_heads: int = 16, attention_head_size: int = 64, feed_forward_multiplier: int = 2, text_embedding_size: int | None = 512, convolution_layers: int = 4, dropout: float = 0.1, query_key_normalization: str | None = None, text_mask_padding: bool = True, rotary_attention_heads: int | None = None, long_skip_connection: bool = False, checkpoint_activations: bool = False, sigma: float = 0.0, audio_drop_probability: float = 0.3, condition_drop_probability: float = 0.2, mask_fraction_range: Tuple[float, float] = (0.7, 1.0), ode_solver_method: str = 'euler')

Bases: Module

F5-TTS flow-matching model, mel front end included.

Build the mel front end, the DiT backbone and the flow-matching wrap.

Hyper-parameters are spelled out rather than carrying upstream F5-TTS’s abbreviations. To port a published F5-TTS config (F5TTS_Base.yaml / F5TTS_Small.yaml, or the paper’s Table 3), translate with:

Upstream F5-TTS	This class
dim	hidden_size
heads	attention_heads
dim_head	attention_head_size
ff_mult	feed_forward_multiplier
text_dim	text_embedding_size
conv_layers	convolution_layers
qk_norm	query_key_normalization
pe_attn_head	rotary_attention_heads
audio_drop_prob	audio_drop_probability
cond_drop_prob	condition_drop_probability
frac_lengths_mask	mask_fraction_range
odeint_kwargs["method"]	ode_solver_method

Upstream’s mel_spec_kwargs has no counterpart: the flow-matching wrapper’s mel is derived from feats_extract_config so the training and reference front ends cannot drift apart.

depth, dropout, sigma, text_mask_padding, long_skip_connection and checkpoint_activations keep their upstream names, which were already unabbreviated.

• Parameters:
  • token_list – Path to the token file, or the token list itself. Its length becomes the vocabulary size.
  • feats_extract_config – Keyword arguments for VocoderMelSpec. Its n_mels becomes the model’s mel dimension, so the recipe never states that dimension twice.
  • hidden_size – Width of the DiT residual stream.
  • depth – Number of DiT blocks.
  • attention_heads – Attention heads per block.
  • attention_head_size – Width of one attention head.
  • feed_forward_multiplier – Feed-forward width as a multiple of hidden_size.
  • text_embedding_size – Width of the text embedding; None falls back to the mel dimension.
  • convolution_layers – ConvNeXt layers in the text encoder.
  • dropout – Dropout probability inside the backbone.
  • query_key_normalization – None or "rms_norm".
  • text_mask_padding – Mask padded text positions in the text encoder.
  • rotary_attention_heads – How many leading attention heads get rotary position embeddings. None applies them to every head.
  • long_skip_connection – Add the DiT long skip connection.
  • checkpoint_activations – Trade compute for memory via activation checkpointing.
  • sigma – Flow-matching noise scale.
  • audio_drop_probability – Probability of dropping the audio condition during training, for classifier-free guidance.
  • condition_drop_probability – Probability of dropping both conditions together, for classifier-free guidance.
  • mask_fraction_range – (min, max) fraction of each utterance to mask for prediction, drawn uniformly per sample.
  • ode_solver_method – "euler" or "midpoint".
• Raises:RuntimeError – If token_list is neither a path nor a sequence.

Example

model:            # F5TTS_Small; omit the sizes for F5TTS_Base
  _target_: espnet3.systems.f5tts.src.f5tts.F5TTS
  token_list: ${data_dir}/tokens/char_tokens.txt
  feats_extract_config:
    fs: 24000
    n_fft: 1024
    hop_length: 256
    win_length: 1024
    n_mels: 100
  hidden_size: 768
  depth: 18
  attention_heads: 12
  mask_fraction_range: [0.7, 1.0]

######## NOTE Changing the backbone sizes changes the parameter shapes, so an existing checkpoint will not load into a resized model.

collect_feats(text: Tensor, text_lengths: Tensor, speech: Tensor, speech_lengths: Tensor, **kwargs) → Dict[str, Tensor]

Return the mel features for the stats-collection stage.

• Parameters:
  • text – Token ids [B, T_text]. Unused, part of the contract.
  • text_lengths – [B]. Unused, part of the contract.
  • speech – Waveform batch [B, T_wav].
  • speech_lengths – [B] valid sample counts.
  • **kwargs – Extra batch fields this model does not use.
• Returns:{"feats": mel[B, T, mel_dim], "feats_lengths": [B]}.

Example

>>> sorted(model.collect_feats(
...     text=text, text_lengths=text_lengths,
...     speech=speech, speech_lengths=speech_lengths,
... ))
['feats', 'feats_lengths']

######## NOTE espnet3.components.data.collect_stats calls this to write the feats_shape files the length-based batch sampler reads. F5 needs no feature normalization, so only the shapes are used.

forward(text: Tensor, text_lengths: Tensor, speech: Tensor, speech_lengths: Tensor, **kwargs) → Tuple[Tensor, Dict[str, Tensor], Tensor | None]

Flow-matching training/validation step.

• Parameters:
  • text – Token ids [B, T_text] (padded with 0 by CommonCollateFn).
  • text_lengths – [B].
  • speech – Waveform batch [B, T_wav].
  • speech_lengths – [B] valid sample counts.
  • **kwargs – Extra batch fields this model does not use.
• Returns: a scalar flow-matching loss, {"loss": detached loss} for logging, and None for the weight.
• Return type: Tuple of (loss, stats, weight), the ESPnet3 training contract

Example

>>> loss, stats, weight = model(
...     text=text, text_lengths=text_lengths,
...     speech=speech, speech_lengths=speech_lengths,
... )

######## NOTE The span masked for prediction is drawn at random each call (mask_fraction_range), so the loss is stochastic: two calls on the same batch differ unless the RNG is reseeded.

inference(text: Tensor, speech: Tensor | None = None, duration: int | None = None, steps: int = 32, guidance_strength: float = 2.0, sway_sampling_coefficient: float = -1.0, **kwargs) → Dict[str, Tensor]

Generate mel for text given a reference, the minimal path.

F5 is zero-shot: it needs a reference. speech is the reference mel or waveform; text should be the (ref + target) token ids, and for the full cross-speaker recipe protocol use espnet3.systems.f5tts.inference.F5TTSInference, which handles reference pairing and vocoding.

• Parameters:
  • text – Reference + target token ids [T_text], unbatched.
  • speech – Reference audio, either mel [T_ref, mel_dim] or raw waveform [T_wav]. A waveform is converted with the same mel front end CFM uses, so the reference prefix is stripped from the result either way. Required.
  • duration – Total mel length to generate, reference included. Defaults to twice the reference length.
  • steps – Number of ODE solver steps.
  • guidance_strength – Classifier-free guidance scale (upstream cfg_strength).
  • sway_sampling_coefficient – Sway sampling coefficient for timestep schedule; negative values front-load the steps.
  • **kwargs – Ignored.
• Returns:{"feat_gen": mel[T_gen, mel_dim]}, with the reference prefix stripped so only the generated span is returned.
• Raises:RuntimeError – If speech is None. F5 is zero-shot and cannot generate without a reference.

Example

>>> out = model.inference(text=token_ids, speech=ref_mel)
>>> out["feat_gen"].shape
torch.Size([T_gen, 100])

######## NOTE This path does no vocoding. The default duration heuristic is a fallback, and output length is driven by the reference length, so use F5TTSInference for the recipe protocol.