dskill/DiffRhythm

1"""2ein notation:3b - batch4n - sequence5nt - text sequence6nw - raw wave length7d - dimension8"""9 10from __future__ import annotations11from typing import Callable12from random import random13 14import torch15from torch import nn16import torch17import torch.nn.functional as F18from torch.nn.utils.rnn import pad_sequence19 20from torchdiffeq import odeint21 22from diffrhythm.model.modules import MelSpec23from diffrhythm.model.utils import (24    default,25    exists,26    list_str_to_idx,27    list_str_to_tensor,28    lens_to_mask,29    mask_from_frac_lengths,30)31 32def custom_mask_from_start_end_indices(seq_len: int["b"], start: int["b"], end: int["b"], device, max_seq_len):  # noqa: F722 F82133    max_seq_len = max_seq_len34    seq = torch.arange(max_seq_len, device=device).long()35    start_mask = seq[None, :] >= start[:, None]36    end_mask = seq[None, :] < end[:, None]37    return start_mask & end_mask38 39class CFM(nn.Module):40    def __init__(41        self,42        transformer: nn.Module,43        sigma=0.0,44        odeint_kwargs: dict = dict(45            method="euler" # 'midpoint'46        ),47        odeint_options: dict = dict(48            min_step=0.0549        ),50        audio_drop_prob=0.3,51        cond_drop_prob=0.2,52        style_drop_prob=0.1,53        lrc_drop_prob=0.1,54        num_channels=None,55        frac_lengths_mask: tuple[float, float] = (0.7, 1.0),56        vocab_char_map: dict[str:int] | None = None,57        use_style_prompt: bool = False58    ):59        super().__init__()60 61        self.frac_lengths_mask = frac_lengths_mask62 63        self.num_channels = num_channels64 65        # classifier-free guidance66        self.audio_drop_prob = audio_drop_prob67        self.cond_drop_prob = cond_drop_prob68        self.style_drop_prob = style_drop_prob69        self.lrc_drop_prob = lrc_drop_prob70 71        # transformer72        self.transformer = transformer73        dim = transformer.dim74        self.dim = dim75 76        # conditional flow related77        self.sigma = sigma78 79        # sampling related80        self.odeint_kwargs = odeint_kwargs81        82        self.odeint_options = odeint_options83 84        # vocab map for tokenization85        self.vocab_char_map = vocab_char_map86 87        self.use_style_prompt = use_style_prompt88 89    @property90    def device(self):91        return next(self.parameters()).device92 93    @torch.no_grad()94    def sample(95        self,96        cond: float["b n d"] | float["b nw"],  # noqa: F72297        text: int["b nt"] | list[str],  # noqa: F72298        duration: int | int["b"],  # noqa: F82199        *,100        style_prompt = None,101        style_prompt_lens = None,102        negative_style_prompt = None,103        lens: int["b"] | None = None,  # noqa: F821104        steps=32,105        cfg_strength=4.0,106        sway_sampling_coef=None,107        seed: int | None = None,108        max_duration=4096,109        vocoder: Callable[[float["b d n"]], float["b nw"]] | None = None,  # noqa: F722110        no_ref_audio=False,111        duplicate_test=False,112        t_inter=0.1,113        edit_mask=None,114        start_time=None,115        latent_pred_start_frame=0,116        latent_pred_end_frame=2048,117        vocal_flag=False,118        odeint_method="euler"119    ):120        self.eval()121        122        self.odeint_kwargs = dict(method=odeint_method)123 124        if next(self.parameters()).dtype == torch.float16:125            cond = cond.half()126 127        # raw wave128        129        if cond.shape[1] > duration:130            cond = cond[:, :duration, :]131 132        if cond.ndim == 2:133            cond = self.mel_spec(cond)134            cond = cond.permute(0, 2, 1)135            assert cond.shape[-1] == self.num_channels136 137        batch, cond_seq_len, device = *cond.shape[:2], cond.device138        if not exists(lens):139            lens = torch.full((batch,), cond_seq_len, device=device, dtype=torch.long)140 141        # text142 143        if isinstance(text, list):144            if exists(self.vocab_char_map):145                text = list_str_to_idx(text, self.vocab_char_map).to(device)146            else:147                text = list_str_to_tensor(text).to(device)148            assert text.shape[0] == batch149 150        if exists(text):151            text_lens = (text != -1).sum(dim=-1)152 153 154        # duration155        cond_mask = lens_to_mask(lens)156        if edit_mask is not None:157            cond_mask = cond_mask & edit_mask158 159        latent_pred_start_frame = torch.tensor([latent_pred_start_frame]).to(cond.device)160        latent_pred_end_frame = duration161        latent_pred_end_frame = torch.tensor([latent_pred_end_frame]).to(cond.device)162        fixed_span_mask = custom_mask_from_start_end_indices(cond_seq_len, latent_pred_start_frame, latent_pred_end_frame, device=cond.device, max_seq_len=duration)163 164        fixed_span_mask = fixed_span_mask.unsqueeze(-1)165        step_cond = torch.where(fixed_span_mask, torch.zeros_like(cond), cond)166 167        if isinstance(duration, int):168            duration = torch.full((batch,), duration, device=device, dtype=torch.long)169 170 171        duration = duration.clamp(max=max_duration)172        max_duration = duration.amax()173 174        # duplicate test corner for inner time step oberservation175        if duplicate_test:176            test_cond = F.pad(cond, (0, 0, cond_seq_len, max_duration - 2 * cond_seq_len), value=0.0)177 178 179        if batch > 1:180            mask = lens_to_mask(duration)181        else:  # save memory and speed up, as single inference need no mask currently182            mask = None183 184        # test for no ref audio185        if no_ref_audio:186            cond = torch.zeros_like(cond)187        188        start_time_embed, positive_text_embed, positive_text_residuals = self.transformer.forward_timestep_invariant(text, step_cond.shape[1], drop_text=False, start_time=start_time)189        _, negative_text_embed, negative_text_residuals = self.transformer.forward_timestep_invariant(text, step_cond.shape[1], drop_text=True, start_time=start_time)190 191        if vocal_flag:192            style_prompt = negative_style_prompt193            negative_style_prompt = torch.zeros_like(style_prompt)194            195        text_embed = torch.cat([positive_text_embed, negative_text_embed], 0)196        text_residuals = [torch.cat([a, b], 0) for a, b in zip(positive_text_residuals, negative_text_residuals)]197        step_cond = torch.cat([step_cond, step_cond], 0)198        style_prompt = torch.cat([style_prompt, negative_style_prompt], 0)199        start_time_embed = torch.cat([start_time_embed, start_time_embed], 0)200            201 202        def fn(t, x):203            x = torch.cat([x, x], 0)204            pred = self.transformer(205                x=x, text_embed=text_embed, text_residuals=text_residuals, cond=step_cond, time=t, 206                drop_audio_cond=True, drop_prompt=False, style_prompt=style_prompt, start_time=start_time_embed207            )208 209            positive_pred, negative_pred = pred.chunk(2, 0)210            cfg_pred = positive_pred + (positive_pred - negative_pred) * cfg_strength211 212            return cfg_pred213 214        # noise input215        # to make sure batch inference result is same with different batch size, and for sure single inference216        # still some difference maybe due to convolutional layers217        y0 = []218        for dur in duration:219            if exists(seed):220                torch.manual_seed(seed)221            y0.append(torch.randn(dur, self.num_channels, device=self.device, dtype=step_cond.dtype))222        y0 = pad_sequence(y0, padding_value=0, batch_first=True)223 224        t_start = 0225 226        # duplicate test corner for inner time step oberservation227        if duplicate_test:228            t_start = t_inter229            y0 = (1 - t_start) * y0 + t_start * test_cond230            steps = int(steps * (1 - t_start))231 232        t = torch.linspace(t_start, 1, steps, device=self.device, dtype=step_cond.dtype)233        if sway_sampling_coef is not None:234            t = t + sway_sampling_coef * (torch.cos(torch.pi / 2 * t) - 1 + t)235 236        trajectory = odeint(fn, y0, t, **self.odeint_kwargs)237 238        sampled = trajectory[-1]239        out = sampled240        out = torch.where(fixed_span_mask, out, cond)241 242        if exists(vocoder):243            out = out.permute(0, 2, 1)244            out = vocoder(out)245 246        return out, trajectory247 248    def forward(249        self,250        inp: float["b n d"] | float["b nw"],  # mel or raw wave  # noqa: F722251        text: int["b nt"] | list[str],  # noqa: F722252        style_prompt = None,253        style_prompt_lens = None,254        lens: int["b"] | None = None,  # noqa: F821255        noise_scheduler: str | None = None,256        grad_ckpt = False,257        start_time = None,258    ):259 260        batch, seq_len, dtype, device, _σ1 = *inp.shape[:2], inp.dtype, self.device, self.sigma261 262        # lens and mask263        if not exists(lens):264            lens = torch.full((batch,), seq_len, device=device)265 266        mask = lens_to_mask(lens, length=seq_len)  # useless here, as collate_fn will pad to max length in batch267 268        # get a random span to mask out for training conditionally269        frac_lengths = torch.zeros((batch,), device=self.device).float().uniform_(*self.frac_lengths_mask)270        rand_span_mask = mask_from_frac_lengths(lens, frac_lengths)271 272        if exists(mask):273            rand_span_mask = mask274            # rand_span_mask &= mask275 276        # mel is x1277        x1 = inp278 279        # x0 is gaussian noise280        x0 = torch.randn_like(x1)281 282        # time step283        time = torch.normal(mean=0, std=1, size=(batch,), device=self.device)284        time = torch.nn.functional.sigmoid(time)285        # TODO. noise_scheduler286 287        # sample xt (φ_t(x) in the paper)288        t = time.unsqueeze(-1).unsqueeze(-1)289        φ = (1 - t) * x0 + t * x1290        flow = x1 - x0291 292        # only predict what is within the random mask span for infilling293        cond = torch.where(rand_span_mask[..., None], torch.zeros_like(x1), x1)294 295        # transformer and cfg training with a drop rate296        drop_audio_cond = random() < self.audio_drop_prob  # p_drop in voicebox paper297        drop_text = random() < self.lrc_drop_prob298        drop_prompt = random() < self.style_drop_prob299 300        # if want rigourously mask out padding, record in collate_fn in dataset.py, and pass in here301        # adding mask will use more memory, thus also need to adjust batchsampler with scaled down threshold for long sequences302        pred = self.transformer(303            x=φ, cond=cond, text=text, time=time, drop_audio_cond=drop_audio_cond, drop_text=drop_text, drop_prompt=drop_prompt,304            style_prompt=style_prompt, style_prompt_lens=style_prompt_lens, grad_ckpt=grad_ckpt, start_time=start_time305        )306 307        # flow matching loss308        loss = F.mse_loss(pred, flow, reduction="none")309        loss = loss[rand_span_mask]310 311        return loss.mean(), cond, pred312