xinference 1.10.0__py3-none-any.whl → 1.11.0__py3-none-any.whl

xinference/thirdparty/indextts/utils/maskgct/models/codec/facodec/modules/quantize.py

@@ -0,0 +1,741 @@
+# Copyright (c) 2023 Amphion.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from modules.dac.nn.quantize import ResidualVectorQuantize
+from torch import nn
+from .wavenet import WN
+from .style_encoder import StyleEncoder
+from .gradient_reversal import GradientReversal
+import torch
+import torchaudio
+import torchaudio.functional as audio_F
+import numpy as np
+from ..alias_free_torch import *
+from torch.nn.utils import weight_norm
+from torch import nn, sin, pow
+from einops.layers.torch import Rearrange
+from modules.dac.model.encodec import SConv1d
+
+
+def init_weights(m):
+    if isinstance(m, nn.Conv1d):
+        nn.init.trunc_normal_(m.weight, std=0.02)
+        nn.init.constant_(m.bias, 0)
+
+
+def WNConv1d(*args, **kwargs):
+    return weight_norm(nn.Conv1d(*args, **kwargs))
+
+
+def WNConvTranspose1d(*args, **kwargs):
+    return weight_norm(nn.ConvTranspose1d(*args, **kwargs))
+
+
+class SnakeBeta(nn.Module):
+    """
+    A modified Snake function which uses separate parameters for the magnitude of the periodic components
+    Shape:
+        - Input: (B, C, T)
+        - Output: (B, C, T), same shape as the input
+    Parameters:
+        - alpha - trainable parameter that controls frequency
+        - beta - trainable parameter that controls magnitude
+    References:
+        - This activation function is a modified version based on this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda:
+        https://arxiv.org/abs/2006.08195
+    Examples:
+        >>> a1 = snakebeta(256)
+        >>> x = torch.randn(256)
+        >>> x = a1(x)
+    """
+
+    def __init__(
+        self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False
+    ):
+        """
+        Initialization.
+        INPUT:
+            - in_features: shape of the input
+            - alpha - trainable parameter that controls frequency
+            - beta - trainable parameter that controls magnitude
+            alpha is initialized to 1 by default, higher values = higher-frequency.
+            beta is initialized to 1 by default, higher values = higher-magnitude.
+            alpha will be trained along with the rest of your model.
+        """
+        super(SnakeBeta, self).__init__()
+        self.in_features = in_features
+
+        # initialize alpha
+        self.alpha_logscale = alpha_logscale
+        if self.alpha_logscale:  # log scale alphas initialized to zeros
+            self.alpha = nn.Parameter(torch.zeros(in_features) * alpha)
+            self.beta = nn.Parameter(torch.zeros(in_features) * alpha)
+        else:  # linear scale alphas initialized to ones
+            self.alpha = nn.Parameter(torch.ones(in_features) * alpha)
+            self.beta = nn.Parameter(torch.ones(in_features) * alpha)
+
+        self.alpha.requires_grad = alpha_trainable
+        self.beta.requires_grad = alpha_trainable
+
+        self.no_div_by_zero = 0.000000001
+
+    def forward(self, x):
+        """
+        Forward pass of the function.
+        Applies the function to the input elementwise.
+        SnakeBeta := x + 1/b * sin^2 (xa)
+        """
+        alpha = self.alpha.unsqueeze(0).unsqueeze(-1)  # line up with x to [B, C, T]
+        beta = self.beta.unsqueeze(0).unsqueeze(-1)
+        if self.alpha_logscale:
+            alpha = torch.exp(alpha)
+            beta = torch.exp(beta)
+        x = x + (1.0 / (beta + self.no_div_by_zero)) * pow(sin(x * alpha), 2)
+
+        return x
+
+
+class ResidualUnit(nn.Module):
+    def __init__(self, dim: int = 16, dilation: int = 1):
+        super().__init__()
+        pad = ((7 - 1) * dilation) // 2
+        self.block = nn.Sequential(
+            Activation1d(activation=SnakeBeta(dim, alpha_logscale=True)),
+            WNConv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad),
+            Activation1d(activation=SnakeBeta(dim, alpha_logscale=True)),
+            WNConv1d(dim, dim, kernel_size=1),
+        )
+
+    def forward(self, x):
+        return x + self.block(x)
+
+
+class CNNLSTM(nn.Module):
+    def __init__(self, indim, outdim, head, global_pred=False):
+        super().__init__()
+        self.global_pred = global_pred
+        self.model = nn.Sequential(
+            ResidualUnit(indim, dilation=1),
+            ResidualUnit(indim, dilation=2),
+            ResidualUnit(indim, dilation=3),
+            Activation1d(activation=SnakeBeta(indim, alpha_logscale=True)),
+            Rearrange("b c t -> b t c"),
+        )
+        self.heads = nn.ModuleList([nn.Linear(indim, outdim) for i in range(head)])
+
+    def forward(self, x):
+        # x: [B, C, T]
+        x = self.model(x)
+        if self.global_pred:
+            x = torch.mean(x, dim=1, keepdim=False)
+        outs = [head(x) for head in self.heads]
+        return outs
+
+
+def sequence_mask(length, max_length=None):
+    if max_length is None:
+        max_length = length.max()
+    x = torch.arange(max_length, dtype=length.dtype, device=length.device)
+    return x.unsqueeze(0) < length.unsqueeze(1)
+
+
+class MFCC(nn.Module):
+    def __init__(self, n_mfcc=40, n_mels=80):
+        super(MFCC, self).__init__()
+        self.n_mfcc = n_mfcc
+        self.n_mels = n_mels
+        self.norm = "ortho"
+        dct_mat = audio_F.create_dct(self.n_mfcc, self.n_mels, self.norm)
+        self.register_buffer("dct_mat", dct_mat)
+
+    def forward(self, mel_specgram):
+        if len(mel_specgram.shape) == 2:
+            mel_specgram = mel_specgram.unsqueeze(0)
+            unsqueezed = True
+        else:
+            unsqueezed = False
+        # (channel, n_mels, time).tranpose(...) dot (n_mels, n_mfcc)
+        # -> (channel, time, n_mfcc).tranpose(...)
+        mfcc = torch.matmul(mel_specgram.transpose(1, 2), self.dct_mat).transpose(1, 2)
+
+        # unpack batch
+        if unsqueezed:
+            mfcc = mfcc.squeeze(0)
+        return mfcc
+
+
+class FAquantizer(nn.Module):
+    def __init__(
+        self,
+        in_dim=1024,
+        n_p_codebooks=1,
+        n_c_codebooks=2,
+        n_t_codebooks=2,
+        n_r_codebooks=3,
+        codebook_size=1024,
+        codebook_dim=8,
+        quantizer_dropout=0.5,
+        causal=False,
+        separate_prosody_encoder=False,
+        timbre_norm=False,
+    ):
+        super(FAquantizer, self).__init__()
+        conv1d_type = SConv1d  # if causal else nn.Conv1d
+        self.prosody_quantizer = ResidualVectorQuantize(
+            input_dim=in_dim,
+            n_codebooks=n_p_codebooks,
+            codebook_size=codebook_size,
+            codebook_dim=codebook_dim,
+            quantizer_dropout=quantizer_dropout,
+        )
+
+        self.content_quantizer = ResidualVectorQuantize(
+            input_dim=in_dim,
+            n_codebooks=n_c_codebooks,
+            codebook_size=codebook_size,
+            codebook_dim=codebook_dim,
+            quantizer_dropout=quantizer_dropout,
+        )
+
+        if not timbre_norm:
+            self.timbre_quantizer = ResidualVectorQuantize(
+                input_dim=in_dim,
+                n_codebooks=n_t_codebooks,
+                codebook_size=codebook_size,
+                codebook_dim=codebook_dim,
+                quantizer_dropout=quantizer_dropout,
+            )
+        else:
+            self.timbre_encoder = StyleEncoder(
+                in_dim=80, hidden_dim=512, out_dim=in_dim
+            )
+            self.timbre_linear = nn.Linear(1024, 1024 * 2)
+            self.timbre_linear.bias.data[:1024] = 1
+            self.timbre_linear.bias.data[1024:] = 0
+            self.timbre_norm = nn.LayerNorm(1024, elementwise_affine=False)
+
+        self.residual_quantizer = ResidualVectorQuantize(
+            input_dim=in_dim,
+            n_codebooks=n_r_codebooks,
+            codebook_size=codebook_size,
+            codebook_dim=codebook_dim,
+            quantizer_dropout=quantizer_dropout,
+        )
+
+        if separate_prosody_encoder:
+            self.melspec_linear = conv1d_type(
+                in_channels=20, out_channels=256, kernel_size=1, causal=causal
+            )
+            self.melspec_encoder = WN(
+                hidden_channels=256,
+                kernel_size=5,
+                dilation_rate=1,
+                n_layers=8,
+                gin_channels=0,
+                p_dropout=0.2,
+                causal=causal,
+            )
+            self.melspec_linear2 = conv1d_type(
+                in_channels=256, out_channels=1024, kernel_size=1, causal=causal
+            )
+        else:
+            pass
+        self.separate_prosody_encoder = separate_prosody_encoder
+
+        self.prob_random_mask_residual = 0.75
+
+        SPECT_PARAMS = {
+            "n_fft": 2048,
+            "win_length": 1200,
+            "hop_length": 300,
+        }
+        MEL_PARAMS = {
+            "n_mels": 80,
+        }
+
+        self.to_mel = torchaudio.transforms.MelSpectrogram(
+            n_mels=MEL_PARAMS["n_mels"], sample_rate=24000, **SPECT_PARAMS
+        )
+        self.mel_mean, self.mel_std = -4, 4
+        self.frame_rate = 24000 / 300
+        self.hop_length = 300
+
+        self.is_timbre_norm = timbre_norm
+        if timbre_norm:
+            self.forward = self.forward_v2
+
+    def preprocess(self, wave_tensor, n_bins=20):
+        mel_tensor = self.to_mel(wave_tensor.squeeze(1))
+        mel_tensor = (torch.log(1e-5 + mel_tensor) - self.mel_mean) / self.mel_std
+        return mel_tensor[:, :n_bins, : int(wave_tensor.size(-1) / self.hop_length)]
+
+    @torch.no_grad()
+    def decode(self, codes):
+        code_c, code_p, code_t = codes.split([1, 1, 2], dim=1)
+
+        z_c = self.content_quantizer.from_codes(code_c)[0]
+        z_p = self.prosody_quantizer.from_codes(code_p)[0]
+        z_t = self.timbre_quantizer.from_codes(code_t)[0]
+
+        z = z_c + z_p + z_t
+
+        return z, [z_c, z_p, z_t]
+
+    @torch.no_grad()
+    def encode(self, x, wave_segments, n_c=1):
+        outs = 0
+        if self.separate_prosody_encoder:
+            prosody_feature = self.preprocess(wave_segments)
+
+            f0_input = prosody_feature  # (B, T, 20)
+            f0_input = self.melspec_linear(f0_input)
+            f0_input = self.melspec_encoder(
+                f0_input,
+                torch.ones(f0_input.shape[0], 1, f0_input.shape[2])
+                .to(f0_input.device)
+                .bool(),
+            )
+            f0_input = self.melspec_linear2(f0_input)
+
+            common_min_size = min(f0_input.size(2), x.size(2))
+            f0_input = f0_input[:, :, :common_min_size]
+
+            x = x[:, :, :common_min_size]
+
+            (
+                z_p,
+                codes_p,
+                latents_p,
+                commitment_loss_p,
+                codebook_loss_p,
+            ) = self.prosody_quantizer(f0_input, 1)
+            outs += z_p.detach()
+        else:
+            (
+                z_p,
+                codes_p,
+                latents_p,
+                commitment_loss_p,
+                codebook_loss_p,
+            ) = self.prosody_quantizer(x, 1)
+            outs += z_p.detach()
+
+        (
+            z_c,
+            codes_c,
+            latents_c,
+            commitment_loss_c,
+            codebook_loss_c,
+        ) = self.content_quantizer(x, n_c)
+        outs += z_c.detach()
+
+        timbre_residual_feature = x - z_p.detach() - z_c.detach()
+
+        (
+            z_t,
+            codes_t,
+            latents_t,
+            commitment_loss_t,
+            codebook_loss_t,
+        ) = self.timbre_quantizer(timbre_residual_feature, 2)
+        outs += z_t  # we should not detach timbre
+
+        residual_feature = timbre_residual_feature - z_t
+
+        (
+            z_r,
+            codes_r,
+            latents_r,
+            commitment_loss_r,
+            codebook_loss_r,
+        ) = self.residual_quantizer(residual_feature, 3)
+
+        return [codes_c, codes_p, codes_t, codes_r], [z_c, z_p, z_t, z_r]
+
+    def forward(
+        self, x, wave_segments, noise_added_flags, recon_noisy_flags, n_c=2, n_t=2
+    ):
+        # timbre = self.timbre_encoder(mels, sequence_mask(mel_lens, mels.size(-1)).unsqueeze(1))
+        # timbre = self.timbre_encoder(mel_segments, torch.ones(mel_segments.size(0), 1, mel_segments.size(2)).bool().to(mel_segments.device))
+        outs = 0
+        if self.separate_prosody_encoder:
+            prosody_feature = self.preprocess(wave_segments)
+
+            f0_input = prosody_feature  # (B, T, 20)
+            f0_input = self.melspec_linear(f0_input)
+            f0_input = self.melspec_encoder(
+                f0_input,
+                torch.ones(f0_input.shape[0], 1, f0_input.shape[2])
+                .to(f0_input.device)
+                .bool(),
+            )
+            f0_input = self.melspec_linear2(f0_input)
+
+            common_min_size = min(f0_input.size(2), x.size(2))
+            f0_input = f0_input[:, :, :common_min_size]
+
+            x = x[:, :, :common_min_size]
+
+            (
+                z_p,
+                codes_p,
+                latents_p,
+                commitment_loss_p,
+                codebook_loss_p,
+            ) = self.prosody_quantizer(f0_input, 1)
+            outs += z_p.detach()
+        else:
+            (
+                z_p,
+                codes_p,
+                latents_p,
+                commitment_loss_p,
+                codebook_loss_p,
+            ) = self.prosody_quantizer(x, 1)
+            outs += z_p.detach()
+
+        (
+            z_c,
+            codes_c,
+            latents_c,
+            commitment_loss_c,
+            codebook_loss_c,
+        ) = self.content_quantizer(x, n_c)
+        outs += z_c.detach()
+
+        timbre_residual_feature = x - z_p.detach() - z_c.detach()
+
+        (
+            z_t,
+            codes_t,
+            latents_t,
+            commitment_loss_t,
+            codebook_loss_t,
+        ) = self.timbre_quantizer(timbre_residual_feature, n_t)
+        outs += z_t  # we should not detach timbre
+
+        residual_feature = timbre_residual_feature - z_t
+
+        (
+            z_r,
+            codes_r,
+            latents_r,
+            commitment_loss_r,
+            codebook_loss_r,
+        ) = self.residual_quantizer(residual_feature, 3)
+
+        bsz = z_r.shape[0]
+        res_mask = np.random.choice(
+            [0, 1],
+            size=bsz,
+            p=[
+                self.prob_random_mask_residual,
+                1 - self.prob_random_mask_residual,
+            ],
+        )
+        res_mask = torch.from_numpy(res_mask).unsqueeze(1).unsqueeze(1)  # (B, 1, 1)
+        res_mask = res_mask.to(device=z_r.device, dtype=z_r.dtype)
+        noise_must_on = noise_added_flags * recon_noisy_flags
+        noise_must_off = noise_added_flags * (~recon_noisy_flags)
+        res_mask[noise_must_on] = 1
+        res_mask[noise_must_off] = 0
+
+        outs += z_r * res_mask
+
+        quantized = [z_p, z_c, z_t, z_r]
+        commitment_losses = (
+            commitment_loss_p
+            + commitment_loss_c
+            + commitment_loss_t
+            + commitment_loss_r
+        )
+        codebook_losses = (
+            codebook_loss_p + codebook_loss_c + codebook_loss_t + codebook_loss_r
+        )
+
+        return outs, quantized, commitment_losses, codebook_losses
+
+    def forward_v2(
+        self,
+        x,
+        wave_segments,
+        n_c=1,
+        n_t=2,
+        full_waves=None,
+        wave_lens=None,
+        return_codes=False,
+    ):
+        # timbre = self.timbre_encoder(x, sequence_mask(mel_lens, mels.size(-1)).unsqueeze(1))
+        if full_waves is None:
+            mel = self.preprocess(wave_segments, n_bins=80)
+            timbre = self.timbre_encoder(
+                mel, torch.ones(mel.size(0), 1, mel.size(2)).bool().to(mel.device)
+            )
+        else:
+            mel = self.preprocess(full_waves, n_bins=80)
+            timbre = self.timbre_encoder(
+                mel,
+                sequence_mask(wave_lens // self.hop_length, mel.size(-1)).unsqueeze(1),
+            )
+        outs = 0
+        if self.separate_prosody_encoder:
+            prosody_feature = self.preprocess(wave_segments)
+
+            f0_input = prosody_feature  # (B, T, 20)
+            f0_input = self.melspec_linear(f0_input)
+            f0_input = self.melspec_encoder(
+                f0_input,
+                torch.ones(f0_input.shape[0], 1, f0_input.shape[2])
+                .to(f0_input.device)
+                .bool(),
+            )
+            f0_input = self.melspec_linear2(f0_input)
+
+            common_min_size = min(f0_input.size(2), x.size(2))
+            f0_input = f0_input[:, :, :common_min_size]
+
+            x = x[:, :, :common_min_size]
+
+            (
+                z_p,
+                codes_p,
+                latents_p,
+                commitment_loss_p,
+                codebook_loss_p,
+            ) = self.prosody_quantizer(f0_input, 1)
+            outs += z_p.detach()
+        else:
+            (
+                z_p,
+                codes_p,
+                latents_p,
+                commitment_loss_p,
+                codebook_loss_p,
+            ) = self.prosody_quantizer(x, 1)
+            outs += z_p.detach()
+
+        (
+            z_c,
+            codes_c,
+            latents_c,
+            commitment_loss_c,
+            codebook_loss_c,
+        ) = self.content_quantizer(x, n_c)
+        outs += z_c.detach()
+
+        residual_feature = x - z_p.detach() - z_c.detach()
+
+        (
+            z_r,
+            codes_r,
+            latents_r,
+            commitment_loss_r,
+            codebook_loss_r,
+        ) = self.residual_quantizer(residual_feature, 3)
+
+        bsz = z_r.shape[0]
+        res_mask = np.random.choice(
+            [0, 1],
+            size=bsz,
+            p=[
+                self.prob_random_mask_residual,
+                1 - self.prob_random_mask_residual,
+            ],
+        )
+        res_mask = torch.from_numpy(res_mask).unsqueeze(1).unsqueeze(1)  # (B, 1, 1)
+        res_mask = res_mask.to(device=z_r.device, dtype=z_r.dtype)
+
+        if not self.training:
+            res_mask = torch.ones_like(res_mask)
+        outs += z_r * res_mask
+
+        quantized = [z_p, z_c, z_r]
+        codes = [codes_p, codes_c, codes_r]
+        commitment_losses = commitment_loss_p + commitment_loss_c + commitment_loss_r
+        codebook_losses = codebook_loss_p + codebook_loss_c + codebook_loss_r
+
+        style = self.timbre_linear(timbre).unsqueeze(2)  # (B, 2d, 1)
+        gamma, beta = style.chunk(2, 1)  # (B, d, 1)
+        outs = outs.transpose(1, 2)
+        outs = self.timbre_norm(outs)
+        outs = outs.transpose(1, 2)
+        outs = outs * gamma + beta
+
+        if return_codes:
+            return outs, quantized, commitment_losses, codebook_losses, timbre, codes
+        else:
+            return outs, quantized, commitment_losses, codebook_losses, timbre
+
+    def voice_conversion(self, z, ref_wave):
+        ref_mel = self.preprocess(ref_wave, n_bins=80)
+        ref_timbre = self.timbre_encoder(
+            ref_mel,
+            sequence_mask(
+                torch.LongTensor([ref_wave.size(-1)]).to(z.device) // self.hop_length,
+                ref_mel.size(-1),
+            ).unsqueeze(1),
+        )
+        style = self.timbre_linear(ref_timbre).unsqueeze(2)  # (B, 2d, 1)
+        gamma, beta = style.chunk(2, 1)  # (B, d, 1)
+        outs = z.transpose(1, 2)
+        outs = self.timbre_norm(outs)
+        outs = outs.transpose(1, 2)
+        outs = outs * gamma + beta
+
+        return outs
+
+
+class FApredictors(nn.Module):
+    def __init__(
+        self,
+        in_dim=1024,
+        use_gr_content_f0=False,
+        use_gr_prosody_phone=False,
+        use_gr_residual_f0=False,
+        use_gr_residual_phone=False,
+        use_gr_timbre_content=True,
+        use_gr_timbre_prosody=True,
+        use_gr_x_timbre=False,
+        norm_f0=True,
+        timbre_norm=False,
+        use_gr_content_global_f0=False,
+    ):
+        super(FApredictors, self).__init__()
+        self.f0_predictor = CNNLSTM(in_dim, 1, 2)
+        self.phone_predictor = CNNLSTM(in_dim, 1024, 1)
+        if timbre_norm:
+            self.timbre_predictor = nn.Linear(in_dim, 20000)
+        else:
+            self.timbre_predictor = CNNLSTM(in_dim, 20000, 1, global_pred=True)
+
+        self.use_gr_content_f0 = use_gr_content_f0
+        self.use_gr_prosody_phone = use_gr_prosody_phone
+        self.use_gr_residual_f0 = use_gr_residual_f0
+        self.use_gr_residual_phone = use_gr_residual_phone
+        self.use_gr_timbre_content = use_gr_timbre_content
+        self.use_gr_timbre_prosody = use_gr_timbre_prosody
+        self.use_gr_x_timbre = use_gr_x_timbre
+
+        self.rev_f0_predictor = nn.Sequential(
+            GradientReversal(alpha=1.0), CNNLSTM(in_dim, 1, 2)
+        )
+        self.rev_content_predictor = nn.Sequential(
+            GradientReversal(alpha=1.0), CNNLSTM(in_dim, 1024, 1)
+        )
+        self.rev_timbre_predictor = nn.Sequential(
+            GradientReversal(alpha=1.0), CNNLSTM(in_dim, 20000, 1, global_pred=True)
+        )
+
+        self.norm_f0 = norm_f0
+        self.timbre_norm = timbre_norm
+        if timbre_norm:
+            self.forward = self.forward_v2
+            self.global_f0_predictor = nn.Linear(in_dim, 1)
+
+        self.use_gr_content_global_f0 = use_gr_content_global_f0
+        if use_gr_content_global_f0:
+            self.rev_global_f0_predictor = nn.Sequential(
+                GradientReversal(alpha=1.0), CNNLSTM(in_dim, 1, 1, global_pred=True)
+            )
+
+    def forward(self, quantized):
+        prosody_latent = quantized[0]
+        content_latent = quantized[1]
+        timbre_latent = quantized[2]
+        residual_latent = quantized[3]
+        content_pred = self.phone_predictor(content_latent)[0]
+
+        if self.norm_f0:
+            spk_pred = self.timbre_predictor(timbre_latent)[0]
+            f0_pred, uv_pred = self.f0_predictor(prosody_latent)
+        else:
+            spk_pred = self.timbre_predictor(timbre_latent + prosody_latent)[0]
+            f0_pred, uv_pred = self.f0_predictor(prosody_latent + timbre_latent)
+
+        prosody_rev_latent = torch.zeros_like(quantized[0])
+        if self.use_gr_content_f0:
+            prosody_rev_latent += quantized[1]
+        if self.use_gr_timbre_prosody:
+            prosody_rev_latent += quantized[2]
+        if self.use_gr_residual_f0:
+            prosody_rev_latent += quantized[3]
+        rev_f0_pred, rev_uv_pred = self.rev_f0_predictor(prosody_rev_latent)
+
+        content_rev_latent = torch.zeros_like(quantized[1])
+        if self.use_gr_prosody_phone:
+            content_rev_latent += quantized[0]
+        if self.use_gr_timbre_content:
+            content_rev_latent += quantized[2]
+        if self.use_gr_residual_phone:
+            content_rev_latent += quantized[3]
+        rev_content_pred = self.rev_content_predictor(content_rev_latent)[0]
+
+        if self.norm_f0:
+            timbre_rev_latent = quantized[0] + quantized[1] + quantized[3]
+        else:
+            timbre_rev_latent = quantized[1] + quantized[3]
+        if self.use_gr_x_timbre:
+            x_spk_pred = self.rev_timbre_predictor(timbre_rev_latent)[0]
+        else:
+            x_spk_pred = None
+
+        preds = {
+            "f0": f0_pred,
+            "uv": uv_pred,
+            "content": content_pred,
+            "timbre": spk_pred,
+        }
+
+        rev_preds = {
+            "rev_f0": rev_f0_pred,
+            "rev_uv": rev_uv_pred,
+            "rev_content": rev_content_pred,
+            "x_timbre": x_spk_pred,
+        }
+        return preds, rev_preds
+
+    def forward_v2(self, quantized, timbre):
+        prosody_latent = quantized[0]
+        content_latent = quantized[1]
+        residual_latent = quantized[2]
+        content_pred = self.phone_predictor(content_latent)[0]
+
+        spk_pred = self.timbre_predictor(timbre)
+        f0_pred, uv_pred = self.f0_predictor(prosody_latent)
+
+        prosody_rev_latent = torch.zeros_like(prosody_latent)
+        if self.use_gr_content_f0:
+            prosody_rev_latent += content_latent
+        if self.use_gr_residual_f0:
+            prosody_rev_latent += residual_latent
+        rev_f0_pred, rev_uv_pred = self.rev_f0_predictor(prosody_rev_latent)
+
+        content_rev_latent = torch.zeros_like(content_latent)
+        if self.use_gr_prosody_phone:
+            content_rev_latent += prosody_latent
+        if self.use_gr_residual_phone:
+            content_rev_latent += residual_latent
+        rev_content_pred = self.rev_content_predictor(content_rev_latent)[0]
+
+        timbre_rev_latent = prosody_latent + content_latent + residual_latent
+        if self.use_gr_x_timbre:
+            x_spk_pred = self.rev_timbre_predictor(timbre_rev_latent)[0]
+        else:
+            x_spk_pred = None
+
+        preds = {
+            "f0": f0_pred,
+            "uv": uv_pred,
+            "content": content_pred,
+            "timbre": spk_pred,
+        }
+
+        rev_preds = {
+            "rev_f0": rev_f0_pred,
+            "rev_uv": rev_uv_pred,
+            "rev_content": rev_content_pred,
+            "x_timbre": x_spk_pred,
+        }
+        return preds, rev_preds