xscdvfaaqqq/DiffRhythm

1# Copyright (c) 2025 ASLP-LAB2#               2025 Ziqian Ning   (ningziqian@mail.nwpu.edu.cn)3#               2025 Huakang Chen  (huakang@mail.nwpu.edu.cn)4#               2025 Yuepeng Jiang (Jiangyp@mail.nwpu.edu.cn)5#6# Licensed under the Apache License, Version 2.0 (the "License");7# you may not use this file except in compliance with the License.8# You may obtain a copy of the License at9 10#     http://www.apache.org/licenses/LICENSE-2.011 12# Unless required by applicable law or agreed to in writing, software13# distributed under the License is distributed on an "AS IS" BASIS,14# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.15# See the License for the specific language governing permissions and16# limitations under the License.17 18""" This implementation is adapted from github repo:19    https://github.com/SWivid/F5-TTS.20"""21 22from __future__ import annotations23 24import torch25from torch import nn26import torch27 28from transformers.models.llama.modeling_llama import LlamaDecoderLayer, LlamaRotaryEmbedding29from transformers.models.llama import LlamaConfig30 31from diffrhythm.model.modules import (32    TimestepEmbedding,33    ConvNeXtV2Block,34    ConvPositionEmbedding,35    AdaLayerNormZero_Final,36    precompute_freqs_cis,37    get_pos_embed_indices,38    _prepare_decoder_attention_mask,39)40 41# Text embedding42class TextEmbedding(nn.Module):43    def __init__(self, text_num_embeds, text_dim, max_pos, conv_layers=0, conv_mult=2):44        super().__init__()45        self.text_embed = nn.Embedding(text_num_embeds + 1, text_dim)  # use 0 as filler token46 47        if conv_layers > 0:48            self.extra_modeling = True49            self.precompute_max_pos = max_pos  # ~44s of 24khz audio50            self.register_buffer("freqs_cis", precompute_freqs_cis(text_dim, self.precompute_max_pos), persistent=False)51            self.text_blocks = nn.Sequential(52                *[ConvNeXtV2Block(text_dim, text_dim * conv_mult) for _ in range(conv_layers)]53            )54        else:55            self.extra_modeling = False56 57    def forward(self, text: int["b nt"], seq_len, drop_text=False):  # noqa: F72258        batch, text_len = text.shape[0], text.shape[1]59 60        if drop_text:  # cfg for text61            text = torch.zeros_like(text)62 63        text = self.text_embed(text)  # b n -> b n d64 65        # possible extra modeling66        if self.extra_modeling:67            # sinus pos emb68            batch_start = torch.zeros((batch,), dtype=torch.long)69            pos_idx = get_pos_embed_indices(batch_start, seq_len, max_pos=self.precompute_max_pos)70            text_pos_embed = self.freqs_cis[pos_idx]71            text = text + text_pos_embed72 73            # convnextv2 blocks74            text = self.text_blocks(text)75 76        return text77 78 79# noised input audio and context mixing embedding80class InputEmbedding(nn.Module):81    def __init__(self, mel_dim, text_dim, out_dim, cond_dim):82        super().__init__()83        self.proj = nn.Linear(mel_dim * 2 + text_dim + cond_dim * 2, out_dim)84        self.conv_pos_embed = ConvPositionEmbedding(dim=out_dim)85 86    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: F72287        if drop_audio_cond:  # cfg for cond audio88            cond = torch.zeros_like(cond)89        style_emb = style_emb.unsqueeze(1).repeat(1, x.shape[1], 1)90        time_emb = time_emb.unsqueeze(1).repeat(1, x.shape[1], 1)91        x = self.proj(torch.cat((x, cond, text_embed, style_emb, time_emb), dim=-1))92        x = self.conv_pos_embed(x) + x93        return x94 95 96# Transformer backbone using Llama blocks97class DiT(nn.Module):98    def __init__(99        self,100        *,101        dim,102        depth=8,103        heads=8,104        dim_head=64,105        dropout=0.1,106        ff_mult=4,107        mel_dim=100,108        text_num_embeds=256,109        text_dim=None,110        conv_layers=0,111        long_skip_connection=False,112        max_frames=2048113    ):114        super().__init__()115        116        self.max_frames = max_frames117 118        cond_dim = 512119        self.time_embed = TimestepEmbedding(cond_dim)120        self.start_time_embed = TimestepEmbedding(cond_dim)121        self.duration_time_embed = TimestepEmbedding(cond_dim) if self.max_frames == 6144 else None122        if text_dim is None:123            text_dim = mel_dim124        self.text_embed = TextEmbedding(text_num_embeds, text_dim, conv_layers=conv_layers, max_pos=self.max_frames)125        self.input_embed = InputEmbedding(mel_dim, text_dim, dim, cond_dim=cond_dim)126 127        self.dim = dim128        self.depth = depth129 130        llama_config = LlamaConfig(hidden_size=dim, intermediate_size=dim * ff_mult, hidden_act='silu', max_position_embeddings=self.max_frames)131        llama_config._attn_implementation = 'sdpa'132        self.transformer_blocks = nn.ModuleList(133            [LlamaDecoderLayer(llama_config, layer_idx=i) for i in range(depth)]134        )135        self.rotary_emb = LlamaRotaryEmbedding(config=llama_config)136        self.long_skip_connection = nn.Linear(dim * 2, dim, bias=False) if long_skip_connection else None137 138        self.text_fusion_linears = nn.ModuleList(139            [140                nn.Sequential(141                    nn.Linear(cond_dim, dim),142                    nn.SiLU()143                ) for i in range(depth // 2)144            ]145        )146        for layer in self.text_fusion_linears:147            for p in layer.parameters():148                p.detach().zero_()149 150        self.norm_out = AdaLayerNormZero_Final(dim, cond_dim)  # final modulation151        self.proj_out = nn.Linear(dim, mel_dim)152 153    def forward_timestep_invariant(self, text, seq_len, drop_text, start_time):154        s_t = self.start_time_embed(start_time)155        text_embed = self.text_embed(text, seq_len, drop_text=drop_text)156        text_residuals = []157        for layer in self.text_fusion_linears:158            text_residual = layer(text_embed)159            text_residuals.append(text_residual)160        return s_t, text_embed, text_residuals161 162 163    def forward(164        self,165        x: float["b n d"],  # nosied input audio  # noqa: F722166        cond: float["b n d"],  # masked cond audio  # noqa: F722167        text: int["b nt"],  # text  # noqa: F722168        time: float["b"] | float[""],  # time step  # noqa: F821 F722169        drop_audio_cond,  # cfg for cond audio170        drop_text,  # cfg for text171        drop_prompt=False,172        style_prompt=None, # [b d t]173        start_time=None,174        duration=None175    ):176 177        batch, seq_len = x.shape[0], x.shape[1]178        if time.ndim == 0:179            time = time.repeat(batch)180 181        # t: conditioning time, c: context (text + masked cond audio), x: noised input audio182        t = self.time_embed(time)183        s_t = self.start_time_embed(start_time)184        d_t = self.duration_time_embed(duration) if self.max_frames == 6144 else torch.zeros_like(s_t)185        c = t + s_t + d_t186        text_embed = self.text_embed(text, seq_len, drop_text=drop_text)187 188        if drop_prompt:189            style_prompt = torch.zeros_like(style_prompt)190        191        style_embed = style_prompt # [b, 512]192 193        x = self.input_embed(x, cond, text_embed, style_embed, c, drop_audio_cond=drop_audio_cond)194 195        if self.long_skip_connection is not None:196            residual = x197 198        pos_ids = torch.arange(x.shape[1], device=x.device)199        pos_ids = pos_ids.unsqueeze(0).repeat(x.shape[0], 1)200        rotary_embed = self.rotary_emb(x, pos_ids)201        202        attention_mask = torch.ones(203            (batch, seq_len),204            dtype=torch.bool,205            device=x.device,206        )207        attention_mask = _prepare_decoder_attention_mask(208            attention_mask,209            (batch, seq_len),210            x,211        )212 213        for i, block in enumerate(self.transformer_blocks):214            x, *_ = block(x, attention_mask=attention_mask, position_embeddings=rotary_embed)215            if i < self.depth // 2:216                x = x + self.text_fusion_linears[i](text_embed)217 218        if self.long_skip_connection is not None:219            x = self.long_skip_connection(torch.cat((x, residual), dim=-1))220 221        x = self.norm_out(x, c)222        output = self.proj_out(x)223 224        return output225