dskill/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 F16from transformers.models.llama.modeling_llama import LlamaDecoderLayer, LlamaRotaryEmbedding17from transformers.models.llama import LlamaConfig18from torch.utils.checkpoint import checkpoint19 20from diffrhythm.model.modules import (21    TimestepEmbedding,22    ConvNeXtV2Block,23    ConvPositionEmbedding,24    DiTBlock,25    AdaLayerNormZero_Final,26    precompute_freqs_cis,27    get_pos_embed_indices,28)29# from liger_kernel.transformers import apply_liger_kernel_to_llama30# apply_liger_kernel_to_llama()31 32# Text embedding33class TextEmbedding(nn.Module):34    def __init__(self, text_num_embeds, text_dim, conv_layers=0, conv_mult=2):35        super().__init__()36        self.text_embed = nn.Embedding(text_num_embeds + 1, text_dim)  # use 0 as filler token37 38        if conv_layers > 0:39            self.extra_modeling = True40            self.precompute_max_pos = 4096  # ~44s of 24khz audio41            self.register_buffer("freqs_cis", precompute_freqs_cis(text_dim, self.precompute_max_pos), persistent=False)42            self.text_blocks = nn.Sequential(43                *[ConvNeXtV2Block(text_dim, text_dim * conv_mult) for _ in range(conv_layers)]44            )45        else:46            self.extra_modeling = False47 48    def forward(self, text: int["b nt"], seq_len, drop_text=False):  # noqa: F72249        batch, text_len = text.shape[0], text.shape[1]50 51        if drop_text:  # cfg for text52            text = torch.zeros_like(text)53 54        text = self.text_embed(text)  # b n -> b n d55 56        # possible extra modeling57        if self.extra_modeling:58            # sinus pos emb59            batch_start = torch.zeros((batch,), dtype=torch.long)60            pos_idx = get_pos_embed_indices(batch_start, seq_len, max_pos=self.precompute_max_pos)61            text_pos_embed = self.freqs_cis[pos_idx]62            text = text + text_pos_embed63 64            # convnextv2 blocks65            text = self.text_blocks(text)66 67        return text68 69 70# noised input audio and context mixing embedding71class InputEmbedding(nn.Module):72    def __init__(self, mel_dim, text_dim, out_dim, cond_dim):73        super().__init__()74        self.proj = nn.Linear(mel_dim * 2 + text_dim + cond_dim * 2, out_dim)75        self.conv_pos_embed = ConvPositionEmbedding(dim=out_dim)76 77    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: F72278        if drop_audio_cond:  # cfg for cond audio79            cond = torch.zeros_like(cond)80 81        style_emb = style_emb.unsqueeze(1).repeat(1, x.shape[1], 1)82        time_emb = time_emb.unsqueeze(1).repeat(1, x.shape[1], 1)83        x = self.proj(torch.cat((x, cond, text_embed, style_emb, time_emb), dim=-1))84        x = self.conv_pos_embed(x) + x85        return x86 87 88# Transformer backbone using DiT blocks89 90 91class DiT(nn.Module):92    def __init__(93        self,94        *,95        dim,96        depth=8,97        heads=8,98        dim_head=64,99        dropout=0.1,100        ff_mult=4,101        mel_dim=100,102        text_num_embeds=256,103        text_dim=None,104        conv_layers=0,105        long_skip_connection=False,106        use_style_prompt=False107    ):108        super().__init__()109 110        cond_dim = 512111        self.time_embed = TimestepEmbedding(cond_dim)112        self.start_time_embed = TimestepEmbedding(cond_dim)113        if text_dim is None:114            text_dim = mel_dim115        self.text_embed = TextEmbedding(text_num_embeds, text_dim, conv_layers=conv_layers)116        self.input_embed = InputEmbedding(mel_dim, text_dim, dim, cond_dim=cond_dim)117 118 119        self.dim = dim120        self.depth = depth121 122        llama_config = LlamaConfig(hidden_size=dim, intermediate_size=dim * ff_mult, hidden_act='silu')123        llama_config._attn_implementation = 'sdpa'124 125        self.transformer_blocks = nn.ModuleList(126            [LlamaDecoderLayer(llama_config, layer_idx=i) for i in range(depth)]127        )128        self.rotary_emb = LlamaRotaryEmbedding(config=llama_config)129        self.long_skip_connection = nn.Linear(dim * 2, dim, bias=False) if long_skip_connection else None130 131        self.text_fusion_linears = nn.ModuleList(132            [133                nn.Sequential(134                    nn.Linear(cond_dim, dim),135                    nn.SiLU()136                ) for i in range(depth // 2)137            ]138        )139        for layer in self.text_fusion_linears:140            for p in layer.parameters():141                p.detach().zero_()142 143        self.norm_out = AdaLayerNormZero_Final(dim, cond_dim)  # final modulation144        self.proj_out = nn.Linear(dim, mel_dim)145 146 147    def forward_timestep_invariant(self, text, seq_len, drop_text, start_time):148        s_t = self.start_time_embed(start_time)149        text_embed = self.text_embed(text, seq_len, drop_text=drop_text)150        text_residuals = []151        for layer in self.text_fusion_linears:152            text_residual = layer(text_embed)153            text_residuals.append(text_residual)154        return s_t, text_embed, text_residuals155 156 157    def forward(158        self,159        x: float["b n d"],  # nosied input audio  # noqa: F722160        text_embed: int["b nt"],  # text  # noqa: F722161        text_residuals,162        cond: float["b n d"],  # masked cond audio  # noqa: F722163        time: float["b"] | float[""],  # time step  # noqa: F821 F722164        drop_audio_cond,  # cfg for cond audio165        drop_prompt=False,166        style_prompt=None, # [b d t]167        start_time=None,168    ):169        batch, seq_len = x.shape[0], x.shape[1]170        if time.ndim == 0:171            time = time.repeat(batch)172 173        t = self.time_embed(time)174        c = t + start_time175 176        if drop_prompt:177            style_prompt = torch.zeros_like(style_prompt)178        179        style_embed = style_prompt # [b, 512]180 181        x = self.input_embed(x, cond, text_embed, style_embed, c, drop_audio_cond=drop_audio_cond)182 183        if self.long_skip_connection is not None:184            residual = x185 186        pos_ids = torch.arange(x.shape[1], device=x.device)187        pos_ids = pos_ids.unsqueeze(0).repeat(x.shape[0], 1)188        rotary_embed = self.rotary_emb(x, pos_ids)189 190        for i, block in enumerate(self.transformer_blocks):191            x, *_ = block(x, position_embeddings=rotary_embed)192            if i < self.depth // 2:193                x = x + text_residuals[i]194 195        if self.long_skip_connection is not None:196            x = self.long_skip_connection(torch.cat((x, residual), dim=-1))197 198        x = self.norm_out(x, c)199        output = self.proj_out(x)200 201        return output202