prabaerode/zero-shot-tts

1"""2ein notation:3b - batch4n - sequence5nt - text sequence6nw - raw wave length7d - dimension8"""9 10from __future__ import annotations11from typing import Literal12 13import torch14from torch import nn15import torch.nn.functional as F16 17from x_transformers import RMSNorm18from x_transformers.x_transformers import RotaryEmbedding19 20from f5_tts.model.modules import (21    TimestepEmbedding,22    ConvNeXtV2Block,23    ConvPositionEmbedding,24    Attention,25    AttnProcessor,26    FeedForward,27    precompute_freqs_cis,28    get_pos_embed_indices,29)30 31 32# Text embedding33 34 35class TextEmbedding(nn.Module):36    def __init__(self, text_num_embeds, text_dim, conv_layers=0, conv_mult=2):37        super().__init__()38        self.text_embed = nn.Embedding(text_num_embeds + 1, text_dim)  # use 0 as filler token39 40        if conv_layers > 0:41            self.extra_modeling = True42            self.precompute_max_pos = 4096  # ~44s of 24khz audio43            self.register_buffer("freqs_cis", precompute_freqs_cis(text_dim, self.precompute_max_pos), persistent=False)44            self.text_blocks = nn.Sequential(45                *[ConvNeXtV2Block(text_dim, text_dim * conv_mult) for _ in range(conv_layers)]46            )47        else:48            self.extra_modeling = False49 50    def forward(self, text: int["b nt"], seq_len, drop_text=False):  # noqa: F72251        text = text + 1  # use 0 as filler token. preprocess of batch pad -1, see list_str_to_idx()52        text = text[:, :seq_len]  # curtail if character tokens are more than the mel spec tokens53        batch, text_len = text.shape[0], text.shape[1]54        text = F.pad(text, (0, seq_len - text_len), value=0)55 56        if drop_text:  # cfg for text57            text = torch.zeros_like(text)58 59        text = self.text_embed(text)  # b n -> b n d60 61        # possible extra modeling62        if self.extra_modeling:63            # sinus pos emb64            batch_start = torch.zeros((batch,), dtype=torch.long)65            pos_idx = get_pos_embed_indices(batch_start, seq_len, max_pos=self.precompute_max_pos)66            text_pos_embed = self.freqs_cis[pos_idx]67            text = text + text_pos_embed68 69            # convnextv2 blocks70            text = self.text_blocks(text)71 72        return text73 74 75# noised input audio and context mixing embedding76 77 78class InputEmbedding(nn.Module):79    def __init__(self, mel_dim, text_dim, out_dim):80        super().__init__()81        self.proj = nn.Linear(mel_dim * 2 + text_dim, out_dim)82        self.conv_pos_embed = ConvPositionEmbedding(dim=out_dim)83 84    def forward(self, x: float["b n d"], cond: float["b n d"], text_embed: float["b n d"], drop_audio_cond=False):  # noqa: F72285        if drop_audio_cond:  # cfg for cond audio86            cond = torch.zeros_like(cond)87 88        x = self.proj(torch.cat((x, cond, text_embed), dim=-1))89        x = self.conv_pos_embed(x) + x90        return x91 92 93# Flat UNet Transformer backbone94 95 96class UNetT(nn.Module):97    def __init__(98        self,99        *,100        dim,101        depth=8,102        heads=8,103        dim_head=64,104        dropout=0.1,105        ff_mult=4,106        mel_dim=100,107        text_num_embeds=256,108        text_dim=None,109        conv_layers=0,110        skip_connect_type: Literal["add", "concat", "none"] = "concat",111    ):112        super().__init__()113        assert depth % 2 == 0, "UNet-Transformer's depth should be even."114 115        self.time_embed = TimestepEmbedding(dim)116        if text_dim is None:117            text_dim = mel_dim118        self.text_embed = TextEmbedding(text_num_embeds, text_dim, conv_layers=conv_layers)119        self.input_embed = InputEmbedding(mel_dim, text_dim, dim)120 121        self.rotary_embed = RotaryEmbedding(dim_head)122 123        # transformer layers & skip connections124 125        self.dim = dim126        self.skip_connect_type = skip_connect_type127        needs_skip_proj = skip_connect_type == "concat"128 129        self.depth = depth130        self.layers = nn.ModuleList([])131 132        for idx in range(depth):133            is_later_half = idx >= (depth // 2)134 135            attn_norm = RMSNorm(dim)136            attn = Attention(137                processor=AttnProcessor(),138                dim=dim,139                heads=heads,140                dim_head=dim_head,141                dropout=dropout,142            )143 144            ff_norm = RMSNorm(dim)145            ff = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh")146 147            skip_proj = nn.Linear(dim * 2, dim, bias=False) if needs_skip_proj and is_later_half else None148 149            self.layers.append(150                nn.ModuleList(151                    [152                        skip_proj,153                        attn_norm,154                        attn,155                        ff_norm,156                        ff,157                    ]158                )159            )160 161        self.norm_out = RMSNorm(dim)162        self.proj_out = nn.Linear(dim, mel_dim)163 164    def forward(165        self,166        x: float["b n d"],  # nosied input audio  # noqa: F722167        cond: float["b n d"],  # masked cond audio  # noqa: F722168        text: int["b nt"],  # text  # noqa: F722169        time: float["b"] | float[""],  # time step  # noqa: F821 F722170        drop_audio_cond,  # cfg for cond audio171        drop_text,  # cfg for text172        mask: bool["b n"] | None = None,  # noqa: F722173    ):174        batch, seq_len = x.shape[0], x.shape[1]175        if time.ndim == 0:176            time = time.repeat(batch)177 178        # t: conditioning time, c: context (text + masked cond audio), x: noised input audio179        t = self.time_embed(time)180        text_embed = self.text_embed(text, seq_len, drop_text=drop_text)181        x = self.input_embed(x, cond, text_embed, drop_audio_cond=drop_audio_cond)182 183        # postfix time t to input x, [b n d] -> [b n+1 d]184        x = torch.cat([t.unsqueeze(1), x], dim=1)  # pack t to x185        if mask is not None:186            mask = F.pad(mask, (1, 0), value=1)187 188        rope = self.rotary_embed.forward_from_seq_len(seq_len + 1)189 190        # flat unet transformer191        skip_connect_type = self.skip_connect_type192        skips = []193        for idx, (maybe_skip_proj, attn_norm, attn, ff_norm, ff) in enumerate(self.layers):194            layer = idx + 1195 196            # skip connection logic197            is_first_half = layer <= (self.depth // 2)198            is_later_half = not is_first_half199 200            if is_first_half:201                skips.append(x)202 203            if is_later_half:204                skip = skips.pop()205                if skip_connect_type == "concat":206                    x = torch.cat((x, skip), dim=-1)207                    x = maybe_skip_proj(x)208                elif skip_connect_type == "add":209                    x = x + skip210 211            # attention and feedforward blocks212            x = attn(attn_norm(x), rope=rope, mask=mask) + x213            x = ff(ff_norm(x)) + x214 215        assert len(skips) == 0216 217        x = self.norm_out(x)[:, 1:, :]  # unpack t from x218 219        return self.proj_out(x)220