prabaerode/zero-shot-tts

1"""2ein notation:3b - batch4n - sequence5nt - text sequence6nw - raw wave length7d - dimension8"""9 10from __future__ import annotations11 12import torch13from torch import nn14import torch.nn.functional as F15 16from x_transformers.x_transformers import RotaryEmbedding17 18from f5_tts.model.modules import (19    TimestepEmbedding,20    ConvNeXtV2Block,21    ConvPositionEmbedding,22    DiTBlock,23    AdaLayerNormZero_Final,24    precompute_freqs_cis,25    get_pos_embed_indices,26)27 28 29# Text embedding30 31 32class TextEmbedding(nn.Module):33    def __init__(self, text_num_embeds, text_dim, conv_layers=0, conv_mult=2):34        super().__init__()35        self.text_embed = nn.Embedding(text_num_embeds + 1, text_dim)  # use 0 as filler token36 37        if conv_layers > 0:38            self.extra_modeling = True39            self.precompute_max_pos = 4096  # ~44s of 24khz audio40            self.register_buffer("freqs_cis", precompute_freqs_cis(text_dim, self.precompute_max_pos), persistent=False)41            self.text_blocks = nn.Sequential(42                *[ConvNeXtV2Block(text_dim, text_dim * conv_mult) for _ in range(conv_layers)]43            )44        else:45            self.extra_modeling = False46 47    def forward(self, text: int["b nt"], seq_len, drop_text=False):  # noqa: F72248        text = text + 1  # use 0 as filler token. preprocess of batch pad -1, see list_str_to_idx()49        text = text[:, :seq_len]  # curtail if character tokens are more than the mel spec tokens50        batch, text_len = text.shape[0], text.shape[1]51        text = F.pad(text, (0, seq_len - text_len), value=0)52 53        if drop_text:  # cfg for text54            text = torch.zeros_like(text)55 56        text = self.text_embed(text)  # b n -> b n d57 58        # possible extra modeling59        if self.extra_modeling:60            # sinus pos emb61            batch_start = torch.zeros((batch,), dtype=torch.long)62            pos_idx = get_pos_embed_indices(batch_start, seq_len, max_pos=self.precompute_max_pos)63            text_pos_embed = self.freqs_cis[pos_idx]64            text = text + text_pos_embed65 66            # convnextv2 blocks67            text = self.text_blocks(text)68 69        return text70 71 72# noised input audio and context mixing embedding73 74 75class InputEmbedding(nn.Module):76    def __init__(self, mel_dim, text_dim, out_dim):77        super().__init__()78        self.proj = nn.Linear(mel_dim * 2 + text_dim, out_dim)79        self.conv_pos_embed = ConvPositionEmbedding(dim=out_dim)80 81    def forward(self, x: float["b n d"], cond: float["b n d"], text_embed: float["b n d"], drop_audio_cond=False):  # noqa: F72282        if drop_audio_cond:  # cfg for cond audio83            cond = torch.zeros_like(cond)84 85        x = self.proj(torch.cat((x, cond, text_embed), dim=-1))86        x = self.conv_pos_embed(x) + x87        return x88 89 90# Transformer backbone using DiT blocks91 92 93class DiT(nn.Module):94    def __init__(95        self,96        *,97        dim,98        depth=8,99        heads=8,100        dim_head=64,101        dropout=0.1,102        ff_mult=4,103        mel_dim=100,104        text_num_embeds=256,105        text_dim=None,106        conv_layers=0,107        long_skip_connection=False,108    ):109        super().__init__()110 111        self.time_embed = TimestepEmbedding(dim)112        if text_dim is None:113            text_dim = mel_dim114        self.text_embed = TextEmbedding(text_num_embeds, text_dim, conv_layers=conv_layers)115        self.input_embed = InputEmbedding(mel_dim, text_dim, dim)116 117        self.rotary_embed = RotaryEmbedding(dim_head)118 119        self.dim = dim120        self.depth = depth121 122        self.transformer_blocks = nn.ModuleList(123            [DiTBlock(dim=dim, heads=heads, dim_head=dim_head, ff_mult=ff_mult, dropout=dropout) for _ in range(depth)]124        )125        self.long_skip_connection = nn.Linear(dim * 2, dim, bias=False) if long_skip_connection else None126 127        self.norm_out = AdaLayerNormZero_Final(dim)  # final modulation128        self.proj_out = nn.Linear(dim, mel_dim)129 130    def forward(131        self,132        x: float["b n d"],  # nosied input audio  # noqa: F722133        cond: float["b n d"],  # masked cond audio  # noqa: F722134        text: int["b nt"],  # text  # noqa: F722135        time: float["b"] | float[""],  # time step  # noqa: F821 F722136        drop_audio_cond,  # cfg for cond audio137        drop_text,  # cfg for text138        mask: bool["b n"] | None = None,  # noqa: F722139    ):140        batch, seq_len = x.shape[0], x.shape[1]141        if time.ndim == 0:142            time = time.repeat(batch)143 144        # t: conditioning time, c: context (text + masked cond audio), x: noised input audio145        t = self.time_embed(time)146        text_embed = self.text_embed(text, seq_len, drop_text=drop_text)147        x = self.input_embed(x, cond, text_embed, drop_audio_cond=drop_audio_cond)148 149        rope = self.rotary_embed.forward_from_seq_len(seq_len)150 151        if self.long_skip_connection is not None:152            residual = x153 154        for block in self.transformer_blocks:155            x = block(x, t, mask=mask, rope=rope)156 157        if self.long_skip_connection is not None:158            x = self.long_skip_connection(torch.cat((x, residual), dim=-1))159 160        x = self.norm_out(x, t)161        output = self.proj_out(x)162 163        return output164