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 Guobin Ma     (guobin.ma@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 annotations23from typing import Callable24from random import random25 26import torch27from torch import nn28import torch29import torch.nn.functional as F30from torch.nn.utils.rnn import pad_sequence31 32from torchdiffeq import odeint33 34from diffrhythm.model.utils import (35    exists,36    list_str_to_idx,37    list_str_to_tensor,38    lens_to_mask,39    mask_from_frac_lengths,40)41 42def custom_mask_from_start_end_indices(43    seq_len: int["b"],  # noqa: F82144    latent_pred_segments,45    device,46    max_seq_len47):48    max_seq_len = max_seq_len49    seq = torch.arange(max_seq_len, device=device).long()50 51    res_mask = torch.zeros(max_seq_len, device=device, dtype=torch.bool)52    53    for start, end in latent_pred_segments:54        start = start.unsqueeze(0)55        end = end.unsqueeze(0)56        start_mask = seq[None, :] >= start[:, None]57        end_mask = seq[None, :] < end[:, None]58        res_mask = res_mask | (start_mask & end_mask)59    60    return res_mask61 62class CFM(nn.Module):63    def __init__(64        self,65        transformer: nn.Module,66        sigma=0.0,67        odeint_kwargs: dict = dict(68            method="euler"69        ),70        odeint_options: dict = dict(71            min_step=0.0572        ),73        audio_drop_prob=0.3,74        cond_drop_prob=0.2,75        style_drop_prob=0.1,76        lrc_drop_prob=0.1,77        num_channels=None,78        frac_lengths_mask: tuple[float, float] = (0.7, 1.0),79        vocab_char_map: dict[str:int] | None = None,80        max_frames=204881    ):82        super().__init__()83 84        self.frac_lengths_mask = frac_lengths_mask85 86        self.num_channels = num_channels87 88        # classifier-free guidance89        self.audio_drop_prob = audio_drop_prob90        self.cond_drop_prob = cond_drop_prob91        self.style_drop_prob = style_drop_prob92        self.lrc_drop_prob = lrc_drop_prob93 94        # transformer95        self.transformer = transformer96        dim = transformer.dim97        self.dim = dim98 99        # conditional flow related100        self.sigma = sigma101 102        # sampling related103        self.odeint_kwargs = odeint_kwargs104        105        self.odeint_options = odeint_options106 107        # vocab map for tokenization108        self.vocab_char_map = vocab_char_map109        110        self.max_frames = max_frames111 112    @property113    def device(self):114        return next(self.parameters()).device115 116    @torch.no_grad()117    def sample(118        self,119        cond: float["b n d"] | float["b nw"],  # noqa: F722120        text: int["b nt"] | list[str],  # noqa: F722121        duration: int | int["b"],  # noqa: F821122        *,123        style_prompt = None,124        style_prompt_lens = None,125        negative_style_prompt = None,126        lens: int["b"] | None = None,  # noqa: F821127        steps=32,128        cfg_strength=4.0,129        sway_sampling_coef=None,130        seed: int | None = None,131        max_duration=6144,132        vocoder: Callable[[float["b d n"]], float["b nw"]] | None = None,  # noqa: F722133        no_ref_audio=False,134        duplicate_test=False,135        t_inter=0.1,136        edit_mask=None,137        start_time=None,138        latent_pred_segments=None,139        vocal_flag=False,140        odeint_method="euler",141        song_duration=None,142        batch_infer_num=5143    ):144        self.eval()145        146        self.odeint_kwargs = dict(method=odeint_method)147 148        if next(self.parameters()).dtype == torch.float16:149            cond = cond.half()150 151        # raw wave152        if cond.shape[1] > duration:153            cond = cond[:, :duration, :]154 155        if cond.ndim == 2:156            cond = self.mel_spec(cond)157            cond = cond.permute(0, 2, 1)158            assert cond.shape[-1] == self.num_channels159 160        batch, cond_seq_len, device = *cond.shape[:2], cond.device161        if not exists(lens):162            lens = torch.full((batch,), cond_seq_len, device=device, dtype=torch.long)163 164        # text165        if isinstance(text, list):166            if exists(self.vocab_char_map):167                text = list_str_to_idx(text, self.vocab_char_map).to(device)168            else:169                text = list_str_to_tensor(text).to(device)170            assert text.shape[0] == batch171 172        # duration173        cond_mask = lens_to_mask(lens)174        if edit_mask is not None:175            cond_mask = cond_mask & edit_mask176 177        latent_pred_segments = torch.tensor(latent_pred_segments).to(cond.device)178        fixed_span_mask = custom_mask_from_start_end_indices(cond_seq_len, latent_pred_segments, device=cond.device, max_seq_len=duration)179        fixed_span_mask = fixed_span_mask.unsqueeze(-1)180        step_cond = torch.where(fixed_span_mask, torch.zeros_like(cond), cond)181 182        if isinstance(duration, int):183            duration = torch.full((batch_infer_num,), duration, device=device, dtype=torch.long)184 185        duration = duration.clamp(max=max_duration)186        max_duration = duration.amax()187 188        # duplicate test corner for inner time step oberservation189        if duplicate_test:190            test_cond = F.pad(cond, (0, 0, cond_seq_len, max_duration - 2 * cond_seq_len), value=0.0)191 192        if batch > 1:193            mask = lens_to_mask(duration)194        else:  # save memory and speed up, as single inference need no mask currently195            mask = None196 197        # test for no ref audio198        if no_ref_audio:199            cond = torch.zeros_like(cond)200            201        if vocal_flag:202            style_prompt = negative_style_prompt203            negative_style_prompt = torch.zeros_like(style_prompt)204 205        cond = cond.repeat(batch_infer_num, 1, 1)206        step_cond = step_cond.repeat(batch_infer_num, 1, 1)207        text = text.repeat(batch_infer_num, 1)208        style_prompt = style_prompt.repeat(batch_infer_num, 1)209        negative_style_prompt = negative_style_prompt.repeat(batch_infer_num, 1)210        start_time = start_time.repeat(batch_infer_num)211        fixed_span_mask = fixed_span_mask.repeat(batch_infer_num, 1, 1)212        song_duration = song_duration.repeat(batch_infer_num)213 214        def fn(t, x):215            # predict flow216            pred = self.transformer(217                x=x, cond=step_cond, text=text, time=t, drop_audio_cond=False, drop_text=False, drop_prompt=False,218                style_prompt=style_prompt, start_time=start_time, duration=song_duration219            )220            if cfg_strength < 1e-5:221                return pred222 223            null_pred = self.transformer(224                x=x, cond=step_cond, text=text, time=t, drop_audio_cond=True, drop_text=True, drop_prompt=False,225                style_prompt=negative_style_prompt, start_time=start_time, duration=song_duration226            )227            return pred + (pred - null_pred) * cfg_strength228 229        # noise input230        # to make sure batch inference result is same with different batch size, and for sure single inference231        # still some difference maybe due to convolutional layers232        y0 = []233        for dur in duration:234            if exists(seed):235                torch.manual_seed(seed)236            y0.append(torch.randn(dur, self.num_channels, device=self.device, dtype=step_cond.dtype))237        y0 = pad_sequence(y0, padding_value=0, batch_first=True)238 239        t_start = 0240 241        # duplicate test corner for inner time step oberservation242        if duplicate_test:243            t_start = t_inter244            y0 = (1 - t_start) * y0 + t_start * test_cond245            steps = int(steps * (1 - t_start))246        247        t = torch.linspace(t_start, 1, steps, device=self.device, dtype=step_cond.dtype)248        if sway_sampling_coef is not None:249            t = t + sway_sampling_coef * (torch.cos(torch.pi / 2 * t) - 1 + t)250 251        trajectory = odeint(fn, y0, t, **self.odeint_kwargs)252 253        sampled = trajectory[-1]254        out = sampled255        out = torch.where(fixed_span_mask, out, cond)256 257        if exists(vocoder):258            out = out.permute(0, 2, 1)259            out = vocoder(out)260 261        out = torch.chunk(out, batch_infer_num, dim=0)262        return out, trajectory263 264    def forward(265        self,266        inp: float["b n d"] | float["b nw"],  # mel or raw wave  # noqa: F722267        text: int["b nt"] | list[str],  # noqa: F722268        style_prompt = None,269        style_prompt_lens = None,270        lens: int["b"] | None = None,  # noqa: F821271        noise_scheduler: str | None = None,272        grad_ckpt = False,273        start_time = None,274    ):275 276        batch, seq_len, dtype, device, _σ1 = *inp.shape[:2], inp.dtype, self.device, self.sigma277 278        # lens and mask279        if not exists(lens):280            lens = torch.full((batch,), seq_len, device=device)281 282        mask = lens_to_mask(lens, length=seq_len)  # useless here, as collate_fn will pad to max length in batch283 284        # get a random span to mask out for training conditionally285        frac_lengths = torch.zeros((batch,), device=self.device).float().uniform_(*self.frac_lengths_mask)286        rand_span_mask = mask_from_frac_lengths(lens, frac_lengths, self.max_frames)287 288        if exists(mask):289            rand_span_mask = mask290 291        # mel is x1292        x1 = inp293 294        # x0 is gaussian noise295        x0 = torch.randn_like(x1)296 297        # time step298        time = torch.normal(mean=0, std=1, size=(batch,), device=self.device)299        time = torch.nn.functional.sigmoid(time)300        # TODO. noise_scheduler301 302        # sample xt (φ_t(x) in the paper)303        t = time.unsqueeze(-1).unsqueeze(-1)304        φ = (1 - t) * x0 + t * x1305        flow = x1 - x0306 307        # only predict what is within the random mask span for infilling308        cond = torch.where(rand_span_mask[..., None], torch.zeros_like(x1), x1)309 310        # transformer and cfg training with a drop rate311        drop_audio_cond = random() < self.audio_drop_prob  # p_drop in voicebox paper312        drop_text = random() < self.lrc_drop_prob313        drop_prompt = random() < self.style_drop_prob314 315        # if want rigourously mask out padding, record in collate_fn in dataset.py, and pass in here316        # adding mask will use more memory, thus also need to adjust batchsampler with scaled down threshold for long sequences317        pred = self.transformer(318            x=φ, cond=cond, text=text, time=time, drop_audio_cond=drop_audio_cond, drop_text=drop_text, drop_prompt=drop_prompt,319            style_prompt=style_prompt, start_time=start_time320        )321 322        # flow matching loss323        loss = F.mse_loss(pred, flow, reduction="none")324        loss = loss[rand_span_mask]325 326        return loss.mean(), cond, pred327