""" ein notation: b - batch n - sequence nt - text sequence nw - raw wave length d - dimension """ from __future__ import annotations import torch from torch import nn from x_transformers.x_transformers import RotaryEmbedding from f5_tts.model_new.modules import (AdaLayerNorm_Final, ConvPositionEmbedding, MMDiTBlock, TimestepEmbedding, get_pos_embed_indices, precompute_freqs_cis) # text embedding class TextEmbedding(nn.Module): def __init__(self, out_dim, text_num_embeds, mask_padding=True): super().__init__() self.text_embed = nn.Embedding( text_num_embeds + 1, out_dim ) # will use 0 as filler token self.mask_padding = mask_padding # mask filler and batch padding tokens or not self.precompute_max_pos = 1024 self.register_buffer( "freqs_cis", precompute_freqs_cis(out_dim, self.precompute_max_pos), persistent=False, ) def forward( self, text: int["b nt"], drop_text=False ) -> int["b nt d"]: # noqa: F722 text = ( text + 1 ) # use 0 as filler token. preprocess of batch pad -1, see list_str_to_idx() if self.mask_padding: text_mask = text == 0 if drop_text: # cfg for text text = torch.zeros_like(text) text = self.text_embed(text) # b nt -> b nt d # sinus pos emb batch_start = torch.zeros((text.shape[0],), dtype=torch.long) batch_text_len = text.shape[1] pos_idx = get_pos_embed_indices( batch_start, batch_text_len, max_pos=self.precompute_max_pos ) text_pos_embed = self.freqs_cis[pos_idx] text = text + text_pos_embed if self.mask_padding: text = text.masked_fill( text_mask.unsqueeze(-1).expand(-1, -1, text.size(-1)), 0.0 ) return text # noised input & masked cond audio embedding class AudioEmbedding(nn.Module): def __init__(self, in_dim, out_dim): super().__init__() self.linear = nn.Linear(2 * in_dim, out_dim) self.conv_pos_embed = ConvPositionEmbedding(out_dim) def forward( self, x: float["b n d"], cond: float["b n d"], drop_audio_cond=False ): # noqa: F722 if drop_audio_cond: cond = torch.zeros_like(cond) x = torch.cat((x, cond), dim=-1) x = self.linear(x) x = self.conv_pos_embed(x) + x return x # Transformer backbone using MM-DiT blocks class MMDiT(nn.Module): def __init__( self, *, dim, depth=8, heads=8, dim_head=64, dropout=0.1, ff_mult=4, mel_dim=100, text_num_embeds=256, text_mask_padding=True, qk_norm=None, ): super().__init__() self.time_embed = TimestepEmbedding(dim) self.text_embed = TextEmbedding( dim, text_num_embeds, mask_padding=text_mask_padding ) self.text_cond, self.text_uncond = None, None # text cache self.audio_embed = AudioEmbedding(mel_dim, dim) self.rotary_embed = RotaryEmbedding(dim_head) self.dim = dim self.depth = depth self.transformer_blocks = nn.ModuleList( [ MMDiTBlock( dim=dim, heads=heads, dim_head=dim_head, dropout=dropout, ff_mult=ff_mult, context_pre_only=i == depth - 1, qk_norm=qk_norm, ) for i in range(depth) ] ) self.norm_out = AdaLayerNorm_Final(dim) # final modulation self.proj_out = nn.Linear(dim, mel_dim) self.initialize_weights() def initialize_weights(self): # Zero-out AdaLN layers in MMDiT blocks: for block in self.transformer_blocks: nn.init.constant_(block.attn_norm_x.linear.weight, 0) nn.init.constant_(block.attn_norm_x.linear.bias, 0) nn.init.constant_(block.attn_norm_c.linear.weight, 0) nn.init.constant_(block.attn_norm_c.linear.bias, 0) # Zero-out output layers: nn.init.constant_(self.norm_out.linear.weight, 0) nn.init.constant_(self.norm_out.linear.bias, 0) nn.init.constant_(self.proj_out.weight, 0) nn.init.constant_(self.proj_out.bias, 0) def get_input_embed( self, x, # b n d cond, # b n d text, # b nt drop_audio_cond: bool = False, drop_text: bool = False, cache: bool = True, ): if cache: if drop_text: if self.text_uncond is None: self.text_uncond = self.text_embed(text, drop_text=True) c = self.text_uncond else: if self.text_cond is None: self.text_cond = self.text_embed(text, drop_text=False) c = self.text_cond else: c = self.text_embed(text, drop_text=drop_text) x = self.audio_embed(x, cond, drop_audio_cond=drop_audio_cond) return x, c def clear_cache(self): self.text_cond, self.text_uncond = None, None def forward( self, x: float["b n d"], # nosied input audio # noqa: F722 cond: float["b n d"], # masked cond audio # noqa: F722 text: int["b nt"], # text # noqa: F722 time: float["b"] | float[""], # time step # noqa: F821 F722 mask: bool["b n"] | None = None, # noqa: F722 drop_audio_cond: bool = False, # cfg for cond audio drop_text: bool = False, # cfg for text cfg_infer: bool = False, # cfg inference, pack cond & uncond forward cache: bool = False, ): batch = x.shape[0] if time.ndim == 0: time = time.repeat(batch) # t: conditioning (time), c: context (text + masked cond audio), x: noised input audio t = self.time_embed(time) if cfg_infer: # pack cond & uncond forward: b n d -> 2b n d x_cond, c_cond = self.get_input_embed( x, cond, text, drop_audio_cond=False, drop_text=False, cache=cache ) x_uncond, c_uncond = self.get_input_embed( x, cond, text, drop_audio_cond=True, drop_text=True, cache=cache ) x = torch.cat((x_cond, x_uncond), dim=0) c = torch.cat((c_cond, c_uncond), dim=0) t = torch.cat((t, t), dim=0) mask = torch.cat((mask, mask), dim=0) if mask is not None else None else: x, c = self.get_input_embed( x, cond, text, drop_audio_cond=drop_audio_cond, drop_text=drop_text, cache=cache, ) seq_len = x.shape[1] text_len = text.shape[1] rope_audio = self.rotary_embed.forward_from_seq_len(seq_len) rope_text = self.rotary_embed.forward_from_seq_len(text_len) for block in self.transformer_blocks: c, x = block(x, c, t, mask=mask, rope=rope_audio, c_rope=rope_text) x = self.norm_out(x, t) output = self.proj_out(x) return output