cocktailpeanut/DiffRhythm

1"""2ein notation:3b - batch4n - sequence5nt - text sequence6nw - raw wave length7d - dimension8"""9 10from __future__ import annotations11 12import torch13from torch import nn14import torch15import torch.nn.functional as F16 17from x_transformers.x_transformers import RotaryEmbedding18from transformers.models.llama.modeling_llama import LlamaDecoderLayer, LlamaRotaryEmbedding19from transformers.models.llama import LlamaConfig20from torch.utils.checkpoint import checkpoint21 22from diffrhythm.model.modules import (23    TimestepEmbedding,24    ConvNeXtV2Block,25    ConvPositionEmbedding,26    DiTBlock,27    AdaLayerNormZero_Final,28    precompute_freqs_cis,29    get_pos_embed_indices,30)31# from liger_kernel.transformers import apply_liger_kernel_to_llama32# apply_liger_kernel_to_llama()33 34# Text embedding35 36 37class TextEmbedding(nn.Module):38    def __init__(self, text_num_embeds, text_dim, max_pos, conv_layers=0, conv_mult=2):39        super().__init__()40        self.text_embed = nn.Embedding(text_num_embeds + 1, text_dim)  # use 0 as filler token41 42        if conv_layers > 0:43            self.extra_modeling = True44            #self.precompute_max_pos = 4096  # ~44s of 24khz audio45            self.precompute_max_pos = max_pos46            self.register_buffer("freqs_cis", precompute_freqs_cis(text_dim, self.precompute_max_pos), persistent=False)47            self.text_blocks = nn.Sequential(48                *[ConvNeXtV2Block(text_dim, text_dim * conv_mult) for _ in range(conv_layers)]49            )50        else:51            self.extra_modeling = False52 53    def forward(self, text: int["b nt"], seq_len, drop_text=False):  # noqa: F72254        #text = text + 1  # use 0 as filler token. preprocess of batch pad -1, see list_str_to_idx()55        #text = text[:, :seq_len]  # curtail if character tokens are more than the mel spec tokens56        batch, text_len = text.shape[0], text.shape[1]57        #text = F.pad(text, (0, seq_len - text_len), value=0)58 59        if drop_text:  # cfg for text60            text = torch.zeros_like(text)61 62        text = self.text_embed(text)  # b n -> b n d63 64        # possible extra modeling65        if self.extra_modeling:66            # sinus pos emb67            batch_start = torch.zeros((batch,), dtype=torch.long)68            pos_idx = get_pos_embed_indices(batch_start, seq_len, max_pos=self.precompute_max_pos)69            text_pos_embed = self.freqs_cis[pos_idx]70            text = text + text_pos_embed71 72            # convnextv2 blocks73            text = self.text_blocks(text)74 75        return text76 77 78# noised input audio and context mixing embedding79 80 81class InputEmbedding(nn.Module):82    def __init__(self, mel_dim, text_dim, out_dim, cond_dim):83        super().__init__()84        self.proj = nn.Linear(mel_dim * 2 + text_dim + cond_dim * 2, out_dim)85        self.conv_pos_embed = ConvPositionEmbedding(dim=out_dim)86 87    def forward(self, x: float["b n d"], cond: float["b n d"], text_embed: float["b n d"], style_emb, time_emb, drop_audio_cond=False):  # noqa: F72288        if drop_audio_cond:  # cfg for cond audio89            cond = torch.zeros_like(cond)90 91        style_emb = style_emb.unsqueeze(1).repeat(1, x.shape[1], 1)92        time_emb = time_emb.unsqueeze(1).repeat(1, x.shape[1], 1)93        # print(x.shape, cond.shape, text_embed.shape, style_emb.shape, time_emb.shape)94        x = self.proj(torch.cat((x, cond, text_embed, style_emb, time_emb), dim=-1))95        x = self.conv_pos_embed(x) + x96        return x97 98 99# Transformer backbone using DiT blocks100 101 102class DiT(nn.Module):103    def __init__(104        self,105        *,106        dim,107        depth=8,108        heads=8,109        dim_head=64,110        dropout=0.1,111        ff_mult=4,112        mel_dim=100,113        text_num_embeds=256,114        text_dim=None,115        conv_layers=0,116        long_skip_connection=False,117        use_style_prompt=False,118        max_pos=2048119    ):120        super().__init__()121 122        cond_dim = 512123        self.time_embed = TimestepEmbedding(cond_dim)124        self.start_time_embed = TimestepEmbedding(cond_dim)125        if text_dim is None:126            text_dim = mel_dim127        self.text_embed = TextEmbedding(text_num_embeds, text_dim, conv_layers=conv_layers, max_pos=max_pos)128        self.input_embed = InputEmbedding(mel_dim, text_dim, dim, cond_dim=cond_dim)129 130        #self.rotary_embed = RotaryEmbedding(dim_head)131 132        self.dim = dim133        self.depth = depth134 135        #self.transformer_blocks = nn.ModuleList(136        #    [DiTBlock(dim=dim, heads=heads, dim_head=dim_head, ff_mult=ff_mult, dropout=dropout, use_style_prompt=use_style_prompt) for _ in range(depth)]137        #)138        llama_config = LlamaConfig(hidden_size=dim, intermediate_size=dim * ff_mult, hidden_act='silu', max_position_embeddings=max_pos)139        llama_config._attn_implementation = 'sdpa'140        #llama_config._attn_implementation = ''141        self.transformer_blocks = nn.ModuleList(142            [LlamaDecoderLayer(llama_config, layer_idx=i) for i in range(depth)]143        )144        self.rotary_emb = LlamaRotaryEmbedding(config=llama_config)145        self.long_skip_connection = nn.Linear(dim * 2, dim, bias=False) if long_skip_connection else None146 147        self.text_fusion_linears = nn.ModuleList(148            [149                nn.Sequential(150                    nn.Linear(cond_dim, dim),151                    nn.SiLU()152                ) for i in range(depth // 2)153            ]154        )155        for layer in self.text_fusion_linears:156            for p in layer.parameters():157                p.detach().zero_()158 159        self.norm_out = AdaLayerNormZero_Final(dim, cond_dim)  # final modulation160        self.proj_out = nn.Linear(dim, mel_dim)161 162        # if use_style_prompt:163        #     self.prompt_rnn = nn.LSTM(64, cond_dim, 1, batch_first=True)164 165    def forward_timestep_invariant(self, text, seq_len, drop_text, start_time):166        s_t = self.start_time_embed(start_time)167        text_embed = self.text_embed(text, seq_len, drop_text=drop_text)168        text_residuals = []169        for layer in self.text_fusion_linears:170            text_residual = layer(text_embed)171            text_residuals.append(text_residual)172        return s_t, text_embed, text_residuals173 174 175    def forward(176        self,177        x: float["b n d"],  # nosied input audio  # noqa: F722178        text_embed: int["b nt"],  # text  # noqa: F722179        text_residuals,180        cond: float["b n d"],  # masked cond audio  # noqa: F722181        time: float["b"] | float[""],  # time step  # noqa: F821 F722182        drop_audio_cond,  # cfg for cond audio183        drop_prompt=False,184        style_prompt=None, # [b d t]185        start_time=None,186    ):187        batch, seq_len = x.shape[0], x.shape[1]188        if time.ndim == 0:189            time = time.repeat(batch)190 191        t = self.time_embed(time)192        c = t + start_time193 194        if drop_prompt:195            style_prompt = torch.zeros_like(style_prompt)196        197        style_embed = style_prompt # [b, 512]198 199        x = self.input_embed(x, cond, text_embed, style_embed, c, drop_audio_cond=drop_audio_cond)200 201        if self.long_skip_connection is not None:202            residual = x203 204        pos_ids = torch.arange(x.shape[1], device=x.device)205        pos_ids = pos_ids.unsqueeze(0).repeat(x.shape[0], 1)206        rotary_embed = self.rotary_emb(x, pos_ids)207 208        for i, block in enumerate(self.transformer_blocks):209            x, *_ = block(x, position_embeddings=rotary_embed)210            if i < self.depth // 2:211                x = x + text_residuals[i]212 213        if self.long_skip_connection is not None:214            x = self.long_skip_connection(torch.cat((x, residual), dim=-1))215 216        x = self.norm_out(x, c)217        output = self.proj_out(x)218 219        return output220