CAMB-AI/MARS5-TTS

1"""2Discrete multinomial diffusion code adapted from https://github.com/RF5/transfusion-asr,3which in turn is adapted from https://github.com/ehoogeboom/multinomial_diffusion.4 5Please see the original repo (https://github.com/ehoogeboom/multinomial_diffusion) and paper for full6details on how multinomial diffusion works -- thanks to the original authors!7"""8 9import torch10from torch import Tensor11from torch.functional import F12import numpy as np13from dataclasses import dataclass14from typing import Union15 16# -------------- Multinomial utility functions -----------17 18MIN_LOG_ARG = 1e-7 # originally was 1e-4019 20def log_1_min_a(a): return torch.log((1 - a.exp()).clamp_(min=1e-30))21 22def log_add_exp(a, b):23    maximum = torch.max(a, b)24    return maximum + torch.log(torch.exp(a - maximum) + torch.exp(b - maximum))25 26def extract(a: Tensor, t, x_shape):27    """ Given 1D vector of alpha/alpha_cum/betas, get index at `t` of shape (bs,), and then28    broadcast it to number of dims in `x_shape`. 29    """30    b, *_ = t.shape31    out = a.gather(-1, t)32    return out.reshape(b, *((1,) * (len(x_shape) - 1)))33 34def index_to_log_onehot(x, num_classes, dim=-1, dtype=torch.float32):35    """ Convert indices `x` (bs, ...) to approx one-hot log-probs of shape (bs, ..., num_classes) """36    assert x.max().item() < num_classes, \37        f'Error: {x.max().item()} >= {num_classes}'38    x_onehot = F.one_hot(x, num_classes)39    if dim == 1:40        permute_order = (0, -1) + tuple(range(1, len(x.size())))41        x_onehot = x_onehot.permute(permute_order)42    else: 43        pass44 45    log_x = torch.log(x_onehot.to(dtype).clamp(min=MIN_LOG_ARG)) # so min(log_x) will be -3046 47    return log_x48 49def sum_except_batch(x: Tensor, num_dims=1) -> Tensor:50    '''51    Sums all dimensions except the first.52    Args:53        x: Tensor, shape (batch_size, ...)54        num_dims: int, number of batch dims (default=1)55    Returns:56        x_sum: Tensor, shape (batch_size,)57    '''58    return x.reshape(*x.shape[:num_dims], -1).sum(-1)59 60# -------------- Multinomial diffusion class -------------61 62class MultinomialDiffusion():63    def __init__(self, num_classes, timesteps=100, diffusion_s=0.008,64                 loss_type='vb_stochastic', parametrization='x0', 65                 dtype=torch.float32,66                 device='cpu'):67        super(MultinomialDiffusion, self).__init__()68        assert loss_type in ('vb_stochastic',)69        assert parametrization in ('x0', 'direct')70 71        self.num_classes = num_classes72        self.loss_type = loss_type73        self.num_timesteps = timesteps74        self.parametrization = parametrization75 76        alphas = self.cosine_beta_schedule(timesteps, diffusion_s)77 78        alphas = alphas.to(torch.float64)79        log_alpha = alphas.log()80        log_cumprod_alpha = torch.cumsum(log_alpha, dim=-1)81 82        log_1_min_alpha = log_1_min_a(log_alpha) # = log(betas)83 84        log_1_min_cumprod_alpha = log_1_min_a(log_cumprod_alpha) # = log(1- \bar{a}) 85        a = log_add_exp(log_alpha, log_1_min_alpha) # log(1-beta + beta) = log(1) = 086 87        assert log_add_exp(log_alpha, log_1_min_alpha).abs().sum().item() < 1.e-588        assert log_add_exp(log_cumprod_alpha, log_1_min_cumprod_alpha).abs().sum().item() < 1e-589        assert (torch.cumsum(log_alpha, dim=-1) - log_cumprod_alpha).abs().sum().item() < 1.e-590 91        # Convert to float32 and register buffers.92        self.log_alpha = log_alpha.to(dtype).to(device)93        self.log_1_min_alpha = log_1_min_alpha.to(dtype).to(device)94        self.log_cumprod_alpha = log_cumprod_alpha.to(dtype).to(device)95        self.log_1_min_cumprod_alpha = log_1_min_cumprod_alpha.to(dtype).to(device)96 97    @staticmethod98    def cosine_beta_schedule(timesteps, s=0.008) -> Tensor:99        """100        cosine schedule as proposed in https://arxiv.org/abs/2102.09672 .101        Returns alpha parameters, NOT Beta102        """103        steps = timesteps + 1104        x = torch.linspace(0, timesteps, steps)105        alphas_cumprod = torch.cos(((x / timesteps) + s) / (1 + s) * torch.pi * 0.5) ** 2106        alphas_cumprod = alphas_cumprod / alphas_cumprod[0]107        alphas = (alphas_cumprod[1:] / alphas_cumprod[:-1])108        alphas = torch.clamp(alphas, 0.001, 1.0)109        return torch.sqrt(alphas)110 111    def multinomial_kl(self, log_prob1: Tensor, log_prob2: Tensor, dim=-1) -> Tensor:112        """ Get KL divergence between two categorical distributions specified with `log_prob1` and `log_prob2`.113        Assumed probability dim is `dim` (i.e. log_prob1.exp().sum(dim=`dim`) should be tensor of ones)114        """115        kl = (log_prob1.exp() * (log_prob1 - log_prob2)).sum(dim=dim)116        return kl117 118    def q_pred_one_timestep(self, log_x_t: Tensor, t: Tensor) -> Tensor:119        """ Compute q(x_t | x_{t-1}) = C(x_t | alpha_t * x_{t-1} + (1-alpha_t)/K in the log-domain120        given `log_x_t` as log one-hot encoding of x_t. 121        122        Recall due to symmetry property we can compute123        this value using x_t instead of x_{t-1} (se appendix A of https://arxiv.org/pdf/2102.05379.pdf)124        """125        dt = log_x_t.dtype126        log_alpha_t = extract(self.log_alpha, t, log_x_t.shape).to(dt)127        log_1_min_alpha_t = extract(self.log_1_min_alpha, t, log_x_t.shape).to(dt)128 129        # alpha_t * E[xt] + (1 - alpha_t) 1 / K130        log_probs = log_add_exp(131            log_x_t + log_alpha_t,132            log_1_min_alpha_t - np.log(self.num_classes)133        )134        return log_probs135 136    def q_pred_one_timestep_scaled(self, log_x_t: Tensor, t: Tensor, c: int, jump_len: int) -> Tensor:137        """ Compute q(x_t | x_{t-1}) = C(x_t | alpha_t * x_{t-1} + (1-alpha_t)/K in the log-domain138        given `log_x_t` as log one-hot encoding of x_t. 139        140        Recall due to symmetry property we can compute141        this value using x_t instead of x_{t-1} (se appendix A of https://arxiv.org/pdf/2102.05379.pdf)142        """143        dt = log_x_t.dtype144        log_alpha_t = extract(self.log_alpha, t, log_x_t.shape).to(dt)145        log_1_min_alpha_t = extract(self.log_1_min_alpha, t, log_x_t.shape).to(dt)146 147        # Magic148        xax = torch.arange(0,log_x_t.shape[1],1).to(log_x_t.device)149        aa=log_x_t.shape[1]*(c/jump_len)150        sig = 1/(1+torch.exp(-(xax-aa+20)/8))151        log_alpha_t = (torch.log(1/sig)[None,:,None] + log_alpha_t).clamp(-torch.inf, 0)152        log_1_min_alpha_t = torch.log(sig)[None,:,None] + log_1_min_alpha_t153 154        # alpha_t * E[xt] + (1 - alpha_t) 1 / K155        log_probs = log_add_exp(156            log_x_t + log_alpha_t,157            log_1_min_alpha_t - np.log(self.num_classes)158        )159        return log_probs160 161    def q_pred(self, log_x_start: Tensor, t) -> Tensor:162        """ Compute q(x_t | x_0) = C(x_t | bar{alpha}_t * x_0 + (1 - bar{alpha}_t)/K ) in log domain,163        given `log_x_start` of log probs of x_0.164        """165        dt = log_x_start.dtype166        log_cumprod_alpha_t = extract(self.log_cumprod_alpha, t, log_x_start.shape).to(dt)167        log_1_min_cumprod_alpha = extract(self.log_1_min_cumprod_alpha, t, log_x_start.shape).to(dt)168        169        log_probs = log_add_exp(170            log_x_start + log_cumprod_alpha_t,171            log_1_min_cumprod_alpha - np.log(self.num_classes)172        )173        174        return log_probs175 176    def q_posterior(self, log_x_start, log_x_t, t):177        """ Compute  `q(xt-1 | xt, x0) = q(xt | xt-1, x0) * q(xt-1 | x0) / q(xt | x0)`178        where q(xt | xt-1, x0) = q(xt | xt-1).179        """180        # q(xt-1 | xt, x0) = q(xt | xt-1, x0) * q(xt-1 | x0) / q(xt | x0)181        # where q(xt | xt-1, x0) = q(xt | xt-1).182 183        t_minus_1 = t - 1184        # Remove negative values, will not be used anyway for final decoder185        t_minus_1 = torch.where(t_minus_1 < 0, torch.zeros_like(t_minus_1), t_minus_1)186        log_EV_qxtmin_x0 = self.q_pred(log_x_start, t_minus_1) # log( q(x_{t-1} | x_0) )187        # if t == 0, then log( q(x_0 | x_0) ) = log( one_hot(x_0) ), not even random at that point.188        # so, where t == 0 189        num_axes = (1,) * (len(log_x_start.size()) - 1) 190        t_broadcast = t.view(-1, *num_axes) * torch.ones_like(log_x_start) # broadcast to non-batch axes191        log_EV_qxtmin_x0 = torch.where(t_broadcast == 0, log_x_start, log_EV_qxtmin_x0) 192        # where it is zero, replace193        # with log one-hot encoding of x0.194 195        # Note: _NOT_ x_tmin1, which is how the formula is typically used!!!196        # Not very easy to see why this is true. But it is :)197        # log_EV_qxtmin_x0 ~ q(x_{t-1} | x_0) 198        # q_pred_one_timestep(log_x_t, t) ~ q(x_t | x_{t-1}) (which due to symmetry can be computed using x_t)199        unnormed_logprobs = log_EV_qxtmin_x0 + self.q_pred_one_timestep(log_x_t, t) # numerator of bayes200        201        # approximate denominator with just a normalizing sum.202        log_EV_xtmin_given_xt_given_xstart = \203            unnormed_logprobs \204            - torch.logsumexp(unnormed_logprobs, dim=-1, keepdim=True)205        206        return log_EV_xtmin_given_xt_given_xstart207 208    def p_pred(self, log_x_t, t, log_x0_pred):209        """ Predict `p(x_{t-1} | x_t)` using `q(xt-1 | xt, hat{x0})`, where `hat{x0}` is given by210        log probabilities from model as `log_x0_pred` (bs, ...., K) and x_t is given by211        `log_x_t` of shape `(bs, ..., K)`212        """213        # log_x_recon = self.predict_start(log_x, t=t) # model itself predicts x_0214        # log_x0_pred215        log_model_pred = self.q_posterior(216            log_x_start=log_x0_pred, log_x_t=log_x_t, t=t)217        return log_model_pred218 219    def log_sample_categorical(self, logprobs: Tensor, dim=-1) -> Tensor:220        """ Sample from categorical `logprobs` (bs, ..., probs), where position of probs is specified221        by `dim`.222 223        Returns sampled long indices of shape `(bs, ...)`224        """225        uniform = torch.rand_like(logprobs)226        gumbel_noise = -torch.log( (-torch.log(uniform.clamp_(min=MIN_LOG_ARG)) ).clamp_(min=MIN_LOG_ARG))227        sample = (gumbel_noise + logprobs).argmax(dim=dim)228        return sample229 230    def q_sample(self, log_x_start, t):231        """ Draw `x_t` ~ q(x_t | x_0) . `log_x_start` is of shape `(bs, ..., K)`, returns result of same shape """232        log_EV_qxt_x0 = self.q_pred(log_x_start, t)233        sample = self.log_sample_categorical(log_EV_qxt_x0)234        # log_sample = index_to_log_onehot(sample, self.num_classes)235 236        return sample #log_sample237 238    def compute_Lt(self, log_x_start: Tensor, log_x_t: Tensor, log_x0_pred: Tensor, t, 239                   detach_mean=False, include_kl_prior=True):240        """ Get loss given one-hot log x_0, one-hot log x_t, t, and model prediction `log_x0_pred`.241        Parameters:242            - `log_x_start`: ground-truth input x0, converted to log one-hot (bs, ..., K)243            - `log_x_t`: sampled noisy input at `x_t`, converted to log one-hot (bs, ..., K)244            - `t`: diffusion timestep (bs,)245            - `log_x0_pred`: model prediction of log probabilities of x0, i.e. hat{x0}.246            - `include_kl_prior`: add last two terms to model loss (does not change optimization problem).247        """248        dtype = log_x_start.dtype249        log_true_prob = self.q_posterior(250            log_x_start=log_x_start, log_x_t=log_x_t, t=t)251 252        log_model_prob = self.p_pred(log_x_t=log_x_t, t=t, log_x0_pred=log_x0_pred)253 254        if detach_mean:255            log_model_prob = log_model_prob.detach()256 257        kl = self.multinomial_kl(log_true_prob, log_model_prob)258        kl = sum_except_batch(kl)259 260        # Add L_0, -log(p(x_0 | x_1))261        decoder_nll = - (log_x_start.exp() * log_model_prob).sum(dim=-1)262        decoder_nll = sum_except_batch(decoder_nll)263 264        mask = (t == torch.zeros_like(t)).to(dtype)265        loss = mask * decoder_nll + (1. - mask) * kl # only add L0 if t == 0.266 267        if include_kl_prior:268            pt = torch.ones_like(t, dtype=dtype)269            kl_prior = self.kl_prior(log_x_start)270            loss = (kl) + kl_prior271 272        return loss273 274    def kl_prior(self, log_x_start: Tensor) -> Tensor:275        """ This function computes -H_{q}(x_T | x_0)+H_{p}(x_T), which 276        by some math (see wiki for KL div relation to conditional entropy).277        So KL(q(x_T | x_0) || 1/K) = -H_{q}(x_T | x_0)+H_{p}(x_T) for categorical distribution.278 279        Given `log_x_start` (bs, ..., probs), return KL prior of shape (bs,)280        """281        b = log_x_start.size(0)282        device = log_x_start.device283        ones = torch.ones(b, device=device, dtype=torch.long)284 285        log_qxT_prob = self.q_pred(log_x_start, t=(self.num_timesteps - 1) * ones) # q(x_T | x_0)286        log_half_prob = -torch.log(self.num_classes * torch.ones_like(log_qxT_prob)) # log(1/K), broadcast to q(x_T|x_0) shape287 288        kl_prior = self.multinomial_kl(log_qxT_prob, log_half_prob)289        return sum_except_batch(kl_prior)290 291 292def index2logit(x: Tensor, vocab_size: int, dtype=torch.float32):293    x = F.one_hot(x, num_classes=vocab_size).to(dtype)294    x = x * (vocab_size/(vocab_size - 1)) - 1/(vocab_size - 1)295    return x296 297 298# ------------------------------299# Functions adapted from the full300 301 302@dataclass303class DSH():304    # Diffusion Sampling Hyperparameters [DSH] (Section 4)305    jump_len: int = 1 # j in RePaint paper [default 10] (Section 4.1)306    jump_n_sample: int = 1 # r in RePaint paper [default 10] (Section 4.1)307    last_greedy: bool = False # whether to not sample at t=0, but take argmax prediction. [default False]308    x_0_temp: float = 1.0 # reweight temp for model prediction of x0309    guidance_w: float = 1.0 # classifier free guidance weight [default 1.5] (Section 4.3)310    enable_kevin_scaled_inference: bool = True # sequentially progressive diffusion [default True] (Section 4.2)311    T_override: Union[None, int] = None # allow variable transcription sizes during inference (Section 4.4)312 313    deep_clone: bool = False # whether to do deep clone. 314    q0_override_steps: int = 0 # number of steps that we allow overriding the input quant level 0 inputs.315    progress: bool = False # whether to show progress bar316 317 318def get_schedule(t_T, jump_len=10, jump_n_sample=10):319    jumps = {}320    for j in range(0, t_T - jump_len, jump_len):321        jumps[j] = jump_n_sample - 1322    t = t_T323    ts = []324    while t >= 1:325        t = t-1326        ts.append(t)327        if jumps.get(t, 0) > 0:328            jumps[t] = jumps[t] - 1329            for _ in range(jump_len):330                t = t + 1331                ts.append(t)332    ts.append(-1)333    return ts334 335 336def forward_diffusion(diff: MultinomialDiffusion, dtype, x, t, c=None, dsh=DSH):337    """Simple forward diffusion process p"""338    log_x_t = index_to_log_onehot(x, diff.num_classes, dtype=dtype)339    if c is not None: x = diff.q_pred_one_timestep_scaled(log_x_t, t, c, dsh.jump_len)340    else: x = diff.q_pred_one_timestep(log_x_t, t)341    x = diff.log_sample_categorical(x)342    return x343 344 345def reverse_diffusion(diff: MultinomialDiffusion, model, batch, x_known=None, m=None, 346                      last_greedy=False, temperature=1.0, alphas=None, ensemble_size=1, dsh=DSH):347    """Reverse diffusion process q: predict x_{t-1} given x, t, x_known, m. Optionally do not sample model output348    for t=0, but rather use the greedy argmax with `last_greedy`.349    """350    x = batch[4]351    t = batch[-1]352    if x_known is None: x_known = torch.zeros_like(x)353    if m is None: m = torch.zeros_like(x)354 355    # Equation 8b356    # for b in batch:357        # print(f"{b.shape}: {b}")358    x_0_pred = model(*batch) # (bs, seq_len, logit_dim, n_quant)359    x_0_pred = x_0_pred.permute(0, 1, 3, 2) # (bs, seq_len, n_quant, dim)360 361    if dsh.guidance_w != 1:362        uncond_x_0_pred = model(*(c.clone() if c is not None else None for c in batch), drop_cond=True)363        uncond_x_0_pred = uncond_x_0_pred.permute(0, 1, 3, 2)364        x_0_pred = dsh.guidance_w*x_0_pred + (1-dsh.guidance_w)*uncond_x_0_pred365 366    x_0_pred = x_0_pred / temperature367    log_x_0_pred = F.log_softmax(x_0_pred, dim=-1)368    log_x_t = index_to_log_onehot(x, diff.num_classes, dtype=x_0_pred.dtype)369 370    # print("PRE: ", log_x_t.shape, t.shape, log_x_0_pred.shape)371    log_model_pred = diff.p_pred(log_x_t, t, log_x_0_pred) # p(x_{t-1} | x_{t})372    373    a_t = alphas[t[0]] if alphas is not None else 0374    mat = torch.eye(ensemble_size, device=x.device)*(1-a_t)375    mat += 1/ensemble_size * a_t376    mat = torch.block_diag(*([mat]*(x.shape[0]//ensemble_size)))377    log_model_pred = ( (mat[..., None, None] ).log().to(x.dtype) + log_model_pred[None])378    log_model_pred = torch.logsumexp(log_model_pred, dim=1)379    380    if (t==0).all() and last_greedy: # Do not sample at t=0381        x_tm1_unknown = log_model_pred.argmax(dim=-1)382    else:383        x_tm1_unknown = diff.log_sample_categorical(log_model_pred)384    385    # Equation 8a386    x_known_log = index_to_log_onehot(x_known, diff.num_classes, dtype=x_0_pred.dtype)387    if (t==0).all(): # Do not sample at t=0388        x_tm1_known = x_known389    else:390        x_tm1_known = diff.q_sample(x_known_log, t)391    392    # Equation 8c393    x_tm1 = x_tm1_known * m.long() + x_tm1_unknown * (1 - m.long())394    return x_tm1, x_0_pred395 396 397 398@torch.inference_mode()399def perform_simple_inference(model: torch.nn.Module, batch: tuple, diff: MultinomialDiffusion, T, dtype=torch.float16,400                             retain_quant0: bool = True, dsh=DSH):401    """ If `retain_quant0`, then do not sample quant0 in each forward or reverse diffusion step. """402 403    # (bs=1, N), (bs, seq_len2, 8), (bs,)404    c_text, c_codes, c_text_lengths, c_codes_lengths, x, x_padding_mask = batch405        406    device = c_text.device407    bs = c_text.shape[0]408    x_quant0 = x[..., 0].clone() # (bs, seq_len) 0th quant level409    x = torch.randint(0, diff.num_classes, x.shape, dtype=x.dtype, device=device)410    # CRITICAL LINE: override quantization level 0 with provided quant0 level.411    x[..., 0] = x_quant0 412 413    # RePaint paper resample scheduling414    times = get_schedule(T, jump_n_sample=dsh.jump_n_sample, jump_len=dsh.jump_len)415 416    x_known = torch.zeros_like(x)417    x_known[..., 0] = x[..., 0] # override L0 codes418    m = torch.zeros_like(x).bool()419    # (bs, seq_len, 8)420    m[..., 0] = True421 422    offset = 0423    if dsh.deep_clone:424        print(f"Note: using deep clone. Assuming input `c_phones` is concatenated prompt and output phones.",425              "Also assuming no padded indices in `c_codes`.")426        prompt = c_codes427        x = torch.cat((prompt, x), dim=1) # (bs=1, sl1 + sl2, 8)428        x_known = torch.cat((prompt, x_known), dim=1)429        x_padding_mask = torch.cat((430            torch.zeros(x_padding_mask.shape[0], c_codes_lengths[0], dtype=torch.bool, device=x_padding_mask.device), 431            x_padding_mask), dim=-1432        )433        # (bs=1, :up to prompt duration, all 8 codebooks) = True/masked.434        m = torch.cat((torch.ones_like(prompt), m), dim=1)435        x_quant0 = torch.cat((prompt[..., 0], x_quant0), dim=-1)436        offset = c_codes_lengths[0]437 438        print(f"New x: {x.shape} | new x_known: {x_known.shape} . Base prompt: {prompt.shape}. New padding mask: {x_padding_mask.shape} | m shape: {m.shape}")439 440    c = 0 # sequentially progressive diffusion offset (Section 4.2)441 442    # ensemble bs (not in paper)443    alphas = torch.linspace(1, 0, T).to(device)444 445    pb = zip(times[:-1], times[1:])446    if dsh.progress:447        from fastprogress import progress_bar448        pb = progress_bar(pb, total=len(times)-1)449 450    # See RePaint paper algorithm451    for t_last, t_cur in pb:452 453        t = torch.ones((bs,), dtype=torch.long, device=x.device) * (t_last)454        if t_cur < t_last:455            if c > dsh.jump_n_sample:456                c = 0457            c += 1/dsh.jump_len458 459            # Reverse diffusion: q460            cbatch = (c_text, c_codes, c_text_lengths, c_codes_lengths, x, x_padding_mask, t) 461            x, x_0_pred = reverse_diffusion(diff, model, cbatch, x_known, m, temperature=dsh.x_0_temp, alphas=alphas, ensemble_size=1, dsh=dsh)462        else:463            # Forward diffusion: p464            if dsh.enable_kevin_scaled_inference: x = forward_diffusion(diff, dtype, x, t, c=c, dsh=dsh)465            else: x = forward_diffusion(diff, dtype, x, t, c=None, dsh=dsh)466 467        if retain_quant0 and dsh.q0_override_steps < t_last:468            x[..., 0] = x_quant0469 470    # crop offset:471    x = x[:, offset:]472    return x473