chore: import upstream snapshot with attribution

modules/aux_decoder/__init__.py

@@ -0,0 +1,70 @@
+import torch.nn
+from torch import nn
+
+from .convnext import ConvNeXtDecoder
+from utils import filter_kwargs
+
+AUX_DECODERS = {
+    'convnext': ConvNeXtDecoder
+}
+AUX_LOSSES = {
+    'convnext': nn.L1Loss
+}
+
+
+def build_aux_decoder(
+        in_dims: int, out_dims: int,
+        aux_decoder_arch: str, aux_decoder_args: dict
+) -> torch.nn.Module:
+    decoder_cls = AUX_DECODERS[aux_decoder_arch]
+    kwargs = filter_kwargs(aux_decoder_args, decoder_cls)
+    return AUX_DECODERS[aux_decoder_arch](in_dims, out_dims, **kwargs)
+
+
+def build_aux_loss(aux_decoder_arch):
+    return AUX_LOSSES[aux_decoder_arch]()
+
+
+class AuxDecoderAdaptor(nn.Module):
+    def __init__(self, in_dims: int, out_dims: int, num_feats: int,
+                 spec_min: list, spec_max: list,
+                 aux_decoder_arch: str, aux_decoder_args: dict):
+        super().__init__()
+        self.decoder = build_aux_decoder(
+            in_dims=in_dims, out_dims=out_dims * num_feats,
+            aux_decoder_arch=aux_decoder_arch,
+            aux_decoder_args=aux_decoder_args
+        )
+        self.out_dims = out_dims
+        self.n_feats = num_feats
+        if spec_min is not None and spec_max is not None:
+            # spec: [B, T, M] or [B, F, T, M]
+            # spec_min and spec_max: [1, 1, M] or [1, 1, F, M] => transpose(-3, -2) => [1, 1, M] or [1, F, 1, M]
+            spec_min = torch.FloatTensor(spec_min)[None, None, :].transpose(-3, -2)
+            spec_max = torch.FloatTensor(spec_max)[None, None, :].transpose(-3, -2)
+            self.register_buffer('spec_min', spec_min, persistent=False)
+            self.register_buffer('spec_max', spec_max, persistent=False)
+
+    def norm_spec(self, x):
+        k = (self.spec_max - self.spec_min) / 2.
+        b = (self.spec_max + self.spec_min) / 2.
+        return (x - b) / k
+
+    def denorm_spec(self, x):
+        k = (self.spec_max - self.spec_min) / 2.
+        b = (self.spec_max + self.spec_min) / 2.
+        return x * k + b
+
+    def forward(self, condition, infer=False):
+        x = self.decoder(condition, infer=infer)  # [B, T, F x C]
+
+        if self.n_feats > 1:
+            # This is the temporary solution since PyTorch 1.13
+            # does not support exporting aten::unflatten to ONNX
+            # x = x.unflatten(dim=2, sizes=(self.n_feats, self.in_dims))
+            x = x.reshape(-1, x.shape[1], self.n_feats, self.out_dims)  # [B, T, F, C]
+            x = x.transpose(1, 2)  # [B, F, T, C]
+        if infer:
+            x = self.denorm_spec(x)
+
+        return x  # [B, T, C] or [B, F, T, C]

modules/aux_decoder/convnext.py

@@ -0,0 +1,89 @@
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+from modules.commons.common_layers import AdamWConv1d
+
+
+class ConvNeXtBlock(nn.Module):
+    """ConvNeXt Block adapted from https://github.com/facebookresearch/ConvNeXt to 1D audio signal.
+
+    Args:
+        dim (int): Number of input channels.
+        intermediate_dim (int): Dimensionality of the intermediate layer.
+        layer_scale_init_value (float, optional): Initial value for the layer scale. None means no scaling.
+            Defaults to None.
+    """
+
+    def __init__(
+            self,
+            dim: int,
+            intermediate_dim: int,
+            layer_scale_init_value: Optional[float] = None, drop_out: float = 0.0
+
+    ):
+        super().__init__()
+        self.dwconv = nn.Conv1d(dim, dim, kernel_size=7, padding=3, groups=dim)  # depthwise conv
+
+        self.norm = nn.LayerNorm(dim, eps=1e-6)
+        self.pwconv1 = nn.Linear(dim, intermediate_dim)  # pointwise/1x1 convs, implemented with linear layers
+        self.act = nn.GELU()
+        self.pwconv2 = nn.Linear(intermediate_dim, dim)
+        self.gamma = (
+            nn.Parameter(layer_scale_init_value * torch.ones(dim), requires_grad=True)
+            if layer_scale_init_value > 0
+            else None
+        )
+        # self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.drop_path = nn.Identity()
+        self.dropout = nn.Dropout(drop_out) if drop_out > 0. else nn.Identity()
+
+    def forward(self, x: torch.Tensor, ) -> torch.Tensor:
+        residual = x
+        x = self.dwconv(x)
+        x = x.transpose(1, 2)  # (B, C, T) -> (B, T, C)
+
+        x = self.norm(x)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.pwconv2(x)
+        if self.gamma is not None:
+            x = self.gamma * x
+        x = x.transpose(1, 2)  # (B, T, C) -> (B, C, T)
+        x = self.dropout(x)
+
+        x = residual + self.drop_path(x)
+        return x
+
+
+class ConvNeXtDecoder(nn.Module):
+    def __init__(
+            self, in_dims, out_dims, /, *,
+            num_channels=512, num_layers=6, kernel_size=7, dropout_rate=0.1
+    ):
+        super().__init__()
+        self.inconv = nn.Conv1d(
+            in_dims, num_channels, kernel_size,
+            stride=1, padding=(kernel_size - 1) // 2
+        )
+        self.conv = nn.ModuleList(
+            ConvNeXtBlock(
+                dim=num_channels, intermediate_dim=num_channels * 4,
+                layer_scale_init_value=1e-6, drop_out=dropout_rate
+            ) for _ in range(num_layers)
+        )
+        self.outconv = AdamWConv1d(
+            num_channels, out_dims, kernel_size,
+            stride=1, padding=(kernel_size - 1) // 2
+        )
+
+    # noinspection PyUnusedLocal
+    def forward(self, x, infer=False):
+        x = x.transpose(1, 2)
+        x = self.inconv(x)
+        for conv in self.conv:
+            x = conv(x)
+        x = self.outconv(x)
+        x = x.transpose(1, 2)
+        return x

modules/backbones/__init__.py

@@ -0,0 +1,20 @@
+import torch.nn
+from modules.backbones.wavenet import WaveNet
+from modules.backbones.lynxnet import LYNXNet
+from modules.backbones.lynxnet2 import LYNXNet2
+from utils import filter_kwargs
+
+BACKBONES = {
+    'wavenet': WaveNet,
+    'lynxnet': LYNXNet,
+    'lynxnet2': LYNXNet2,
+}
+
+
+def build_backbone(
+        out_dims: int, num_feats: int,
+        backbone_type: str, backbone_args: dict
+) -> torch.nn.Module:
+    backbone = BACKBONES[backbone_type]
+    kwargs = filter_kwargs(backbone_args, backbone)
+    return BACKBONES[backbone_type](out_dims, num_feats, **kwargs)

modules/backbones/lynxnet.py

@@ -0,0 +1,147 @@
+# refer to： 
+# https://github.com/CNChTu/Diffusion-SVC/blob/v2.0_dev/diffusion/naive_v2/model_conformer_naive.py
+# https://github.com/CNChTu/Diffusion-SVC/blob/v2.0_dev/diffusion/naive_v2/naive_v2_diff.py
+
+import torch.nn as nn
+import torch.nn.functional as F
+
+from modules.commons.common_layers import SinusoidalPosEmb, SwiGLU, Transpose, AdamWConv1d
+from modules.commons.common_layers import KaimingNormalConv1d as Conv1d
+from utils.hparams import hparams
+
+
+class LYNXConvModule(nn.Module):
+    @staticmethod
+    def calc_same_padding(kernel_size):
+        pad = kernel_size // 2
+        return pad, pad - (kernel_size + 1) % 2
+
+    def __init__(self, dim, expansion_factor, kernel_size=31, activation='PReLU', dropout=0.0):
+        super().__init__()
+        inner_dim = dim * expansion_factor
+        activation_classes = {
+            'SiLU': nn.SiLU,
+            'ReLU': nn.ReLU,
+            'PReLU': lambda: nn.PReLU(inner_dim)
+        }
+        activation = activation if activation is not None else 'PReLU'
+        if activation not in activation_classes:
+            raise ValueError(f'{activation} is not a valid activation')
+        _activation = activation_classes[activation]()
+        padding = self.calc_same_padding(kernel_size)
+        if float(dropout) > 0.:
+            _dropout = nn.Dropout(dropout)
+        else:
+            _dropout = nn.Identity()
+        self.net = nn.Sequential(
+            nn.LayerNorm(dim),
+            Transpose((1, 2)),
+            nn.Conv1d(dim, inner_dim * 2, 1),
+            SwiGLU(dim=1),
+            nn.Conv1d(inner_dim, inner_dim, kernel_size=kernel_size, padding=padding[0], groups=inner_dim),
+            _activation,
+            nn.Conv1d(inner_dim, dim, 1),
+            Transpose((1, 2)),
+            _dropout
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class LYNXNetResidualLayer(nn.Module):
+    def __init__(self, dim_cond, dim, expansion_factor, kernel_size=31, activation='PReLU', dropout=0.0):
+        super().__init__()
+        self.diffusion_projection = nn.Conv1d(dim, dim, 1)
+        self.conditioner_projection = nn.Conv1d(dim_cond, dim, 1)
+        self.convmodule = LYNXConvModule(dim=dim, expansion_factor=expansion_factor, kernel_size=kernel_size,
+                                         activation=activation, dropout=dropout)
+
+    def forward(self, x, conditioner, diffusion_step, front_cond_inject=False):
+        if front_cond_inject:
+            x = x + self.conditioner_projection(conditioner)
+            res_x = x
+        else:
+            res_x = x
+            x = x + self.conditioner_projection(conditioner)
+        x = x + self.diffusion_projection(diffusion_step)
+        x = x.transpose(1, 2)
+        x = self.convmodule(x)  # (#batch, dim, length) 
+        x = x.transpose(1, 2) + res_x
+        return x  # (#batch, length, dim)
+
+
+class LYNXNet(nn.Module):
+    def __init__(self, in_dims, n_feats, *, num_layers=6, num_channels=512, expansion_factor=2, kernel_size=31,
+                 activation='PReLU', dropout_rate=0.0, strong_cond=False):
+        """
+        LYNXNet(Linear Gated Depthwise Separable Convolution Network)
+        TIPS:You can control the style of the generated results by modifying the 'activation', 
+            - 'PReLU'(default) : Similar to WaveNet
+            - 'SiLU' : Voice will be more pronounced, not recommended for use under DDPM
+            - 'ReLU' : Contrary to 'SiLU', Voice will be weakened
+        """
+        super().__init__()
+        self.in_dims = in_dims
+        self.n_feats = n_feats
+        self.input_projection = Conv1d(in_dims * n_feats, num_channels, 1)
+        self.diffusion_embedding = nn.Sequential(
+            SinusoidalPosEmb(num_channels),
+            nn.Linear(num_channels, num_channels * 4),
+            nn.GELU(),
+            nn.Linear(num_channels * 4, num_channels),
+        )
+        self.residual_layers = nn.ModuleList(
+            [
+                LYNXNetResidualLayer(
+                    dim_cond=hparams['hidden_size'],
+                    dim=num_channels,
+                    expansion_factor=expansion_factor,
+                    kernel_size=kernel_size,
+                    activation=activation,
+                    dropout=dropout_rate
+                )
+                for _ in range(num_layers)
+            ]
+        )
+        self.norm = nn.LayerNorm(num_channels)
+        self.output_projection = AdamWConv1d(num_channels, in_dims * n_feats, kernel_size=1)
+        self.strong_cond = strong_cond
+        nn.init.zeros_(self.output_projection.weight)
+
+    def forward(self, spec, diffusion_step, cond):
+        """
+        :param spec: [B, F, M, T]
+        :param diffusion_step: [B, 1]
+        :param cond: [B, H, T]
+        :return:
+        """
+
+        if self.n_feats == 1:
+            x = spec[:, 0]  # [B, M, T]
+        else:
+            x = spec.flatten(start_dim=1, end_dim=2)  # [B, F x M, T]
+
+        x = self.input_projection(x)  # x [B, residual_channel, T]
+        if not self.strong_cond:
+            x = F.gelu(x)
+
+        diffusion_step = self.diffusion_embedding(diffusion_step).unsqueeze(-1)
+
+        for layer in self.residual_layers:
+            x = layer(x, cond, diffusion_step, front_cond_inject=self.strong_cond)
+
+        # post-norm
+        x = self.norm(x.transpose(1, 2)).transpose(1, 2)
+
+        # output_projection
+        x = self.output_projection(x)  # [B, 128, T]
+
+        if self.n_feats == 1:
+            x = x[:, None, :, :]
+        else:
+            # This is the temporary solution since PyTorch 1.13
+            # does not support exporting aten::unflatten to ONNX
+            # x = x.unflatten(dim=1, sizes=(self.n_feats, self.in_dims))
+            x = x.reshape(-1, self.n_feats, self.in_dims, x.shape[2])
+        return x

modules/backbones/lynxnet2.py

@@ -0,0 +1,115 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from modules.commons.common_layers import SinusoidalPosEmb, SwiGLU, ATanGLU, Transpose, AdamWLinear
+from utils.hparams import hparams
+
+
+class LYNXNet2Block(nn.Module):
+    def __init__(self, dim, expansion_factor, kernel_size=31, dropout=0., glu_type='swiglu'):
+        super().__init__()
+        inner_dim = int(dim * expansion_factor)
+        if glu_type == 'swiglu':
+            _glu = SwiGLU()
+        elif glu_type == 'atanglu':
+            _glu = ATanGLU()
+        else:
+            raise ValueError(f'{glu_type} is not a valid activation')
+        if float(dropout) > 0.:
+            _dropout = nn.Dropout(dropout)
+        else:
+            _dropout = nn.Identity()
+        self.net = nn.Sequential(
+            nn.LayerNorm(dim),
+            Transpose((1, 2)),
+            nn.Conv1d(dim, dim, kernel_size=kernel_size, padding=kernel_size // 2, groups=dim),
+            Transpose((1, 2)),
+            nn.Linear(dim, inner_dim * 2),
+            _glu,
+            nn.Linear(inner_dim, inner_dim * 2),
+            _glu,
+            nn.Linear(inner_dim, dim),
+            _dropout
+        )
+
+    def forward(self, x):
+        return x + self.net(x)
+
+
+class LYNXNet2(nn.Module):
+    def __init__(self, in_dims, n_feats, *, num_layers=6, num_channels=512, expansion_factor=1, kernel_size=31,
+                 dropout_rate=0.0, use_conditioner_cache=False, glu_type='swiglu'):
+        """
+        LYNXNet2(Linear Gated Depthwise Separable Convolution Network Version 2)
+        """
+        super().__init__()
+        self.in_dims = in_dims
+        self.n_feats = n_feats
+        self.input_projection = nn.Linear(in_dims * n_feats, num_channels)
+        self.use_conditioner_cache = use_conditioner_cache
+        if self.use_conditioner_cache:
+            # Conv1d is used for condition cache compatibility
+            self.conditioner_projection = nn.Conv1d(hparams['hidden_size'], num_channels, 1)
+        else:
+            self.conditioner_projection = nn.Linear(hparams['hidden_size'], num_channels)
+        self.diffusion_embedding = nn.Sequential(
+            SinusoidalPosEmb(num_channels),
+            nn.Linear(num_channels, num_channels * 4),
+            nn.GELU(),
+            nn.Linear(num_channels * 4, num_channels),
+        )
+        self.residual_layers = nn.ModuleList(
+            [
+                LYNXNet2Block(
+                    dim=num_channels,
+                    expansion_factor=expansion_factor,
+                    kernel_size=kernel_size,
+                    dropout=dropout_rate,
+                    glu_type=glu_type
+                )
+                for _ in range(num_layers)
+            ]
+        )
+        self.norm = nn.LayerNorm(num_channels)
+        self.output_projection = AdamWLinear(num_channels, in_dims * n_feats)
+        nn.init.kaiming_normal_(self.input_projection.weight)
+        nn.init.kaiming_normal_(self.conditioner_projection.weight)
+        nn.init.zeros_(self.output_projection.weight)
+
+    def forward(self, spec, diffusion_step, cond):
+        """
+        :param spec: [B, F, M, T]
+        :param diffusion_step: [B, 1]
+        :param cond: [B, H, T]
+        :return:
+        """
+
+        if self.n_feats == 1:
+            x = spec[:, 0]  # [B, M, T]
+        else:
+            x = spec.flatten(start_dim=1, end_dim=2)  # [B, F x M, T]
+
+        x = self.input_projection(x.transpose(1, 2)) # [B, T, F x M]
+        if self.use_conditioner_cache:
+            x = x + self.conditioner_projection(cond).transpose(1, 2)
+        else:
+            x = x + self.conditioner_projection(cond.transpose(1, 2))
+        x = x + self.diffusion_embedding(diffusion_step).unsqueeze(1)
+
+        for layer in self.residual_layers:
+            x = layer(x)
+
+        # post-norm
+        x = self.norm(x)
+
+        # output projection
+        x = self.output_projection(x).transpose(1, 2)  # [B, 128, T]
+
+        if self.n_feats == 1:
+            x = x[:, None, :, :]
+        else:
+            # Using reshape instead of unflatten for ONNX export compatibility
+            # x = x.unflatten(dim=1, sizes=(self.n_feats, self.in_dims))
+            x = x.reshape(-1, self.n_feats, self.in_dims, x.shape[2])
+        return x

modules/backbones/wavenet.py

@@ -0,0 +1,104 @@
+import math
+from math import sqrt
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from modules.commons.common_layers import SinusoidalPosEmb, AdamWConv1d
+from modules.commons.common_layers import KaimingNormalConv1d as Conv1d
+from utils.hparams import hparams
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, encoder_hidden, residual_channels, dilation):
+        super().__init__()
+        self.residual_channels = residual_channels
+        self.dilated_conv = nn.Conv1d(
+            residual_channels,
+            2 * residual_channels,
+            kernel_size=3,
+            padding=dilation,
+            dilation=dilation
+        )
+        self.diffusion_projection = nn.Linear(residual_channels, residual_channels)
+        self.conditioner_projection = nn.Conv1d(encoder_hidden, 2 * residual_channels, 1)
+        self.output_projection = nn.Conv1d(residual_channels, 2 * residual_channels, 1)
+
+    def forward(self, x, conditioner, diffusion_step):
+        diffusion_step = self.diffusion_projection(diffusion_step).unsqueeze(-1)
+        conditioner = self.conditioner_projection(conditioner)
+        y = x + diffusion_step
+
+        y = self.dilated_conv(y) + conditioner
+
+        # Using torch.split instead of torch.chunk to avoid using onnx::Slice
+        gate, filter = torch.split(y, [self.residual_channels, self.residual_channels], dim=1)
+        y = torch.sigmoid(gate) * torch.tanh(filter)
+
+        y = self.output_projection(y)
+
+        # Using torch.split instead of torch.chunk to avoid using onnx::Slice
+        residual, skip = torch.split(y, [self.residual_channels, self.residual_channels], dim=1)
+        return (x + residual) / math.sqrt(2.0), skip
+
+
+class WaveNet(nn.Module):
+    def __init__(self, in_dims, n_feats, *, num_layers=20, num_channels=256, dilation_cycle_length=4):
+        super().__init__()
+        self.in_dims = in_dims
+        self.n_feats = n_feats
+        self.input_projection = Conv1d(in_dims * n_feats, num_channels, 1)
+        self.diffusion_embedding = SinusoidalPosEmb(num_channels)
+        self.mlp = nn.Sequential(
+            nn.Linear(num_channels, num_channels * 4),
+            nn.Mish(),
+            nn.Linear(num_channels * 4, num_channels)
+        )
+        self.residual_layers = nn.ModuleList([
+            ResidualBlock(
+                encoder_hidden=hparams['hidden_size'],
+                residual_channels=num_channels,
+                dilation=2 ** (i % dilation_cycle_length)
+            )
+            for i in range(num_layers)
+        ])
+        self.skip_projection = Conv1d(num_channels, num_channels, 1)
+        self.output_projection = AdamWConv1d(num_channels, in_dims * n_feats, 1)
+        nn.init.zeros_(self.output_projection.weight)
+
+    def forward(self, spec, diffusion_step, cond):
+        """
+        :param spec: [B, F, M, T]
+        :param diffusion_step: [B, 1]
+        :param cond: [B, H, T]
+        :return:
+        """
+        if self.n_feats == 1:
+            # Use indexing instead of squeeze to avoid emitting an onnx::If
+            # whose branches have different rank, which breaks shape inference
+            # for the downstream Conv on PyTorch >= 2.0.
+            x = spec[:, 0]  # [B, M, T]
+        else:
+            x = spec.flatten(start_dim=1, end_dim=2)  # [B, F x M, T]
+        x = self.input_projection(x)  # [B, C, T]
+
+        x = F.relu(x)
+        diffusion_step = self.diffusion_embedding(diffusion_step)
+        diffusion_step = self.mlp(diffusion_step)
+        skip = []
+        for layer in self.residual_layers:
+            x, skip_connection = layer(x, cond, diffusion_step)
+            skip.append(skip_connection)
+
+        x = torch.sum(torch.stack(skip), dim=0) / sqrt(len(self.residual_layers))
+        x = self.skip_projection(x)
+        x = F.relu(x)
+        x = self.output_projection(x)  # [B, M, T]
+        if self.n_feats == 1:
+            x = x[:, None, :, :]
+        else:
+            # Using reshape instead of unflatten for ONNX export compatibility
+            # x = x.unflatten(dim=1, sizes=(self.n_feats, self.in_dims))
+            x = x.reshape(-1, self.n_feats, self.in_dims, x.shape[2])
+        return x

modules/commons/common_layers.py

@@ -0,0 +1,427 @@
+from __future__ import annotations
+
+import math
+import numpy as np
+import torch
+import torch.nn.functional as F
+import torch.onnx.operators
+from torch import nn
+from torch.nn import LayerNorm, ReLU, GELU, SiLU
+
+import utils
+
+
+class NormalInitEmbedding(torch.nn.Embedding):
+    def __init__(
+            self,
+            num_embeddings: int,
+            embedding_dim: int,
+            padding_idx: int | None = None,
+            *args,
+            **kwargs
+    ):
+        super().__init__(num_embeddings, embedding_dim, *args, padding_idx=padding_idx, **kwargs)
+        nn.init.normal_(self.weight, mean=0, std=self.embedding_dim ** -0.5)
+        if padding_idx is not None:
+            nn.init.constant_(self.weight[padding_idx], 0)
+
+
+class AdamWLinear(torch.nn.Linear):
+    def __init__(
+            self,
+            in_features: int,
+            out_features: int,
+            *args,
+            bias: bool = True,
+            **kwargs
+    ):
+        super().__init__(in_features, out_features, *args, bias=bias, **kwargs)
+        nn.init.xavier_uniform_(self.weight)
+        if bias:
+            nn.init.constant_(self.bias, 0.)
+
+
+class XavierUniformInitLinear(torch.nn.Linear):
+    def __init__(
+            self,
+            in_features: int,
+            out_features: int,
+            *args,
+            bias: bool = True,
+            **kwargs
+    ):
+        super().__init__(in_features, out_features, *args, bias=bias, **kwargs)
+        nn.init.xavier_uniform_(self.weight)
+        if bias:
+            nn.init.constant_(self.bias, 0.)
+
+
+class SinusoidalPositionalEmbedding(nn.Module):
+    """This module produces sinusoidal positional embeddings of any length.
+
+    Padding symbols are ignored.
+    """
+
+    def __init__(self, embedding_dim, padding_idx, init_size=1024):
+        super().__init__()
+        self.embedding_dim = embedding_dim
+        self.padding_idx = padding_idx
+        self.weights = SinusoidalPositionalEmbedding.get_embedding(
+            init_size,
+            embedding_dim,
+            padding_idx,
+        )
+        self.register_buffer('_float_tensor', torch.FloatTensor(1))
+
+    @staticmethod
+    def get_embedding(num_embeddings, embedding_dim, padding_idx=None):
+        """Build sinusoidal embeddings.
+
+        This matches the implementation in tensor2tensor, but differs slightly
+        from the description in Section 3.5 of "Attention Is All You Need".
+        """
+        half_dim = embedding_dim // 2
+        emb = math.log(10000) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim, dtype=torch.float) * -emb)
+        emb = torch.arange(num_embeddings, dtype=torch.float).unsqueeze(1) * emb.unsqueeze(0)
+        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1).view(num_embeddings, -1)
+        if embedding_dim % 2 == 1:
+            # zero pad
+            emb = torch.cat([emb, torch.zeros(num_embeddings, 1)], dim=1)
+        if padding_idx is not None:
+            emb[padding_idx, :] = 0
+        return emb
+
+    def forward(self, x, incremental_state=None, timestep=None, positions=None):
+        """Input is expected to be of size [bsz x seqlen]."""
+        bsz, seq_len = x.shape[:2]
+        max_pos = self.padding_idx + 1 + seq_len
+        if self.weights is None or max_pos > self.weights.size(0):
+            # recompute/expand embeddings if needed
+            self.weights = SinusoidalPositionalEmbedding.get_embedding(
+                max_pos,
+                self.embedding_dim,
+                self.padding_idx,
+            )
+        self.weights = self.weights.to(self._float_tensor)
+
+        if incremental_state is not None:
+            # positions is the same for every token when decoding a single step
+            pos = timestep.view(-1)[0] + 1 if timestep is not None else seq_len
+            return self.weights[self.padding_idx + pos, :].expand(bsz, 1, -1)
+
+        positions = utils.make_positions(x, self.padding_idx) if positions is None else positions
+        return self.weights.index_select(0, positions.view(-1)).view(bsz, seq_len, -1).detach()
+
+    @staticmethod
+    def max_positions():
+        """Maximum number of supported positions."""
+        return int(1e5)  # an arbitrary large number
+
+
+class SwiGLU(nn.Module):
+    # Swish-Applies the gated linear unit function.
+    def __init__(self, dim=-1):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, x):
+        # out, gate = x.chunk(2, dim=self.dim)
+        # Using torch.split instead of chunk for ONNX export compatibility.
+        out, gate = torch.split(x, x.size(self.dim) // 2, dim=self.dim)
+        gate = F.silu(gate)
+        if x.dtype == torch.float16:
+            out_min, out_max = torch.aminmax(out.detach())
+            gate_min, gate_max = torch.aminmax(gate.detach())
+            max_abs_out = torch.max(-out_min, out_max).float()
+            max_abs_gate = torch.max(-gate_min, gate_max).float()
+            max_abs_value = max_abs_out * max_abs_gate
+            if max_abs_value > 1000:
+                ratio = (1000 / max_abs_value).half()
+                gate = gate * ratio
+                return (out * gate).clamp(-1000 * ratio, 1000 * ratio) / ratio
+        return out * gate
+
+
+class ATanGLUFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, out, gate):
+        atan_gate = torch.atan(gate)
+        decay_out = out / gate.square().add(1.0)
+        ctx.save_for_backward(decay_out, atan_gate)
+        return out * atan_gate
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        decay_out, atan_gate = ctx.saved_tensors
+        grad_out_part = grad_output * atan_gate
+        grad_gate_part = grad_output * decay_out
+        return grad_out_part, grad_gate_part
+
+
+class ATanGLU(nn.Module):
+    # ArcTan-Applies the gated linear unit function.
+    def __init__(self, dim=-1):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, x):
+        # out, gate = x.chunk(2, dim=self.dim)
+        # Using torch.split instead of chunk for ONNX export compatibility.        
+        out, gate = torch.split(x, x.size(self.dim) // 2, dim=self.dim)
+        if self.training:
+            return ATanGLUFunction.apply(out, gate)
+        else:
+            return out * torch.atan(gate)
+
+
+class AdamWConv1d(torch.nn.Conv1d):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        nn.init.kaiming_normal_(self.weight)
+
+
+class KaimingNormalConv1d(torch.nn.Conv1d):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        nn.init.kaiming_normal_(self.weight)
+
+
+class Transpose(nn.Module):
+    def __init__(self, dims):
+        super().__init__()
+        assert len(dims) == 2, 'dims must be a tuple of two dimensions'
+        self.dims = dims
+
+    def forward(self, x):
+        return x.transpose(*self.dims)
+
+
+class Mixed_LayerNorm(nn.Module):
+    def __init__(
+            self,
+            channels: int,
+            condition_channels: int,
+            beta_distribution_concentration: float = 0.2,
+            eps: float = 1e-5,
+            bias: bool = True
+    ):
+        super().__init__()
+        self.channels = channels
+        self.eps = eps
+
+        self.beta_distribution = torch.distributions.Beta(
+            beta_distribution_concentration,
+            beta_distribution_concentration
+        )
+
+        self.affine = XavierUniformInitLinear(condition_channels, channels * 2, bias=bias)
+        if self.affine.bias is not None:
+            self.affine.bias.data[:channels] = 0  # betas (shift)
+            self.affine.bias.data[channels:] = 1  # gammas (scale)
+
+    def forward(
+            self,
+            x: torch.FloatTensor,
+            condition: torch.FloatTensor  # -> shape [Batch, Cond_d]
+    ) -> torch.FloatTensor:
+        x = F.layer_norm(x, normalized_shape=(self.channels,), weight=None, bias=None, eps=self.eps)
+
+        affine_params = self.affine(condition)
+        if affine_params.ndim == 2:
+            affine_params = affine_params.unsqueeze(1)
+        betas, gammas = torch.split(affine_params, self.channels, dim=-1)
+
+        if not self.training or x.size(0) == 1:
+            return gammas * x + betas
+
+        shuffle_indices = torch.randperm(x.size(0), device=x.device)
+        shuffled_betas = betas[shuffle_indices]
+        shuffled_gammas = gammas[shuffle_indices]
+
+        beta_samples = self.beta_distribution.sample((x.size(0), 1, 1)).to(x.device)
+        mixed_betas = beta_samples * betas + (1 - beta_samples) * shuffled_betas
+        mixed_gammas = beta_samples * gammas + (1 - beta_samples) * shuffled_gammas
+
+        return mixed_gammas * x + mixed_betas
+
+
+class TransformerFFNLayer(nn.Module):
+    def __init__(self, hidden_size, filter_size, kernel_size=1, dropout=0., act='gelu'):
+        super().__init__()
+        self.kernel_size = kernel_size
+        self.dropout = dropout
+        self.act = act
+        filter_size_1 = filter_size
+        if self.act == 'relu':
+            self.act_fn = ReLU()
+        elif self.act == 'gelu':
+            self.act_fn = GELU()
+        elif self.act == 'swish':
+            self.act_fn = SiLU()
+        elif self.act == 'swiglu':
+            self.act_fn = SwiGLU()
+            filter_size_1 = filter_size * 2
+        elif self.act == 'atanglu':
+            self.act_fn = ATanGLU()
+            filter_size_1 = filter_size * 2
+        else:
+            raise ValueError(f'{act} is not a valid activation')
+        self.ffn_1 = nn.Conv1d(hidden_size, filter_size_1, kernel_size, padding=kernel_size // 2)
+        self.ffn_2 = XavierUniformInitLinear(filter_size, hidden_size)
+
+    def forward(self, x):
+        # x: B x T x C
+        x = self.ffn_1(x.transpose(1, 2)).transpose(1, 2)
+        x = x * self.kernel_size ** -0.5
+
+        x = self.act_fn(x)
+        x = F.dropout(x, self.dropout, training=self.training)
+        x = self.ffn_2(x)
+        return x
+
+
+class MultiheadSelfAttentionWithRoPE(nn.Module):
+    def __init__(self, embed_dim, num_heads, dropout=0.1, bias=False, rotary_embed=None):
+        super().__init__()
+        assert embed_dim % num_heads == 0, "Embedding dimension must be divisible by number of heads"
+
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.head_dim = embed_dim // num_heads
+
+        # Linear layers for Q, K, V projections
+        self.in_proj = nn.Linear(embed_dim, embed_dim * 3, bias=bias)
+
+        # Final linear layer after concatenation
+        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+
+        # Dropout layer
+        self.dropout = nn.Dropout(dropout)
+
+        # Rotary Embeddings
+        self.rotary_embed = rotary_embed
+
+        # Initialization parameters
+        nn.init.xavier_uniform_(self.in_proj.weight)
+        nn.init.xavier_uniform_(self.out_proj.weight)
+        if bias:
+            nn.init.constant_(self.in_proj.bias, 0.0)
+            nn.init.constant_(self.out_proj.bias, 0.0)
+
+    def forward(self, x, key_padding_mask=None):
+        # x: (B, L, C)
+        # key_padding_mask: (B, L)
+        batch_size, seq_len, embed_dim = x.size()
+
+        # Project inputs to Q, K, V
+        Q, K, V = torch.split(self.in_proj(x), self.embed_dim, dim=-1)
+
+        # Reshape Q, K, V for multi-head attention
+        Q = Q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)  # (B, H, L, D)
+        K = K.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)  # (B, H, L, D)
+        V = V.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)  # (B, H, L, D)
+
+        # Apply RoPE
+        if self.rotary_embed is not None:
+            Q = self.rotary_embed.rotate_queries_or_keys(Q)
+            K = self.rotary_embed.rotate_queries_or_keys(K)
+
+        # Compute attention scores
+        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.head_dim)  # (B, H, L, L)
+
+        # Apply key padding mask if provided
+        if key_padding_mask is not None:
+            # Expand mask to match attention scores shape
+            mask = key_padding_mask.unsqueeze(1).unsqueeze(1)  # (B, 1, 1, L)
+            scores = scores.masked_fill(mask == 1, -np.inf)  # Masked positions are set to -inf
+
+        # Compute attention weights
+        attn_weights = F.softmax(scores, dim=-1)  # (B, H, L, L)
+        attn_weights = self.dropout(attn_weights)
+
+        # Apply attention weights to V
+        attn_output = torch.matmul(attn_weights, V)  # (B, H, L, D)
+
+        # Reshape and concatenate heads
+        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, embed_dim)  # (B, L, C)
+
+        # Final linear projection
+        output = self.out_proj(attn_output)  # (B, L, C)
+
+        return output
+
+
+class EncSALayer(nn.Module):
+    def __init__(self, c, num_heads, dropout, attention_dropout=0.1,
+                 relu_dropout=0.1, kernel_size=9, act='gelu', rotary_embed=None,
+                 layer_idx=None, mix_ln_layer=None
+                 ):
+        super().__init__()
+        self.dropout = dropout
+        self.use_mix_ln = (
+                layer_idx is not None
+                and mix_ln_layer is not None
+                and layer_idx in mix_ln_layer
+        )
+        if self.use_mix_ln:
+            self.layer_norm1 = Mixed_LayerNorm(c, c)
+        else:
+            self.layer_norm1 = LayerNorm(c)
+        # Always use the in-house manual attention. With rotary_embed=None this
+        # is a plain multi-head self-attention that is ONNX-export safe across
+        # dynamic sequence lengths. Using torch.nn.MultiheadAttention here was
+        # the source of the "Reshape baked tgt_len" bug on PyTorch >= 2.0
+        # because its SDPA-branched implementation specializes tgt_len to a
+        # Python int and re-injects it into the output Reshape.
+        self.self_attn = MultiheadSelfAttentionWithRoPE(
+            c, num_heads, dropout=attention_dropout, bias=False, rotary_embed=rotary_embed
+        )
+        if self.use_mix_ln:
+            self.layer_norm2 = Mixed_LayerNorm(c, c)
+        else:
+            self.layer_norm2 = LayerNorm(c)
+        self.ffn = TransformerFFNLayer(
+            c, 4 * c, kernel_size=kernel_size, dropout=relu_dropout, act=act
+        )
+
+    def forward(self, x, encoder_padding_mask=None, cond=None, **kwargs):
+        layer_norm_training = kwargs.get('layer_norm_training', None)
+        if layer_norm_training is not None:
+            self.layer_norm1.training = layer_norm_training
+            self.layer_norm2.training = layer_norm_training
+        residual = x
+        if self.use_mix_ln:
+            x = self.layer_norm1(x, cond)
+        else:
+            x = self.layer_norm1(x)
+        x = self.self_attn(x, key_padding_mask=encoder_padding_mask)
+        x = F.dropout(x, self.dropout, training=self.training)
+        x = residual + x
+        x = x * (1 - encoder_padding_mask.float())[..., None]
+
+        residual = x
+        if self.use_mix_ln:
+            x = self.layer_norm2(x, cond)
+        else:
+            x = self.layer_norm2(x)
+        x = self.ffn(x)
+        x = F.dropout(x, self.dropout, training=self.training)
+        x = residual + x
+        x = x * (1 - encoder_padding_mask.float())[..., None]
+        return x
+
+
+class SinusoidalPosEmb(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, x):
+        device = x.device
+        half_dim = self.dim // 2
+        emb = math.log(10000) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
+        emb = x.unsqueeze(-1) * emb.unsqueeze(0)
+        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
+        return emb

modules/commons/espnet_positional_embedding.py

@@ -0,0 +1,113 @@
+import math
+import torch
+
+
+class PositionalEncoding(torch.nn.Module):
+    """Positional encoding.
+    Args:
+        d_model (int): Embedding dimension.
+        dropout_rate (float): Dropout rate.
+        max_len (int): Maximum input length.
+        reverse (bool): Whether to reverse the input position.
+    """
+
+    def __init__(self, d_model, dropout_rate, max_len=5000, reverse=False):
+        """Construct an PositionalEncoding object."""
+        super(PositionalEncoding, self).__init__()
+        self.d_model = d_model
+        self.reverse = reverse
+        self.xscale = math.sqrt(self.d_model)
+        self.dropout = torch.nn.Dropout(p=dropout_rate)
+        self.pe = None
+        self.extend_pe(torch.tensor(0.0).expand(1, max_len))
+
+    def extend_pe(self, x):
+        """Reset the positional encodings."""
+        if self.pe is not None:
+            if self.pe.size(1) >= x.size(1):
+                if self.pe.dtype != x.dtype or self.pe.device != x.device:
+                    self.pe = self.pe.to(dtype=x.dtype, device=x.device)
+                return
+        if self.reverse:
+            position = torch.arange(
+                x.size(1) - 1, -1, -1.0, dtype=torch.float32
+            ).unsqueeze(1)
+        else:
+            position = torch.arange(0, x.size(1), dtype=torch.float32).unsqueeze(1)
+        div_term = torch.exp(
+            torch.arange(0, self.d_model, 2, dtype=torch.float32)
+            * -(math.log(10000.0) / self.d_model)
+        )
+        pe = torch.stack([
+            torch.sin(position * div_term),
+            torch.cos(position * div_term)
+        ], dim=2).view(-1, self.d_model).unsqueeze(0)
+        self.pe = pe.to(device=x.device, dtype=x.dtype)
+
+    def forward(self, x: torch.Tensor):
+        """Add positional encoding.
+        Args:
+            x (torch.Tensor): Input tensor (batch, time, `*`).
+        Returns:
+            torch.Tensor: Encoded tensor (batch, time, `*`).
+        """
+        self.extend_pe(x)
+        x = x * self.xscale + self.pe[:, : x.size(1)]
+        return self.dropout(x)
+
+
+class ScaledPositionalEncoding(PositionalEncoding):
+    """Scaled positional encoding module.
+    See Sec. 3.2  https://arxiv.org/abs/1809.08895
+    Args:
+        d_model (int): Embedding dimension.
+        dropout_rate (float): Dropout rate.
+        max_len (int): Maximum input length.
+    """
+
+    def __init__(self, d_model, dropout_rate, max_len=5000):
+        """Initialize class."""
+        super().__init__(d_model=d_model, dropout_rate=dropout_rate, max_len=max_len)
+        self.alpha = torch.nn.Parameter(torch.tensor(1.0))
+
+    def reset_parameters(self):
+        """Reset parameters."""
+        self.alpha.data = torch.tensor(1.0)
+
+    def forward(self, x):
+        """Add positional encoding.
+        Args:
+            x (torch.Tensor): Input tensor (batch, time, `*`).
+        Returns:
+            torch.Tensor: Encoded tensor (batch, time, `*`).
+        """
+        self.extend_pe(x)
+        x = x + self.alpha * self.pe[:, : x.size(1)]
+        return self.dropout(x)
+
+
+class RelPositionalEncoding(PositionalEncoding):
+    """Relative positional encoding module.
+    See : Appendix B in https://arxiv.org/abs/1901.02860
+    Args:
+        d_model (int): Embedding dimension.
+        dropout_rate (float): Dropout rate.
+        max_len (int): Maximum input length.
+    """
+
+    def __init__(self, d_model, dropout_rate, max_len=5000):
+        """Initialize class."""
+        super().__init__(d_model, dropout_rate, max_len, reverse=True)
+
+    def forward(self, x):
+        """Compute positional encoding.
+        Args:
+            x (torch.Tensor): Input tensor (batch, time, `*`).
+        Returns:
+            torch.Tensor: Encoded tensor (batch, time, `*`).
+            torch.Tensor: Positional embedding tensor (1, time, `*`).
+        """
+        self.extend_pe(x)
+        x = x * self.xscale
+        pos_emb = self.pe[:, : x.size(1)]
+        return self.dropout(x) + self.dropout(pos_emb)

modules/commons/rotary_embedding_torch.py

@@ -0,0 +1,63 @@
+import torch
+from einops import rearrange, repeat
+from torch import einsum, Tensor
+from torch.nn import Module
+
+
+def rotate_half(x: Tensor, interleaved=True) -> Tensor:
+    if not interleaved:
+        # x_half1, x_half2 = x.chunk(2, dim=-1)
+        # Using torch.split instead of chunk for ONNX export compatibility.
+        x1, x2 = torch.split(x, x.size(-1) // 2, dim=-1)
+        return torch.cat((-x2, x1), dim=-1)
+    else:
+        x = rearrange(x, '... (d r) -> ... d r', r=2)
+        x1, x2 = x.unbind(dim=-1)
+        x = torch.stack((-x2, x1), dim=-1)
+        return rearrange(x, '... d r -> ... (d r)')
+
+
+def apply_rotary_emb(freqs: Tensor, t: Tensor, interleaved=True) -> Tensor:
+    rot_dim = freqs.shape[-1]
+    t_to_rotate = t[..., :rot_dim]
+    t_pass_through = t[..., rot_dim:]
+
+    t_rotated = (t_to_rotate * freqs.cos()) + (rotate_half(t_to_rotate, interleaved) * freqs.sin())
+
+    return torch.cat((t_rotated, t_pass_through), dim=-1)
+
+
+class RotaryEmbedding(Module):
+    def __init__(
+            self,
+            dim,
+            theta=10000,
+            max_seq_len=8192,
+            interleaved: bool = True
+    ):
+        super().__init__()
+        self.interleaved = interleaved
+        self.cached_freqs_seq_len = max_seq_len
+        inv_freq = 1. / (theta ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer('inv_freq', inv_freq, persistent=False)
+        self.register_buffer('cached_freqs', self._precompute_cache(max_seq_len), persistent=False)
+
+    def _precompute_cache(self, seq_len: int):
+        seq = torch.arange(seq_len, device=self.inv_freq.device, dtype=self.inv_freq.dtype)
+        freqs = einsum('i, j -> i j', seq, self.inv_freq)
+        if self.interleaved:
+            freqs = repeat(freqs, '... n -> ... (n r)', r=2)
+        else:
+            freqs = torch.cat((freqs, freqs), dim=-1)
+        return freqs
+
+    def forward(self, seq_len: int) -> Tensor:
+        if seq_len > self.cached_freqs_seq_len:
+            raise RuntimeError("sequence exceeds RoPE max_seq_len!")
+        return self.cached_freqs[0: seq_len].detach()
+
+    def rotate_queries_or_keys(self, t: Tensor) -> Tensor:
+        device, dtype, seq_len = t.device, t.dtype, t.shape[-2]
+        freqs = self.forward(seq_len=seq_len)
+
+        return apply_rotary_emb(freqs.to(device=device, dtype=dtype), t, self.interleaved)

modules/compat.py

@@ -0,0 +1,24 @@
+def get_backbone_type(root_config: dict, nested_config: dict = None):
+    if nested_config is None:
+        nested_config = root_config
+    return nested_config.get(
+        'backbone_type',
+        root_config.get(
+            'backbone_type',
+            root_config.get('diff_decoder_type', 'wavenet')
+        )
+    )
+
+
+def get_backbone_args(config: dict, backbone_type: str):
+    args = config.get('backbone_args')
+    if args is not None:
+        return args
+    elif backbone_type == 'wavenet':
+        return {
+            'num_layers': config.get('residual_layers'),
+            'num_channels': config.get('residual_channels'),
+            'dilation_cycle_length': config.get('dilation_cycle_length'),
+        }
+    else:
+        return None

modules/core/__init__.py

@@ -0,0 +1,2 @@
+from .ddpm import GaussianDiffusion, PitchDiffusion, MultiVarianceDiffusion
+from .reflow import RectifiedFlow, PitchRectifiedFlow, MultiVarianceRectifiedFlow

modules/core/ddpm.py

@@ -0,0 +1,509 @@
+from __future__ import annotations
+
+from collections import deque
+from functools import partial
+from typing import List, Tuple
+
+import numpy as np
+import torch
+from torch import nn
+from tqdm import tqdm
+
+from modules.backbones import build_backbone
+from utils.hparams import hparams
+
+
+def extract(a, t, x_shape):
+    b, *_ = t.shape
+    out = a.gather(-1, t)
+    return out.reshape(b, *((1,) * (len(x_shape) - 1)))
+
+
+def noise_like(shape, device, repeat=False):
+    repeat_noise = lambda: torch.randn((1, *shape[1:]), device=device).repeat(shape[0], *((1,) * (len(shape) - 1)))
+    noise = lambda: torch.randn(shape, device=device)
+    return repeat_noise() if repeat else noise()
+
+
+def linear_beta_schedule(timesteps, max_beta=0.01):
+    """
+    linear schedule
+    """
+    betas = np.linspace(1e-4, max_beta, timesteps)
+    return betas
+
+
+def cosine_beta_schedule(timesteps, s=0.008):
+    """
+    cosine schedule
+    as proposed in https://openreview.net/forum?id=-NEXDKk8gZ
+    """
+    steps = timesteps + 1
+    x = np.linspace(0, steps, steps)
+    alphas_cumprod = np.cos(((x / steps) + s) / (1 + s) * np.pi * 0.5) ** 2
+    alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
+    betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
+    return np.clip(betas, a_min=0, a_max=0.999)
+
+
+beta_schedule = {
+    "cosine": cosine_beta_schedule,
+    "linear": linear_beta_schedule,
+}
+
+
+class GaussianDiffusion(nn.Module):
+    def __init__(self, out_dims, num_feats=1, timesteps=1000, k_step=1000,
+                 backbone_type=None, backbone_args=None, betas=None,
+                 spec_min=None, spec_max=None):
+        super().__init__()
+        self.denoise_fn: nn.Module = build_backbone(out_dims, num_feats, backbone_type, backbone_args)
+        self.out_dims = out_dims
+        self.num_feats = num_feats
+
+        if betas is not None:
+            betas = betas.detach().cpu().numpy() if isinstance(betas, torch.Tensor) else betas
+        else:
+            schedule_args = {}
+            if hparams['schedule_type'] == 'linear':
+                schedule_args['max_beta'] = hparams.get('max_beta', 0.01)
+            betas = beta_schedule[hparams['schedule_type']](timesteps, **schedule_args)
+
+        alphas = 1. - betas
+        alphas_cumprod = np.cumprod(alphas, axis=0)
+        alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1])
+
+        self.use_shallow_diffusion = hparams.get('use_shallow_diffusion', False)
+        if self.use_shallow_diffusion:
+            assert k_step <= timesteps, 'K_step should not be larger than timesteps.'
+        self.timesteps = timesteps
+        self.k_step = k_step if self.use_shallow_diffusion else timesteps
+        self.noise_list = deque(maxlen=4)
+
+        to_torch = partial(torch.tensor, dtype=torch.float32)
+
+        self.register_buffer('betas', to_torch(betas))
+        self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
+        self.register_buffer('alphas_cumprod_prev', to_torch(alphas_cumprod_prev))
+
+        # calculations for diffusion q(x_t | x_{t-1}) and others
+        self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod)))
+        self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod)))
+        self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod)))
+        self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod)))
+        self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod - 1)))
+
+        # calculations for posterior q(x_{t-1} | x_t, x_0)
+        posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod)
+        # above: equal to 1. / (1. / (1. - alpha_cumprod_tm1) + alpha_t / beta_t)
+        self.register_buffer('posterior_variance', to_torch(posterior_variance))
+        # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain
+        self.register_buffer('posterior_log_variance_clipped', to_torch(np.log(np.maximum(posterior_variance, 1e-20))))
+        self.register_buffer('posterior_mean_coef1', to_torch(
+            betas * np.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod)))
+        self.register_buffer('posterior_mean_coef2', to_torch(
+            (1. - alphas_cumprod_prev) * np.sqrt(alphas) / (1. - alphas_cumprod)))
+
+        # spec: [B, T, M] or [B, F, T, M]
+        # spec_min and spec_max: [1, 1, M] or [1, 1, F, M] => transpose(-3, -2) => [1, 1, M] or [1, F, 1, M]
+        spec_min = torch.FloatTensor(spec_min)[None, None, :out_dims].transpose(-3, -2)
+        spec_max = torch.FloatTensor(spec_max)[None, None, :out_dims].transpose(-3, -2)
+        self.register_buffer('spec_min', spec_min)
+        self.register_buffer('spec_max', spec_max)
+
+        # for compatibility with ONNX continuous acceleration
+        self.time_scale_factor = self.timesteps
+        self.t_start = 1 - self.k_step / self.timesteps
+        factors = torch.LongTensor([i for i in range(1, self.timesteps + 1) if self.timesteps % i == 0])
+        self.register_buffer('timestep_factors', factors, persistent=False)
+
+    def q_mean_variance(self, x_start, t):
+        mean = extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
+        variance = extract(1. - self.alphas_cumprod, t, x_start.shape)
+        log_variance = extract(self.log_one_minus_alphas_cumprod, t, x_start.shape)
+        return mean, variance, log_variance
+
+    def predict_start_from_noise(self, x_t, t, noise):
+        return (
+                extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t -
+                extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise
+        )
+
+    def q_posterior(self, x_start, x_t, t):
+        posterior_mean = (
+                extract(self.posterior_mean_coef1, t, x_t.shape) * x_start +
+                extract(self.posterior_mean_coef2, t, x_t.shape) * x_t
+        )
+        posterior_variance = extract(self.posterior_variance, t, x_t.shape)
+        posterior_log_variance_clipped = extract(self.posterior_log_variance_clipped, t, x_t.shape)
+        return posterior_mean, posterior_variance, posterior_log_variance_clipped
+
+    def p_mean_variance(self, x, t, cond):
+        noise_pred = self.denoise_fn(x, t, cond=cond)
+        x_recon = self.predict_start_from_noise(x, t=t, noise=noise_pred)
+
+        # This is previously inherited from original DiffSinger repository
+        # and disabled due to some loudness issues when speedup = 1.
+        # x_recon.clamp_(-1., 1.)
+
+        model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t)
+        return model_mean, posterior_variance, posterior_log_variance
+
+    @torch.no_grad()
+    def p_sample(self, x, t, cond, clip_denoised=True, repeat_noise=False):
+        b, *_, device = *x.shape, x.device
+        model_mean, _, model_log_variance = self.p_mean_variance(x=x, t=t, cond=cond)
+        noise = noise_like(x.shape, device, repeat_noise)
+        # no noise when t == 0
+        nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))
+        return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise
+
+    @torch.no_grad()
+    def p_sample_ddim(self, x, t, interval, cond):
+        a_t = extract(self.alphas_cumprod, t, x.shape)
+        a_prev = extract(self.alphas_cumprod, torch.max(t - interval, torch.zeros_like(t)), x.shape)
+
+        noise_pred = self.denoise_fn(x, t, cond=cond)
+        x_prev = a_prev.sqrt() * (
+                x / a_t.sqrt() + (((1 - a_prev) / a_prev).sqrt() - ((1 - a_t) / a_t).sqrt()) * noise_pred
+        )
+        return x_prev
+
+    @torch.no_grad()
+    def p_sample_plms(self, x, t, interval, cond, clip_denoised=True, repeat_noise=False):
+        """
+        Use the PLMS method from
+        [Pseudo Numerical Methods for Diffusion Models on Manifolds](https://arxiv.org/abs/2202.09778).
+        """
+
+        def get_x_pred(x, noise_t, t):
+            a_t = extract(self.alphas_cumprod, t, x.shape)
+            a_prev = extract(self.alphas_cumprod, torch.max(t - interval, torch.zeros_like(t)), x.shape)
+            a_t_sq, a_prev_sq = a_t.sqrt(), a_prev.sqrt()
+
+            x_delta = (a_prev - a_t) * ((1 / (a_t_sq * (a_t_sq + a_prev_sq))) * x - 1 / (
+                    a_t_sq * (((1 - a_prev) * a_t).sqrt() + ((1 - a_t) * a_prev).sqrt())) * noise_t)
+            x_pred = x + x_delta
+
+            return x_pred
+
+        noise_list = self.noise_list
+        noise_pred = self.denoise_fn(x, t, cond=cond)
+
+        if len(noise_list) == 0:
+            x_pred = get_x_pred(x, noise_pred, t)
+            noise_pred_prev = self.denoise_fn(x_pred, max(t - interval, 0), cond=cond)
+            noise_pred_prime = (noise_pred + noise_pred_prev) / 2
+        elif len(noise_list) == 1:
+            noise_pred_prime = (3 * noise_pred - noise_list[-1]) / 2
+        elif len(noise_list) == 2:
+            noise_pred_prime = (23 * noise_pred - 16 * noise_list[-1] + 5 * noise_list[-2]) / 12
+        else:
+            noise_pred_prime = (55 * noise_pred - 59 * noise_list[-1] + 37 * noise_list[-2] - 9 * noise_list[-3]) / 24
+
+        x_prev = get_x_pred(x, noise_pred_prime, t)
+        noise_list.append(noise_pred)
+
+        return x_prev
+
+    def q_sample(self, x_start, t, noise):
+        return (
+                extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +
+                extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
+        )
+
+    def p_losses(self, x_start, t, cond, noise=None):
+        if noise is None:
+            noise = torch.randn_like(x_start)
+
+        x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
+        x_recon = self.denoise_fn(x_noisy, t, cond)
+
+        return x_recon, noise
+
+    def inference(self, cond, b=1, x_start=None, device=None):
+        depth = hparams.get('K_step_infer', self.k_step)
+        speedup = hparams['diff_speedup']
+        if speedup > 0:
+            assert depth % speedup == 0, f'Acceleration ratio must be a factor of diffusion depth {depth}.'
+
+        noise = torch.randn(b, self.num_feats, self.out_dims, cond.shape[2], device=device)
+        if self.use_shallow_diffusion:
+            t_max = min(depth, self.k_step)
+        else:
+            t_max = self.k_step
+
+        if t_max >= self.timesteps:
+            x = noise
+        elif t_max > 0:
+            assert x_start is not None, 'Missing shallow diffusion source.'
+            x = self.q_sample(
+                x_start, torch.full((b,), t_max - 1, device=device, dtype=torch.long), noise
+            )
+        else:
+            assert x_start is not None, 'Missing shallow diffusion source.'
+            x = x_start
+
+        if speedup > 1 and t_max > 0:
+            algorithm = hparams['diff_accelerator']
+            if algorithm == 'dpm-solver':
+                from inference.dpm_solver_pytorch import NoiseScheduleVP, model_wrapper, DPM_Solver
+                # 1. Define the noise schedule.
+                noise_schedule = NoiseScheduleVP(schedule='discrete', betas=self.betas[:t_max])
+
+                # 2. Convert your discrete-time `model` to the continuous-time
+                # noise prediction model. Here is an example for a diffusion model
+                # `model` with the noise prediction type ("noise") .
+                def my_wrapper(fn):
+                    def wrapped(x, t, **kwargs):
+                        ret = fn(x, t, **kwargs)
+                        self.bar.update(1)
+                        return ret
+
+                    return wrapped
+
+                model_fn = model_wrapper(
+                    my_wrapper(self.denoise_fn),
+                    noise_schedule,
+                    model_type="noise",  # or "x_start" or "v" or "score"
+                    model_kwargs={"cond": cond}
+                )
+
+                # 3. Define dpm-solver and sample by singlestep DPM-Solver.
+                # (We recommend singlestep DPM-Solver for unconditional sampling)
+                # You can adjust the `steps` to balance the computation
+                # costs and the sample quality.
+                dpm_solver = DPM_Solver(model_fn, noise_schedule, algorithm_type="dpmsolver++")
+
+                steps = t_max // hparams["diff_speedup"]
+                self.bar = tqdm(desc="sample time step", total=steps, disable=not hparams['infer'], leave=False)
+                x = dpm_solver.sample(
+                    x,
+                    steps=steps,
+                    order=2,
+                    skip_type="time_uniform",
+                    method="multistep",
+                )
+                self.bar.close()
+            elif algorithm == 'unipc':
+                from inference.uni_pc import NoiseScheduleVP, model_wrapper, UniPC
+                # 1. Define the noise schedule.
+                noise_schedule = NoiseScheduleVP(schedule='discrete', betas=self.betas[:t_max])
+
+                # 2. Convert your discrete-time `model` to the continuous-time
+                # noise prediction model. Here is an example for a diffusion model
+                # `model` with the noise prediction type ("noise") .
+                def my_wrapper(fn):
+                    def wrapped(x, t, **kwargs):
+                        ret = fn(x, t, **kwargs)
+                        self.bar.update(1)
+                        return ret
+
+                    return wrapped
+
+                model_fn = model_wrapper(
+                    my_wrapper(self.denoise_fn),
+                    noise_schedule,
+                    model_type="noise",  # or "x_start" or "v" or "score"
+                    model_kwargs={"cond": cond}
+                )
+
+                # 3. Define uni_pc and sample by multistep UniPC.
+                # You can adjust the `steps` to balance the computation
+                # costs and the sample quality.
+                uni_pc = UniPC(model_fn, noise_schedule, variant='bh2')
+
+                steps = t_max // hparams["diff_speedup"]
+                self.bar = tqdm(desc="sample time step", total=steps, disable=not hparams['infer'], leave=False)
+                x = uni_pc.sample(
+                    x,
+                    steps=steps,
+                    order=2,
+                    skip_type="time_uniform",
+                    method="multistep",
+                )
+                self.bar.close()
+            elif algorithm == 'pndm':
+                self.noise_list = deque(maxlen=4)
+                iteration_interval = speedup
+                for i in tqdm(
+                        reversed(range(0, t_max, iteration_interval)), desc='sample time step',
+                        total=t_max // iteration_interval, disable=not hparams['infer'], leave=False
+                ):
+                    x = self.p_sample_plms(
+                        x, torch.full((b,), i, device=device, dtype=torch.long),
+                        iteration_interval, cond=cond
+                    )
+            elif algorithm == 'ddim':
+                iteration_interval = speedup
+                for i in tqdm(
+                        reversed(range(0, t_max, iteration_interval)), desc='sample time step',
+                        total=t_max // iteration_interval, disable=not hparams['infer'], leave=False
+                ):
+                    x = self.p_sample_ddim(
+                        x, torch.full((b,), i, device=device, dtype=torch.long),
+                        iteration_interval, cond=cond
+                    )
+            else:
+                raise ValueError(f"Unsupported acceleration algorithm for DDPM: {algorithm}.")
+        else:
+            for i in tqdm(reversed(range(0, t_max)), desc='sample time step', total=t_max,
+                          disable=not hparams['infer'], leave=False):
+                x = self.p_sample(x, torch.full((b,), i, device=device, dtype=torch.long), cond)
+        x = x.transpose(2, 3).squeeze(1)  # [B, F, M, T] => [B, T, M] or [B, F, T, M]
+        return x
+
+    def forward(self, condition, gt_spec=None, src_spec=None, infer=True):
+        """
+            conditioning diffusion, use fastspeech2 encoder output as the condition
+        """
+        cond = condition.transpose(1, 2)
+        b, device = condition.shape[0], condition.device
+
+        if not infer:
+            # gt_spec: [B, T, M] or [B, F, T, M]
+            spec = self.norm_spec(gt_spec).transpose(-2, -1)  # [B, M, T] or [B, F, M, T]
+            if self.num_feats == 1:
+                spec = spec[:, None, :, :]  # [B, F=1, M, T]
+            t = torch.randint(0, self.k_step, (b,), device=device).long()
+            x_recon, noise = self.p_losses(spec, t, cond=cond)
+            return x_recon, noise
+        else:
+            # src_spec: [B, T, M] or [B, F, T, M]
+            if src_spec is not None:
+                spec = self.norm_spec(src_spec).transpose(-2, -1)
+                if self.num_feats == 1:
+                    spec = spec[:, None, :, :]
+            else:
+                spec = None
+            x = self.inference(cond, b=b, x_start=spec, device=device)
+            return self.denorm_spec(x)
+
+    def norm_spec(self, x):
+        return (x - self.spec_min) / (self.spec_max - self.spec_min) * 2 - 1
+
+    def denorm_spec(self, x):
+        return (x + 1) / 2 * (self.spec_max - self.spec_min) + self.spec_min
+
+
+class RepetitiveDiffusion(GaussianDiffusion):
+    def __init__(self, vmin: float | int | list, vmax: float | int | list,
+                 repeat_bins: int, timesteps=1000, k_step=1000,
+                 backbone_type=None, backbone_args=None,
+                 betas=None):
+        assert (isinstance(vmin, (float, int)) and isinstance(vmax, (float, int))) or len(vmin) == len(vmax)
+        num_feats = 1 if isinstance(vmin, (float, int)) else len(vmin)
+        spec_min = [vmin] if num_feats == 1 else [[v] for v in vmin]
+        spec_max = [vmax] if num_feats == 1 else [[v] for v in vmax]
+        self.repeat_bins = repeat_bins
+        super().__init__(
+            out_dims=repeat_bins, num_feats=num_feats,
+            timesteps=timesteps, k_step=k_step,
+            backbone_type=backbone_type, backbone_args=backbone_args,
+            betas=betas, spec_min=spec_min, spec_max=spec_max
+        )
+
+    def norm_spec(self, x):
+        """
+
+        :param x: [B, T] or [B, F, T]
+        :return [B, T, R] or [B, F, T, R]
+        """
+        if self.num_feats == 1:
+            repeats = [1, 1, self.repeat_bins]
+        else:
+            repeats = [1, 1, 1, self.repeat_bins]
+        return super().norm_spec(x.unsqueeze(-1).repeat(repeats))
+
+    def denorm_spec(self, x):
+        """
+
+        :param x: [B, T, R] or [B, F, T, R]
+        :return [B, T] or [B, F, T]
+        """
+        return super().denorm_spec(x).mean(dim=-1)
+
+
+class PitchDiffusion(RepetitiveDiffusion):
+    def __init__(self, vmin: float, vmax: float,
+                 cmin: float, cmax: float, repeat_bins,
+                 timesteps=1000, k_step=1000,
+                 backbone_type=None, backbone_args=None,
+                 betas=None):
+        self.vmin = vmin  # norm min
+        self.vmax = vmax  # norm max
+        self.cmin = cmin  # clip min
+        self.cmax = cmax  # clip max
+        super().__init__(
+            vmin=vmin, vmax=vmax, repeat_bins=repeat_bins,
+            timesteps=timesteps, k_step=k_step,
+            backbone_type=backbone_type, backbone_args=backbone_args,
+            betas=betas
+        )
+
+    def norm_spec(self, x):
+        return super().norm_spec(x.clamp(min=self.cmin, max=self.cmax))
+
+    def denorm_spec(self, x):
+        return super().denorm_spec(x).clamp(min=self.cmin, max=self.cmax)
+
+
+class MultiVarianceDiffusion(RepetitiveDiffusion):
+    def __init__(
+            self, ranges: List[Tuple[float, float]],
+            clamps: List[Tuple[float | None, float | None] | None],
+            repeat_bins, timesteps=1000, k_step=1000,
+            backbone_type=None, backbone_args=None,
+            betas=None
+    ):
+        assert len(ranges) == len(clamps)
+        self.clamps = clamps
+        vmin = [r[0] for r in ranges]
+        vmax = [r[1] for r in ranges]
+        if len(vmin) == 1:
+            vmin = vmin[0]
+        if len(vmax) == 1:
+            vmax = vmax[0]
+        super().__init__(
+            vmin=vmin, vmax=vmax, repeat_bins=repeat_bins,
+            timesteps=timesteps, k_step=k_step,
+            backbone_type=backbone_type, backbone_args=backbone_args,
+            betas=betas
+        )
+
+    def clamp_spec(self, xs: list | tuple):
+        clamped = []
+        for x, c in zip(xs, self.clamps):
+            if c is None:
+                clamped.append(x)
+                continue
+            clamped.append(x.clamp(min=c[0], max=c[1]))
+        return clamped
+
+    def norm_spec(self, xs: list | tuple):
+        """
+
+        :param xs: sequence of [B, T]
+        :return: [B, F, T] => super().norm_spec(xs) => [B, F, T, R]
+        """
+        assert len(xs) == self.num_feats
+        clamped = self.clamp_spec(xs)
+        xs = torch.stack(clamped, dim=1)  # [B, F, T]
+        if self.num_feats == 1:
+            xs = xs.squeeze(1)  # [B, T]
+        return super().norm_spec(xs)
+
+    def denorm_spec(self, xs):
+        """
+
+        :param xs: [B, T, R] or [B, F, T, R] => super().denorm_spec(xs) => [B, T] or [B, F, T]
+        :return: sequence of [B, T]
+        """
+        xs = super().denorm_spec(xs)
+        if self.num_feats == 1:
+            xs = [xs]
+        else:
+            xs = xs.unbind(dim=1)
+        assert len(xs) == self.num_feats
+        return self.clamp_spec(xs)
+

modules/core/reflow.py

@@ -0,0 +1,262 @@
+from __future__ import annotations
+
+from typing import List, Tuple
+
+import torch
+import torch.nn as nn
+from tqdm import tqdm
+
+from modules.backbones import build_backbone
+from utils.hparams import hparams
+
+
+class RectifiedFlow(nn.Module):
+    def __init__(self, out_dims, num_feats=1, t_start=0., time_scale_factor=1000,
+                 backbone_type=None, backbone_args=None,
+                 spec_min=None, spec_max=None):
+        super().__init__()
+        self.velocity_fn: nn.Module = build_backbone(out_dims, num_feats, backbone_type, backbone_args)
+        self.out_dims = out_dims
+        self.num_feats = num_feats
+        self.use_shallow_diffusion = hparams.get('use_shallow_diffusion', False)
+        if self.use_shallow_diffusion:
+            assert 0. <= t_start <= 1., 'T_start should be in [0, 1].'
+        else:
+            t_start = 0.
+        self.t_start = t_start
+        self.time_scale_factor = time_scale_factor
+
+        # spec: [B, T, M] or [B, F, T, M]
+        # spec_min and spec_max: [1, 1, M] or [1, 1, F, M] => transpose(-3, -2) => [1, 1, M] or [1, F, 1, M]
+        spec_min = torch.FloatTensor(spec_min)[None, None, :out_dims].transpose(-3, -2)
+        spec_max = torch.FloatTensor(spec_max)[None, None, :out_dims].transpose(-3, -2)
+        self.register_buffer('spec_min', spec_min, persistent=False)
+        self.register_buffer('spec_max', spec_max, persistent=False)
+
+    def p_losses(self, x_end, t, cond):
+        x_start = torch.randn_like(x_end)
+        x_t = x_start + t[:, None, None, None] * (x_end - x_start)
+        v_pred = self.velocity_fn(x_t, t * self.time_scale_factor, cond)
+
+        return v_pred, x_end - x_start
+
+    def forward(self, condition, gt_spec=None, src_spec=None, infer=True):
+        cond = condition.transpose(1, 2)
+        b, device = condition.shape[0], condition.device
+
+        if not infer:
+            # gt_spec: [B, T, M] or [B, F, T, M]
+            spec = self.norm_spec(gt_spec).transpose(-2, -1)  # [B, M, T] or [B, F, M, T]
+            if self.num_feats == 1:
+                spec = spec[:, None, :, :]  # [B, F=1, M, T]
+            t = self.t_start + (1.0 - self.t_start) * torch.rand((b,), device=device)
+            v_pred, v_gt = self.p_losses(spec, t, cond=cond)
+            return v_pred, v_gt, t
+        else:
+            # src_spec: [B, T, M] or [B, F, T, M]
+            if src_spec is not None:
+                spec = self.norm_spec(src_spec).transpose(-2, -1)
+                if self.num_feats == 1:
+                    spec = spec[:, None, :, :]
+            else:
+                spec = None
+            x = self.inference(cond, b=b, x_end=spec, device=device)
+            return self.denorm_spec(x)
+
+    @torch.no_grad()
+    def sample_euler(self, x, t, dt, cond):
+        x += self.velocity_fn(x, self.time_scale_factor * t, cond) * dt
+        t += dt
+        return x, t
+
+    @torch.no_grad()
+    def sample_rk2(self, x, t, dt, cond):
+        k_1 = self.velocity_fn(x, self.time_scale_factor * t, cond)
+        k_2 = self.velocity_fn(x + 0.5 * k_1 * dt, self.time_scale_factor * (t + 0.5 * dt), cond)
+        x += k_2 * dt
+        t += dt
+        return x, t
+
+    @torch.no_grad()
+    def sample_rk4(self, x, t, dt, cond):
+        k_1 = self.velocity_fn(x, self.time_scale_factor * t, cond)
+        k_2 = self.velocity_fn(x + 0.5 * k_1 * dt, self.time_scale_factor * (t + 0.5 * dt), cond)
+        k_3 = self.velocity_fn(x + 0.5 * k_2 * dt, self.time_scale_factor * (t + 0.5 * dt), cond)
+        k_4 = self.velocity_fn(x + k_3 * dt, self.time_scale_factor * (t + dt), cond)
+        x += (k_1 + 2 * k_2 + 2 * k_3 + k_4) * dt / 6
+        t += dt
+        return x, t
+
+    @torch.no_grad()
+    def sample_rk5(self, x, t, dt, cond):
+        k_1 = self.velocity_fn(x, self.time_scale_factor * t, cond)
+        k_2 = self.velocity_fn(x + 0.25 * k_1 * dt, self.time_scale_factor * (t + 0.25 * dt), cond)
+        k_3 = self.velocity_fn(x + 0.125 * (k_2 + k_1) * dt, self.time_scale_factor * (t + 0.25 * dt), cond)
+        k_4 = self.velocity_fn(x + 0.5 * (-k_2 + 2 * k_3) * dt, self.time_scale_factor * (t + 0.5 * dt), cond)
+        k_5 = self.velocity_fn(x + 0.0625 * (3 * k_1 + 9 * k_4) * dt, self.time_scale_factor * (t + 0.75 * dt), cond)
+        k_6 = self.velocity_fn(x + (-3 * k_1 + 2 * k_2 + 12 * k_3 - 12 * k_4 + 8 * k_5) * dt / 7,
+                               self.time_scale_factor * (t + dt),
+                               cond)
+        x += (7 * k_1 + 32 * k_3 + 12 * k_4 + 32 * k_5 + 7 * k_6) * dt / 90
+        t += dt
+        return x, t
+
+    @torch.no_grad()
+    def inference(self, cond, b=1, x_end=None, device=None):
+        noise = torch.randn(b, self.num_feats, self.out_dims, cond.shape[2], device=device)
+        t_start = hparams.get('T_start_infer', self.t_start)
+        if self.use_shallow_diffusion and t_start > 0:
+            assert x_end is not None, 'Missing shallow diffusion source.'
+            if t_start >= 1.:
+                t_start = 1.
+                x = x_end
+            else:
+                x = t_start * x_end + (1 - t_start) * noise
+        else:
+            t_start = 0.
+            x = noise
+
+        algorithm = hparams['sampling_algorithm']
+        infer_step = hparams['sampling_steps']
+
+        if t_start < 1:
+            dt = (1.0 - t_start) / max(1, infer_step)
+            algorithm_fn = {
+                'euler': self.sample_euler,
+                'rk2': self.sample_rk2,
+                'rk4': self.sample_rk4,
+                'rk5': self.sample_rk5,
+            }.get(algorithm)
+            if algorithm_fn is None:
+                raise ValueError(f'Unsupported algorithm for Rectified Flow: {algorithm}.')
+            dts = torch.tensor([dt]).to(x)
+            for i in tqdm(range(infer_step), desc='sample time step', total=infer_step,
+                          disable=not hparams['infer'], leave=False):
+                x, _ = algorithm_fn(x, t_start + i * dts, dt, cond)
+            x = x.float()
+        x = x.transpose(2, 3).squeeze(1)  # [B, F, M, T] => [B, T, M] or [B, F, T, M]
+        return x
+
+    def norm_spec(self, x):
+        return (x - self.spec_min) / (self.spec_max - self.spec_min) * 2 - 1
+
+    def denorm_spec(self, x):
+        return (x + 1) / 2 * (self.spec_max - self.spec_min) + self.spec_min
+
+
+class RepetitiveRectifiedFlow(RectifiedFlow):
+    def __init__(self, vmin: float | int | list, vmax: float | int | list,
+                 repeat_bins: int, time_scale_factor=1000,
+                 backbone_type=None, backbone_args=None):
+        assert (isinstance(vmin, (float, int)) and isinstance(vmax, (float, int))) or len(vmin) == len(vmax)
+        num_feats = 1 if isinstance(vmin, (float, int)) else len(vmin)
+        spec_min = [vmin] if num_feats == 1 else [[v] for v in vmin]
+        spec_max = [vmax] if num_feats == 1 else [[v] for v in vmax]
+        self.repeat_bins = repeat_bins
+        super().__init__(
+            out_dims=repeat_bins, num_feats=num_feats,
+            time_scale_factor=time_scale_factor,
+            backbone_type=backbone_type, backbone_args=backbone_args,
+            spec_min=spec_min, spec_max=spec_max
+        )
+
+    def norm_spec(self, x):
+        """
+
+        :param x: [B, T] or [B, F, T]
+        :return [B, T, R] or [B, F, T, R]
+        """
+        if self.num_feats == 1:
+            repeats = [1, 1, self.repeat_bins]
+        else:
+            repeats = [1, 1, 1, self.repeat_bins]
+        return super().norm_spec(x.unsqueeze(-1).repeat(repeats))
+
+    def denorm_spec(self, x):
+        """
+
+        :param x: [B, T, R] or [B, F, T, R]
+        :return [B, T] or [B, F, T]
+        """
+        return super().denorm_spec(x).mean(dim=-1)
+
+
+class PitchRectifiedFlow(RepetitiveRectifiedFlow):
+    def __init__(self, vmin: float, vmax: float,
+                 cmin: float, cmax: float, repeat_bins,
+                 time_scale_factor=1000,
+                 backbone_type=None, backbone_args=None):
+        self.vmin = vmin  # norm min
+        self.vmax = vmax  # norm max
+        self.cmin = cmin  # clip min
+        self.cmax = cmax  # clip max
+        super().__init__(
+            vmin=vmin, vmax=vmax, repeat_bins=repeat_bins,
+            time_scale_factor=time_scale_factor,
+            backbone_type=backbone_type, backbone_args=backbone_args
+        )
+
+    def norm_spec(self, x):
+        return super().norm_spec(x.clamp(min=self.cmin, max=self.cmax))
+
+    def denorm_spec(self, x):
+        return super().denorm_spec(x).clamp(min=self.cmin, max=self.cmax)
+
+
+class MultiVarianceRectifiedFlow(RepetitiveRectifiedFlow):
+    def __init__(
+            self, ranges: List[Tuple[float, float]],
+            clamps: List[Tuple[float | None, float | None] | None],
+            repeat_bins, time_scale_factor=1000,
+            backbone_type=None, backbone_args=None
+    ):
+        assert len(ranges) == len(clamps)
+        self.clamps = clamps
+        vmin = [r[0] for r in ranges]
+        vmax = [r[1] for r in ranges]
+        if len(vmin) == 1:
+            vmin = vmin[0]
+        if len(vmax) == 1:
+            vmax = vmax[0]
+        super().__init__(
+            vmin=vmin, vmax=vmax, repeat_bins=repeat_bins,
+            time_scale_factor=time_scale_factor,
+            backbone_type=backbone_type, backbone_args=backbone_args
+        )
+
+    def clamp_spec(self, xs: list | tuple):
+        clamped = []
+        for x, c in zip(xs, self.clamps):
+            if c is None:
+                clamped.append(x)
+                continue
+            clamped.append(x.clamp(min=c[0], max=c[1]))
+        return clamped
+
+    def norm_spec(self, xs: list | tuple):
+        """
+
+        :param xs: sequence of [B, T]
+        :return: [B, F, T] => super().norm_spec(xs) => [B, F, T, R]
+        """
+        assert len(xs) == self.num_feats
+        clamped = self.clamp_spec(xs)
+        xs = torch.stack(clamped, dim=1)  # [B, F, T]
+        if self.num_feats == 1:
+            xs = xs.squeeze(1)  # [B, T]
+        return super().norm_spec(xs)
+
+    def denorm_spec(self, xs):
+        """
+
+        :param xs: [B, T, R] or [B, F, T, R] => super().denorm_spec(xs) => [B, T] or [B, F, T]
+        :return: sequence of [B, T]
+        """
+        xs = super().denorm_spec(xs)
+        if self.num_feats == 1:
+            xs = [xs]
+        else:
+            xs = xs.unbind(dim=1)
+        assert len(xs) == self.num_feats
+        return self.clamp_spec(xs)
+

modules/fastspeech/acoustic_encoder.py

@@ -0,0 +1,185 @@
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+from modules.commons.common_layers import (
+    NormalInitEmbedding as Embedding,
+    SinusoidalPosEmb,
+    AdamWLinear,
+)
+from modules.fastspeech.tts_modules import FastSpeech2Encoder, mel2ph_to_dur, StretchRegulator
+from utils.hparams import hparams
+from utils.phoneme_utils import PAD_INDEX
+
+
+class FastSpeech2Acoustic(nn.Module):
+    def __init__(self, vocab_size):
+        super().__init__()
+        self.txt_embed = Embedding(vocab_size, hparams['hidden_size'], PAD_INDEX)
+        self.use_lang_id = hparams.get('use_lang_id', False)
+        if self.use_lang_id:
+            self.lang_embed = Embedding(hparams['num_lang'] + 1, hparams['hidden_size'], padding_idx=0)
+
+        self.use_stretch_embed = hparams.get('use_stretch_embed', False)
+        if self.use_stretch_embed:
+            self.sr = StretchRegulator()
+            self.stretch_embed = nn.Sequential(
+                SinusoidalPosEmb(hparams['hidden_size']),
+                nn.Linear(hparams['hidden_size'], hparams['hidden_size'] * 4),
+                nn.GELU(),
+                nn.Linear(hparams['hidden_size'] * 4, hparams['hidden_size']),
+            )
+            self.stretch_embed_rnn = nn.GRU(hparams['hidden_size'], hparams['hidden_size'], 1, batch_first=True)
+            self._stretch_embed_rnn_flattened = False
+
+        self.dur_embed = AdamWLinear(1, hparams['hidden_size'])
+        self.use_mix_ln = hparams.get('use_mix_ln', False)
+        if self.use_mix_ln:
+            self.mix_ln_layer = hparams['mix_ln_layer']
+        else:
+            self.mix_ln_layer = []
+        self.encoder = FastSpeech2Encoder(
+            hidden_size=hparams['hidden_size'], num_layers=hparams['enc_layers'],
+            ffn_kernel_size=hparams['enc_ffn_kernel_size'], ffn_act=hparams['ffn_act'],
+            dropout=hparams['dropout'], num_heads=hparams['num_heads'],
+            use_pos_embed=hparams['use_pos_embed'], rel_pos=hparams.get('rel_pos', False),
+            use_rope=hparams.get('use_rope', False), rope_interleaved=hparams.get('rope_interleaved', True),
+            mix_ln_layer=self.mix_ln_layer
+        )
+
+        self.pitch_embed = AdamWLinear(1, hparams['hidden_size'])
+        self.variance_embed_list = []
+        self.use_energy_embed = hparams.get('use_energy_embed', False)
+        self.use_breathiness_embed = hparams.get('use_breathiness_embed', False)
+        self.use_voicing_embed = hparams.get('use_voicing_embed', False)
+        self.use_tension_embed = hparams.get('use_tension_embed', False)
+        if self.use_energy_embed:
+            self.variance_embed_list.append('energy')
+        if self.use_breathiness_embed:
+            self.variance_embed_list.append('breathiness')
+        if self.use_voicing_embed:
+            self.variance_embed_list.append('voicing')
+        if self.use_tension_embed:
+            self.variance_embed_list.append('tension')
+
+        self.use_variance_embeds = len(self.variance_embed_list) > 0
+        if self.use_variance_embeds:
+            self.variance_embeds = nn.ModuleDict({
+                v_name: AdamWLinear(1, hparams['hidden_size'])
+                for v_name in self.variance_embed_list
+            })
+
+        self.use_variance_scaling = hparams.get('use_variance_scaling', False)
+        if self.use_variance_scaling:
+            self.variance_scaling_factor = {
+                'energy': 1. / 96,  # 96 dB — max dynamic range of 16-bit audio
+                'breathiness': 1. / 96,
+                'voicing': 1. / 96,
+                'tension': 0.1,  # 1 / 10; tension logits are roughly [-10, 10]
+                'key_shift': 1. / 12,  # one octave — max key shift in most editors
+                'speed': 1.
+            }
+        else:
+            self.variance_scaling_factor = {
+                'energy': 1.,
+                'breathiness': 1.,
+                'voicing': 1.,
+                'tension': 1.,
+                'key_shift': 1.,
+                'speed': 1.
+            }
+
+        self.use_key_shift_embed = hparams.get('use_key_shift_embed', False)
+        if self.use_key_shift_embed:
+            self.key_shift_embed = AdamWLinear(1, hparams['hidden_size'])
+
+        self.use_speed_embed = hparams.get('use_speed_embed', False)
+        if self.use_speed_embed:
+            self.speed_embed = AdamWLinear(1, hparams['hidden_size'])
+
+        self.use_spk_id = hparams['use_spk_id']
+        if self.use_spk_id:
+            self.spk_embed = Embedding(hparams['num_spk'], hparams['hidden_size'])
+
+    def forward_variance_embedding(self, condition, key_shift=None, speed=None, **variances):
+        if self.use_variance_embeds:
+            variance_embeds = torch.stack([
+                self.variance_embeds[v_name](variances[v_name][:, :, None] * self.variance_scaling_factor[v_name])
+                for v_name in self.variance_embed_list
+            ], dim=-1).sum(-1)
+            condition += variance_embeds
+
+        if self.use_key_shift_embed:
+            key_shift_embed = self.key_shift_embed(key_shift[:, :, None] * self.variance_scaling_factor['key_shift'])
+            condition += key_shift_embed
+
+        if self.use_speed_embed:
+            speed_embed = self.speed_embed(speed[:, :, None] * self.variance_scaling_factor['speed'])
+            condition += speed_embed
+
+        return condition
+
+    def forward(
+            self, txt_tokens, mel2ph, f0,
+            key_shift=None, speed=None,
+            spk_embed_id=None, languages=None,
+            **kwargs
+    ):
+        spk_embed = None
+        if self.use_spk_id:
+            spk_mix_embed = kwargs.get('spk_mix_embed')
+            if spk_mix_embed is not None:
+                spk_embed = spk_mix_embed
+            else:
+                spk_embed = self.spk_embed(spk_embed_id)[:, None, :]
+        txt_embed = self.txt_embed(txt_tokens)
+        dur = mel2ph_to_dur(mel2ph, txt_tokens.shape[1])
+        if self.use_variance_scaling:
+            dur_embed = self.dur_embed(torch.log(1 + dur[:, :, None].float()))
+        else:
+            dur_embed = self.dur_embed(dur[:, :, None].float())
+        if self.use_lang_id:
+            lang_embed = self.lang_embed(languages)
+            extra_embed = dur_embed + lang_embed
+        else:
+            extra_embed = dur_embed
+        encoder_out = self.encoder(txt_embed, extra_embed, txt_tokens == 0, spk_embed)
+
+        encoder_out = F.pad(encoder_out, [0, 0, 1, 0])
+        mel2ph_ = mel2ph[..., None].repeat([1, 1, encoder_out.shape[-1]])
+        condition = torch.gather(encoder_out, 1, mel2ph_)
+
+        if self.use_stretch_embed:
+            stretch = torch.round(1000 * self.sr(mel2ph, dur))
+            if self.training and stretch.numel() > 1000:
+                # construct a phoneme stretching index lookup table with a total of 1001 indexes (0~1000)
+                table = self.stretch_embed(torch.arange(0, 1001, device=stretch.device))
+                stretch_embed = torch.index_select(table, 0, stretch.view(-1).long()).view_as(condition)
+            else:
+                stretch_embed = self.stretch_embed(stretch)
+            condition += stretch_embed
+            # flatten_parameters fuses the GRU weights into a contiguous buffer for cuDNN.
+            # It only needs to happen once after weight init, device change, or load_state_dict.
+            # We guard with a flag to avoid the redundant call on every forward.
+            # Limitation: the flag lives on this module and is invisible to PyTorch. After
+            # load_state_dict() or model.to(device) replaces the GRU weights, the flag stays
+            # True and flatten_parameters is skipped — cuDNN will fall back to the slower path.
+            # To restore the fast path, reset the flag manually: model._stretch_embed_rnn_flattened = False
+            if not self._stretch_embed_rnn_flattened:
+                self.stretch_embed_rnn.flatten_parameters()
+                self._stretch_embed_rnn_flattened = True
+            stretch_embed_rnn_out, _ = self.stretch_embed_rnn(condition)
+            condition = condition + stretch_embed_rnn_out
+
+        if self.use_spk_id:
+            condition += spk_embed
+
+        f0_mel = (1 + f0 / 700).log()
+        pitch_embed = self.pitch_embed(f0_mel[:, :, None])
+        condition += pitch_embed
+
+        condition = self.forward_variance_embedding(
+            condition, key_shift=key_shift, speed=speed, **kwargs
+        )
+
+        return condition

modules/fastspeech/param_adaptor.py

@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+import torch
+
+import modules.compat as compat
+from modules.core.ddpm import MultiVarianceDiffusion
+from utils import filter_kwargs
+from utils.hparams import hparams
+
+VARIANCE_CHECKLIST = ['energy', 'breathiness', 'voicing', 'tension']
+
+
+class ParameterAdaptorModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.variance_prediction_list = []
+        self.predict_energy = hparams.get('predict_energy', False)
+        self.predict_breathiness = hparams.get('predict_breathiness', False)
+        self.predict_voicing = hparams.get('predict_voicing', False)
+        self.predict_tension = hparams.get('predict_tension', False)
+        if self.predict_energy:
+            self.variance_prediction_list.append('energy')
+        if self.predict_breathiness:
+            self.variance_prediction_list.append('breathiness')
+        if self.predict_voicing:
+            self.variance_prediction_list.append('voicing')
+        if self.predict_tension:
+            self.variance_prediction_list.append('tension')
+        self.predict_variances = len(self.variance_prediction_list) > 0
+
+    def build_adaptor(self, cls=MultiVarianceDiffusion):
+        ranges = []
+        clamps = []
+
+        if self.predict_energy:
+            ranges.append((
+                hparams['energy_db_min'],
+                hparams['energy_db_max']
+            ))
+            clamps.append((hparams['energy_db_min'], 0.))
+
+        if self.predict_breathiness:
+            ranges.append((
+                hparams['breathiness_db_min'],
+                hparams['breathiness_db_max']
+            ))
+            clamps.append((hparams['breathiness_db_min'], 0.))
+
+        if self.predict_voicing:
+            ranges.append((
+                hparams['voicing_db_min'],
+                hparams['voicing_db_max']
+            ))
+            clamps.append((hparams['voicing_db_min'], 0.))
+
+        if self.predict_tension:
+            ranges.append((
+                hparams['tension_logit_min'],
+                hparams['tension_logit_max']
+            ))
+            clamps.append((
+                hparams['tension_logit_min'],
+                hparams['tension_logit_max']
+            ))
+
+        variances_hparams = hparams['variances_prediction_args']
+        total_repeat_bins = variances_hparams['total_repeat_bins']
+        assert total_repeat_bins % len(self.variance_prediction_list) == 0, \
+            f'Total number of repeat bins must be divisible by number of ' \
+            f'variance parameters ({len(self.variance_prediction_list)}).'
+        repeat_bins = total_repeat_bins // len(self.variance_prediction_list)
+        backbone_type = compat.get_backbone_type(hparams, nested_config=variances_hparams)
+        backbone_args = compat.get_backbone_args(variances_hparams, backbone_type=backbone_type)
+        kwargs = filter_kwargs(
+            {
+                'ranges': ranges,
+                'clamps': clamps,
+                'repeat_bins': repeat_bins,
+                'timesteps': hparams.get('timesteps'),
+                'time_scale_factor': hparams.get('time_scale_factor'),
+                'backbone_type': backbone_type,
+                'backbone_args': backbone_args
+            },
+            cls
+        )
+        return cls(**kwargs)
+
+    def collect_variance_inputs(self, **kwargs) -> list:
+        return [kwargs.get(name) for name in self.variance_prediction_list]
+
+    def collect_variance_outputs(self, variances: list | tuple) -> dict:
+        return {
+            name: pred
+            for name, pred in zip(self.variance_prediction_list, variances)
+        }

modules/fastspeech/tts_modules.py

@@ -0,0 +1,455 @@
+import math
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from modules.commons.rotary_embedding_torch import RotaryEmbedding
+from modules.commons.common_layers import SinusoidalPositionalEmbedding, EncSALayer, AdamWLinear
+from modules.commons.espnet_positional_embedding import RelPositionalEncoding
+
+DEFAULT_MAX_SOURCE_POSITIONS = 2000
+DEFAULT_MAX_TARGET_POSITIONS = 2000
+
+
+class TransformerEncoderLayer(nn.Module):
+    def __init__(self, hidden_size, dropout, kernel_size=None, act='gelu', num_heads=2, rotary_embed=None,
+                 layer_idx=None, mix_ln_layer=None):
+        super().__init__()
+        self.op = EncSALayer(
+            hidden_size, num_heads, dropout=dropout,
+            attention_dropout=0.0, relu_dropout=dropout,
+            kernel_size=kernel_size,
+            act=act, rotary_embed=rotary_embed,
+            layer_idx=layer_idx, mix_ln_layer=mix_ln_layer
+        )
+
+    def forward(self, x, **kwargs):
+        return self.op(x, **kwargs)
+
+
+######################
+# fastspeech modules
+######################
+class LayerNorm(torch.nn.LayerNorm):
+    """Layer normalization module.
+    :param int nout: output dim size
+    :param int dim: dimension to be normalized
+    """
+
+    def __init__(self, nout, dim=-1):
+        """Construct an LayerNorm object."""
+        super(LayerNorm, self).__init__(nout, eps=1e-12)
+        self.dim = dim
+
+    def forward(self, x):
+        """Apply layer normalization.
+        :param torch.Tensor x: input tensor
+        :return: layer normalized tensor
+        :rtype torch.Tensor
+        """
+        if self.dim == -1:
+            return super(LayerNorm, self).forward(x)
+        return super(LayerNorm, self).forward(x.transpose(1, -1)).transpose(1, -1)
+
+
+class DurationPredictor(torch.nn.Module):
+    """Duration predictor module.
+    This is a module of duration predictor described in `FastSpeech: Fast, Robust and Controllable Text to Speech`_.
+    The duration predictor predicts a duration of each frame in log domain from the hidden embeddings of encoder.
+    .. _`FastSpeech: Fast, Robust and Controllable Text to Speech`:
+        https://arxiv.org/pdf/1905.09263.pdf
+    Note:
+        The calculation domain of outputs is different between in `forward` and in `inference`. In `forward`,
+        the outputs are calculated in log domain but in `inference`, those are calculated in linear domain.
+    """
+
+    def __init__(self, in_dims, n_layers=2, n_chans=384, kernel_size=3,
+                 dropout_rate=0.1, offset=1.0, dur_loss_type='mse', arch='resnet'):
+        """Initialize duration predictor module.
+        Args:
+            in_dims (int): Input dimension.
+            n_layers (int, optional): Number of convolutional layers.
+            n_chans (int, optional): Number of channels of convolutional layers.
+            kernel_size (int, optional): Kernel size of convolutional layers.
+            dropout_rate (float, optional): Dropout rate.
+            offset (float, optional): Offset value to avoid nan in log domain.
+        """
+        super(DurationPredictor, self).__init__()
+        self.offset = offset
+        self.conv = torch.nn.ModuleList()
+        self.kernel_size = kernel_size
+        self.use_resnet = (arch == 'resnet')
+        for idx in range(n_layers):
+            in_chans = in_dims if idx == 0 else n_chans
+            if self.use_resnet:
+                self.conv.append(nn.Sequential(
+                    LayerNorm(in_chans, dim=1),
+                    nn.Conv1d(in_chans, n_chans, kernel_size, stride=1, padding=kernel_size // 2),
+                    nn.ReLU(),
+                    nn.Conv1d(n_chans, n_chans, 1),
+                    nn.Dropout(dropout_rate)
+                ))
+            else:
+                self.conv.append(nn.Sequential(
+                    nn.Identity(),  # this is a placeholder for ConstantPad1d which is now merged into Conv1d
+                    nn.Conv1d(in_chans, n_chans, kernel_size, stride=1, padding=kernel_size // 2),
+                    nn.ReLU(),
+                    LayerNorm(n_chans, dim=1),
+                    nn.Dropout(dropout_rate)
+                ))
+        if self.use_resnet and in_dims != n_chans:
+            self.res_conv = nn.Conv1d(in_dims, n_chans, 1)
+        else:
+            self.res_conv = None
+        self.loss_type = dur_loss_type
+        if self.loss_type in ['mse', 'huber']:
+            self.out_dims = 1
+        # elif hparams['dur_loss_type'] == 'mog':
+        #     out_dims = 15
+        # elif hparams['dur_loss_type'] == 'crf':
+        #     out_dims = 32
+        #     from torchcrf import CRF
+        #     self.crf = CRF(out_dims, batch_first=True)
+        else:
+            raise NotImplementedError()
+        self.linear = AdamWLinear(n_chans, self.out_dims)
+
+    def out2dur(self, xs):
+        if self.loss_type in ['mse', 'huber']:
+            # NOTE: calculate loss in log domain
+            dur = xs.squeeze(-1).exp() - self.offset  # (B, Tmax)
+        # elif hparams['dur_loss_type'] == 'crf':
+        #     dur = torch.LongTensor(self.crf.decode(xs)).cuda()
+        else:
+            raise NotImplementedError()
+        return dur
+
+    def forward(self, xs, x_masks=None, infer=True):
+        """Calculate forward propagation.
+        Args:
+            xs (Tensor): Batch of input sequences (B, Tmax, idim).
+            x_masks (BoolTensor, optional): Batch of masks indicating padded part (B, Tmax).
+            infer (bool): Whether inference
+        Returns:
+            (train) FloatTensor, (infer) LongTensor: Batch of predicted durations in linear domain (B, Tmax).
+        """
+        xs = xs.transpose(1, -1)  # (B, idim, Tmax)
+        masks = 1 - x_masks.float()
+        masks_ = masks[:, None, :]
+        for idx, f in enumerate(self.conv):
+            if self.use_resnet:
+                residual = self.res_conv(xs) if idx == 0 and self.res_conv is not None else xs
+                xs = residual + f(xs)
+            else:
+                xs = f(xs)
+            if x_masks is not None:
+                xs = xs * masks_
+        xs = self.linear(xs.transpose(1, -1))  # [B, T, C]
+        xs = xs * masks[:, :, None]  # (B, T, C)
+
+        dur_pred = self.out2dur(xs)
+        if infer:
+            dur_pred = dur_pred.clamp(min=0.)  # avoid negative value
+        return dur_pred
+
+
+class VariancePredictor(torch.nn.Module):
+    def __init__(self, vmin, vmax, in_dims,
+                 n_layers=5, n_chans=512, kernel_size=5,
+                 dropout_rate=0.1):
+        """Initialize variance predictor module.
+        Args:
+            in_dims (int): Input dimension.
+            n_layers (int, optional): Number of convolutional layers.
+            n_chans (int, optional): Number of channels of convolutional layers.
+            kernel_size (int, optional): Kernel size of convolutional layers.
+            dropout_rate (float, optional): Dropout rate.
+        """
+        super(VariancePredictor, self).__init__()
+
+        self.vmin = vmin
+        self.vmax = vmax
+        self.conv = torch.nn.ModuleList()
+        self.kernel_size = kernel_size
+        for idx in range(n_layers):
+            in_chans = in_dims if idx == 0 else n_chans
+            self.conv.append(torch.nn.Sequential(
+                torch.nn.Conv1d(in_chans, n_chans, kernel_size, stride=1, padding=kernel_size // 2),
+                torch.nn.ReLU(),
+                LayerNorm(n_chans, dim=1),
+                torch.nn.Dropout(dropout_rate)
+            ))
+        self.linear = torch.nn.Linear(n_chans, 1)
+        self.embed_positions = SinusoidalPositionalEmbedding(in_dims, 0, init_size=4096)
+        self.pos_embed_alpha = nn.Parameter(torch.Tensor([1]))
+
+    def out2value(self, xs):
+        return (xs + 1) / 2 * (self.vmax - self.vmin) + self.vmin
+
+    def forward(self, xs, infer=True):
+        """
+        :param xs: [B, T, H]
+        :param infer: whether inference
+        :return: [B, T]
+        """
+        positions = self.pos_embed_alpha * self.embed_positions(xs[..., 0])
+        xs = xs + positions
+        xs = xs.transpose(1, -1)  # (B, idim, Tmax)
+        for f in self.conv:
+            xs = f(xs)  # (B, C, Tmax)
+        xs = self.linear(xs.transpose(1, -1)).squeeze(-1)  # (B, Tmax)
+        if infer:
+            xs = self.out2value(xs)
+        return xs
+
+
+class PitchPredictor(torch.nn.Module):
+    def __init__(self, vmin, vmax, num_bins, deviation,
+                 in_dims, n_layers=5, n_chans=384, kernel_size=5,
+                 dropout_rate=0.1):
+        """Initialize pitch predictor module.
+        Args:
+            in_dims (int): Input dimension.
+            n_layers (int, optional): Number of convolutional layers.
+            n_chans (int, optional): Number of channels of convolutional layers.
+            kernel_size (int, optional): Kernel size of convolutional layers.
+            dropout_rate (float, optional): Dropout rate.
+        """
+        super(PitchPredictor, self).__init__()
+        self.vmin = vmin
+        self.vmax = vmax
+        self.interval = (vmax - vmin) / (num_bins - 1)  # align with centers of bins
+        self.sigma = deviation / self.interval
+        self.register_buffer('x', torch.arange(num_bins).float().reshape(1, 1, -1))  # [1, 1, N]
+
+        self.base_pitch_embed = torch.nn.Linear(1, in_dims)
+        self.conv = torch.nn.ModuleList()
+        self.kernel_size = kernel_size
+        for idx in range(n_layers):
+            in_chans = in_dims if idx == 0 else n_chans
+            self.conv.append(torch.nn.Sequential(
+                torch.nn.Conv1d(in_chans, n_chans, kernel_size, stride=1, padding=kernel_size // 2),
+                torch.nn.ReLU(),
+                LayerNorm(n_chans, dim=1),
+                torch.nn.Dropout(dropout_rate)
+            ))
+        self.linear = torch.nn.Linear(n_chans, num_bins)
+        self.embed_positions = SinusoidalPositionalEmbedding(in_dims, 0, init_size=4096)
+        self.pos_embed_alpha = nn.Parameter(torch.Tensor([1]))
+
+    def bins_to_values(self, bins):
+        return bins * self.interval + self.vmin
+
+    def out2pitch(self, probs):
+        logits = probs.sigmoid()  # [B, T, N]
+        # return logits
+        # logits_sum = logits.sum(dim=2)  # [B, T]
+        bins = torch.sum(self.x * logits, dim=2) / torch.sum(logits, dim=2)  # [B, T]
+        pitch = self.bins_to_values(bins)
+        # uv = logits_sum / (self.sigma * math.sqrt(2 * math.pi)) < 0.3
+        # pitch[uv] = torch.nan
+        return pitch
+
+    def forward(self, xs, base):
+        """
+        :param xs: [B, T, H]
+        :param base: [B, T]
+        :return: [B, T, N]
+        """
+        xs = xs + self.base_pitch_embed(base[..., None])
+        positions = self.pos_embed_alpha * self.embed_positions(xs[..., 0])
+        xs = xs + positions
+        xs = xs.transpose(1, -1)  # (B, idim, Tmax)
+        for f in self.conv:
+            xs = f(xs)  # (B, C, Tmax)
+        xs = self.linear(xs.transpose(1, -1))  # (B, Tmax, H)
+        return self.out2pitch(xs) + base, xs
+
+
+class RhythmRegulator(torch.nn.Module):
+    def __init__(self, eps=1e-5):
+        super().__init__()
+        self.eps = eps
+
+    def forward(self, ph_dur, ph2word, word_dur):
+        """
+        Example (no batch dim version):
+            1. ph_dur = [4,2,3,2]
+            2. word_dur = [3,4,2], ph2word = [1,2,2,3]
+            3. word_dur_in = [4,5,2]
+            4. alpha_w = [0.75,0.8,1], alpha_ph = [0.75,0.8,0.8,1]
+            5. ph_dur_out = [3,1.6,2.4,2]
+        :param ph_dur: [B, T_ph]
+        :param ph2word: [B, T_ph]
+        :param word_dur: [B, T_w]
+        """
+        ph_dur = ph_dur.float() * (ph2word > 0)
+        word_dur = word_dur.float()
+        word_dur_in = ph_dur.new_zeros(ph_dur.shape[0], ph2word.max() + 1).scatter_add(
+            1, ph2word, ph_dur
+        )[:, 1:]  # [B, T_ph] => [B, T_w]
+        alpha_w = word_dur / word_dur_in.clamp(min=self.eps)  # avoid dividing by zero
+        alpha_ph = torch.gather(F.pad(alpha_w, [1, 0]), 1, ph2word)  # [B, T_w] => [B, T_ph]
+        ph_dur_out = ph_dur * alpha_ph
+        return ph_dur_out.round().long()
+
+
+class LengthRegulator(torch.nn.Module):
+    # noinspection PyMethodMayBeStatic
+    def forward(self, dur, dur_padding=None, alpha=None):
+        """
+        Example (no batch dim version):
+            1. dur = [2,2,3]
+            2. token_idx = [[1],[2],[3]], dur_cumsum = [2,4,7], dur_cumsum_prev = [0,2,4]
+            3. token_mask = [[1,1,0,0,0,0,0],
+                             [0,0,1,1,0,0,0],
+                             [0,0,0,0,1,1,1]]
+            4. token_idx * token_mask = [[1,1,0,0,0,0,0],
+                                         [0,0,2,2,0,0,0],
+                                         [0,0,0,0,3,3,3]]
+            5. (token_idx * token_mask).sum(0) = [1,1,2,2,3,3,3]
+
+        :param dur: Batch of durations of each frame (B, T_txt)
+        :param dur_padding: Batch of padding of each frame (B, T_txt)
+        :param alpha: duration rescale coefficient
+        :return:
+            mel2ph (B, T_speech)
+        """
+        assert alpha is None or alpha > 0
+        if alpha is not None:
+            dur = torch.round(dur.float() * alpha).long()
+        if dur_padding is not None:
+            dur = dur * (1 - dur_padding.long())
+        token_idx = torch.arange(1, dur.shape[1] + 1)[None, :, None].to(dur.device)
+        dur_cumsum = torch.cumsum(dur, 1)
+        dur_cumsum_prev = F.pad(dur_cumsum, [1, -1], mode='constant', value=0)
+
+        pos_idx = torch.arange(dur.sum(-1).max())[None, None].to(dur.device)
+        token_mask = (pos_idx >= dur_cumsum_prev[:, :, None]) & (pos_idx < dur_cumsum[:, :, None])
+        mel2ph = (token_idx * token_mask.long()).sum(1)
+        return mel2ph
+
+
+class StretchRegulator(torch.nn.Module):
+    # noinspection PyMethodMayBeStatic
+    def forward(self, mel2ph, dur=None):
+        """
+        Example (no batch dim version):
+            1. dur = [2,4,3]
+            2. mel2ph = [1,1,2,2,2,2,3,3,3]
+            3. mel2dur = [2,2,4,4,4,4,3,3,3]
+            4. bound_mask = [0,1,0,0,0,1,0,0,1]
+            5. 1 - bound_mask * mel2dur = [1,-1,1,1,1,-3,1,1,-2] => pad => [0,1,-1,1,1,1,-3,1,1]
+            6. stretch_denorm = [0,1,0,1,2,3,0,1,2]
+
+        :param dur: Batch of durations of each frame (B, T_txt)
+        :param mel2ph: Batch of mel2ph (B, T_speech)
+        :return:
+            stretch (B, T_speech)
+        """
+        if dur is None:
+            dur = mel2ph_to_dur(mel2ph, mel2ph.max())
+        dur = torch.cat([torch.ones_like(dur[:, :1]), dur], dim=1)  # Avoid dividing by zero
+        mel2dur = torch.gather(dur, 1, mel2ph)
+        bound_mask = torch.gt(mel2ph[:, 1:], mel2ph[:, :-1])
+        stretch_delta = 1 - bound_mask * mel2dur[:, :-1]
+        stretch_delta = F.pad(stretch_delta, [1, 0])
+        stretch_denorm = torch.cumsum(stretch_delta, dim=1)
+        stretch = stretch_denorm.float() / mel2dur
+        return stretch * (mel2ph > 0)
+
+
+def mel2ph_to_dur(mel2ph, T_txt, max_dur=None):
+    B, _ = mel2ph.shape
+    dur = mel2ph.new_zeros(B, T_txt + 1).scatter_add(1, mel2ph, torch.ones_like(mel2ph))
+    dur = dur[:, 1:]
+    if max_dur is not None:
+        dur = dur.clamp(max=max_dur)
+    return dur
+
+
+class FastSpeech2Encoder(nn.Module):
+    def __init__(
+            self, hidden_size, num_layers,
+            ffn_kernel_size=9, ffn_act='gelu',
+            dropout=None, num_heads=2, use_pos_embed=True, rel_pos=True,
+            use_rope=False, rope_interleaved=True, mix_ln_layer=None
+    ):
+        super().__init__()
+        self.num_layers = num_layers
+        embed_dim = self.hidden_size = hidden_size
+        self.dropout = dropout
+        self.use_pos_embed = use_pos_embed
+        if use_pos_embed and use_rope:
+            if embed_dim % (num_heads * 2) != 0:
+                raise ValueError(
+                    "RoPE requires the hidden size to be multiple of "
+                    f"num_heads * 2 = {num_heads * 2}, but got {embed_dim}."
+                )
+            rotary_embed = RotaryEmbedding(dim=embed_dim // num_heads, interleaved=rope_interleaved)
+        else:
+            rotary_embed = None
+        self.layers = nn.ModuleList([
+            TransformerEncoderLayer(
+                self.hidden_size, self.dropout,
+                kernel_size=ffn_kernel_size, act=ffn_act,
+                num_heads=num_heads, rotary_embed=rotary_embed,
+                layer_idx=i, mix_ln_layer=mix_ln_layer
+            )
+            for i in range(self.num_layers)
+        ])
+        self.layer_norm = nn.LayerNorm(embed_dim)
+
+        self.embed_scale = math.sqrt(hidden_size)
+        self.padding_idx = 0
+        self.rel_pos = rel_pos
+        if use_rope:
+            self.embed_positions = None
+        elif self.rel_pos:
+            self.embed_positions = RelPositionalEncoding(hidden_size, dropout_rate=0.0)
+        else:
+            self.embed_positions = SinusoidalPositionalEmbedding(
+                hidden_size, self.padding_idx, init_size=DEFAULT_MAX_TARGET_POSITIONS,
+            )
+
+    def forward_embedding(self, main_embed, extra_embed=None, padding_mask=None):
+        # embed tokens and positions
+        x = self.embed_scale * main_embed
+        if extra_embed is not None:
+            x = x + extra_embed
+        if self.use_pos_embed and self.embed_positions is not None:
+            if self.rel_pos:
+                x = self.embed_positions(x)
+            else:
+                positions = self.embed_positions(~padding_mask)
+                x = x + positions
+        x = F.dropout(x, p=self.dropout, training=self.training)
+        return x
+
+    def forward(self, main_embed, extra_embed, padding_mask, spk_embed=None, attn_mask=None, return_hiddens=False):
+        x = self.forward_embedding(main_embed, extra_embed, padding_mask=padding_mask)  # [B, T, H]
+        nonpadding_mask_BT = 1 - padding_mask.float()[:, :, None]  # [B, T, 1]
+
+        # NOTICE:
+        # The following codes are commented out because
+        # `self.use_pos_embed` is always False in the older versions,
+        # and this argument did not compat with `hparams['use_pos_embed']`,
+        # which defaults to True. The new version fixed this inconsistency,
+        # resulting in temporary removal of pos_embed_alpha, which has actually
+        # never been used before.
+
+        # if self.use_pos_embed:
+        #     positions = self.pos_embed_alpha * self.embed_positions(x[..., 0])
+        #     x = x + positions
+        #     x = F.dropout(x, p=self.dropout, training=self.training)
+
+        x = x * nonpadding_mask_BT
+        hiddens = []
+        for layer in self.layers:
+            x = layer(x, encoder_padding_mask=padding_mask, cond=spk_embed, attn_mask=attn_mask) * nonpadding_mask_BT
+            if return_hiddens:
+                hiddens.append(x)
+        x = self.layer_norm(x) * nonpadding_mask_BT
+        if return_hiddens:
+            x = torch.stack(hiddens, 0)  # [L, B, T, C]
+        return x

modules/fastspeech/variance_encoder.py

@@ -0,0 +1,158 @@
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+from modules.commons.common_layers import (
+    NormalInitEmbedding as Embedding,
+    XavierUniformInitLinear as Linear,
+    AdamWLinear,
+)
+from modules.fastspeech.tts_modules import FastSpeech2Encoder, DurationPredictor
+from utils.hparams import hparams
+from utils.phoneme_utils import PAD_INDEX
+
+
+class FastSpeech2Variance(nn.Module):
+    def __init__(self, vocab_size):
+        super().__init__()
+        self.predict_dur = hparams['predict_dur']
+        self.linguistic_mode = 'word' if hparams['predict_dur'] else 'phoneme'
+        self.use_lang_id = hparams['use_lang_id']
+        self.use_variance_scaling = hparams.get('use_variance_scaling', False)
+        self.txt_embed = Embedding(vocab_size, hparams['hidden_size'], PAD_INDEX)
+        if self.use_lang_id:
+            self.lang_embed = Embedding(hparams['num_lang'] + 1, hparams['hidden_size'], padding_idx=0)
+
+        if self.predict_dur:
+            self.onset_embed = Embedding(2, hparams['hidden_size'])
+            self.word_dur_embed = AdamWLinear(1, hparams['hidden_size'])
+        else:
+            self.ph_dur_embed = AdamWLinear(1, hparams['hidden_size'])
+
+        self.encoder = FastSpeech2Encoder(
+            hidden_size=hparams['hidden_size'], num_layers=hparams['enc_layers'],
+            ffn_kernel_size=hparams['enc_ffn_kernel_size'], ffn_act=hparams['ffn_act'],
+            dropout=hparams['dropout'], num_heads=hparams['num_heads'],
+            use_pos_embed=hparams['use_pos_embed'], rel_pos=hparams.get('rel_pos', False), 
+            use_rope=hparams.get('use_rope', False), rope_interleaved=hparams.get('rope_interleaved', True)
+        )
+
+        dur_hparams = hparams['dur_prediction_args']
+        if self.predict_dur:
+            self.midi_embed = Embedding(128, hparams['hidden_size'])
+            self.dur_predictor = DurationPredictor(
+                in_dims=hparams['hidden_size'],
+                n_chans=dur_hparams['hidden_size'],
+                n_layers=dur_hparams['num_layers'],
+                dropout_rate=dur_hparams['dropout'],
+                kernel_size=dur_hparams['kernel_size'],
+                offset=dur_hparams['log_offset'],
+                dur_loss_type=dur_hparams['loss_type'],
+                arch=dur_hparams['arch']
+            )
+
+    def forward(
+            self, txt_tokens, midi, ph2word,
+            ph_dur=None, word_dur=None,
+            spk_embed=None, languages=None,
+            infer=True
+    ):
+        """
+        :param txt_tokens: (train, infer) [B, T_ph]
+        :param midi: (train, infer) [B, T_ph]
+        :param ph2word: (train, infer) [B, T_ph]
+        :param ph_dur: (train, [infer]) [B, T_ph]
+        :param word_dur: (infer) [B, T_w]
+        :param spk_embed: (train) [B, T_ph, H]
+        :param languages (train, infer) [B, T_ph]
+        :param infer: whether inference
+        :return: encoder_out, ph_dur_pred
+        """
+        txt_embed = self.txt_embed(txt_tokens)
+        if self.linguistic_mode == 'word':
+            b = txt_tokens.shape[0]
+            onset = torch.diff(ph2word, dim=1, prepend=ph2word.new_zeros(b, 1)) > 0
+            onset_embed = self.onset_embed(onset.long())  # [B, T_ph, H]
+
+            if word_dur is None or not infer:
+                word_dur = ph_dur.new_zeros(b, ph2word.max() + 1).scatter_add(
+                    1, ph2word, ph_dur
+                )[:, 1:]  # [B, T_ph] => [B, T_w]
+            word_dur = torch.gather(F.pad(word_dur, [1, 0], value=0), 1, ph2word)  # [B, T_w] => [B, T_ph]
+            word_dur_embed = self.word_dur_embed(word_dur.float()[:, :, None])
+            extra_embed = onset_embed + word_dur_embed
+        elif self.use_variance_scaling:
+            extra_embed = self.ph_dur_embed(torch.log(1 + ph_dur.float())[:, :, None])
+        else:
+            extra_embed = self.ph_dur_embed(ph_dur.float()[:, :, None])
+            
+        if self.use_lang_id:
+            lang_embed = self.lang_embed(languages)
+            extra_embed += lang_embed
+        encoder_out = self.encoder(txt_embed, extra_embed, txt_tokens == 0)
+
+        if self.predict_dur:
+            midi_embed = self.midi_embed(midi)  # => [B, T_ph, H]
+            dur_cond = encoder_out + midi_embed
+            if spk_embed is not None:
+                dur_cond += spk_embed
+            ph_dur_pred = self.dur_predictor(dur_cond, x_masks=txt_tokens == PAD_INDEX, infer=infer)
+
+            return encoder_out, ph_dur_pred
+        else:
+            return encoder_out, None
+
+
+class MelodyEncoder(nn.Module):
+    def __init__(self, enc_hparams: dict):
+        super().__init__()
+
+        def get_hparam(key):
+            return enc_hparams.get(key, hparams.get(key))
+
+        # MIDI inputs
+        hidden_size = get_hparam('hidden_size')
+        self.use_variance_scaling = hparams.get('use_variance_scaling', False)
+        self.note_midi_embed = AdamWLinear(1, hidden_size)
+        self.note_dur_embed = AdamWLinear(1, hidden_size)
+
+        # ornament inputs
+        self.use_glide_embed = hparams['use_glide_embed']
+        self.glide_embed_scale = hparams['glide_embed_scale']
+        if self.use_glide_embed:
+            # 0: none, 1: up, 2: down
+            self.note_glide_embed = Embedding(len(hparams['glide_types']) + 1, hidden_size, padding_idx=0)
+
+        self.encoder = FastSpeech2Encoder(
+            hidden_size=hidden_size, num_layers=get_hparam('enc_layers'),
+            ffn_kernel_size=get_hparam('enc_ffn_kernel_size'), ffn_act=get_hparam('ffn_act'),
+            dropout=get_hparam('dropout'), num_heads=get_hparam('num_heads'),
+            use_pos_embed=get_hparam('use_pos_embed'), rel_pos=get_hparam('rel_pos'),
+            use_rope=get_hparam('use_rope'), rope_interleaved=hparams.get('rope_interleaved', True)
+        )
+        self.out_proj = Linear(hidden_size, hparams['hidden_size'])
+
+    def forward(self, note_midi, note_rest, note_dur, glide=None):
+        """
+        :param note_midi: float32 [B, T_n], -1: padding
+        :param note_rest: bool [B, T_n]
+        :param note_dur: int64 [B, T_n]
+        :param glide: int64 [B, T_n]
+        :return: [B, T_n, H]
+        """
+        if self.use_variance_scaling:
+            midi_embed = self.note_midi_embed(note_midi[:, :, None] / 128)
+            dur_embed = self.note_dur_embed(torch.log(1 + note_dur.float())[:, :, None])
+        else:
+            midi_embed = self.note_midi_embed(note_midi[:, :, None])
+            dur_embed = self.note_dur_embed(note_dur.float()[:, :, None])
+        midi_embed *= ~note_rest[:, :, None]
+        ornament_embed = 0
+        if self.use_glide_embed:
+            ornament_embed += self.note_glide_embed(glide) * self.glide_embed_scale
+        encoder_out = self.encoder(
+            midi_embed, dur_embed + ornament_embed,
+            padding_mask=note_midi < 0
+        )
+        encoder_out = self.out_proj(encoder_out)
+        return encoder_out

modules/hnsep/vr/__init__.py

@@ -0,0 +1,35 @@
+import pathlib
+
+import torch
+import yaml
+
+from .nets import CascadedNet
+
+
+class DotDict(dict):
+    def __getattr__(*args):
+        val = dict.get(*args)
+        return DotDict(val) if type(val) is dict else val
+
+    __setattr__ = dict.__setitem__
+    __delattr__ = dict.__delitem__
+
+
+def load_sep_model(model_path, device='cpu'):
+    model_path = pathlib.Path(model_path)
+    config_file = model_path.with_name('config.yaml')
+    with open(config_file, "r") as config:
+        args = yaml.safe_load(config)
+    args = DotDict(args)
+    model = CascadedNet(
+        args.n_fft,
+        args.hop_length,
+        args.n_out,
+        args.n_out_lstm,
+        True,
+        is_mono=args.is_mono
+    )
+    model.to(device)
+    model.load_state_dict(torch.load(model_path, map_location='cpu'))
+    model.eval()
+    return model

modules/hnsep/vr/layers.py

@@ -0,0 +1,166 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+
+def crop_center(h1, h2):
+    h1_shape = h1.size()
+    h2_shape = h2.size()
+
+    if h1_shape[3] == h2_shape[3]:
+        return h1
+    elif h1_shape[3] < h2_shape[3]:
+        raise ValueError('h1_shape[3] must be greater than h2_shape[3]')
+
+    # s_freq = (h2_shape[2] - h1_shape[2]) // 2
+    # e_freq = s_freq + h1_shape[2]
+    s_time = (h1_shape[3] - h2_shape[3]) // 2
+    e_time = s_time + h2_shape[3]
+    h1 = h1[:, :, :, s_time:e_time]
+
+    return h1
+
+
+class Conv2DBNActiv(nn.Module):
+
+    def __init__(self, nin, nout, ksize=3, stride=1, pad=1, dilation=1, activ=nn.ReLU):
+        super(Conv2DBNActiv, self).__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(
+                nin, nout,
+                kernel_size=ksize,
+                stride=stride,
+                padding=pad,
+                dilation=dilation,
+                bias=False
+            ),
+            nn.BatchNorm2d(nout),
+            activ()
+        )
+
+    def forward(self, x):
+        return self.conv(x)
+
+
+class Encoder(nn.Module):
+
+    def __init__(self, nin, nout, ksize=3, stride=1, pad=1, activ=nn.LeakyReLU):
+        super(Encoder, self).__init__()
+        self.conv1 = Conv2DBNActiv(nin, nout, ksize, stride, pad, activ=activ)
+        self.conv2 = Conv2DBNActiv(nout, nout, ksize, 1, pad, activ=activ)
+
+    def forward(self, x):
+        h = self.conv1(x)
+        h = self.conv2(h)
+
+        return h
+
+
+class Decoder(nn.Module):
+
+    def __init__(self, nin, nout, ksize=3, stride=1, pad=1, activ=nn.ReLU, dropout=False):
+        super(Decoder, self).__init__()
+        self.conv1 = Conv2DBNActiv(nin, nout, ksize, 1, pad, activ=activ)
+        # self.conv2 = Conv2DBNActiv(nout, nout, ksize, 1, pad, activ=activ)
+        self.dropout = nn.Dropout2d(0.1) if dropout else None
+
+    def forward(self, x, skip=None, fixed_length=True):
+        if fixed_length:
+            x = F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=True)
+        else:
+            _, _, h, w = x.size() 
+            x = F.pad(x, (0, 1, 0, 1), mode='replicate')
+            x = F.interpolate(x, size=(2*h+1,2*w+1), mode='bilinear', align_corners=True)
+            x = x[:, :, :-1, :-1]
+
+        if skip is not None:
+            skip = crop_center(skip, x)
+            x = torch.cat([x, skip], dim=1)
+
+        h = self.conv1(x)
+        # h = self.conv2(h)
+
+        if self.dropout is not None:
+            h = self.dropout(h)
+
+        return h
+
+
+class Mean(nn.Module):
+    def __init__(self, dim, keepdims=False):
+        super(Mean, self).__init__()
+        self.dim = dim
+        self.keepdims = keepdims
+
+    def forward(self, x):
+        return x.mean(self.dim, keepdims=self.keepdims)
+
+
+class ASPPModule(nn.Module):
+
+    def __init__(self, nin, nout, dilations=(4, 8, 12), activ=nn.ReLU, dropout=False):
+        super(ASPPModule, self).__init__()
+        self.conv1 = nn.Sequential(
+            Mean(dim=-2, keepdims=True),  # nn.AdaptiveAvgPool2d((1, None)),
+            Conv2DBNActiv(nin, nout, 1, 1, 0, activ=activ)
+        )
+        self.conv2 = Conv2DBNActiv(
+            nin, nout, 1, 1, 0, activ=activ
+        )
+        self.conv3 = Conv2DBNActiv(
+            nin, nout, 3, 1, dilations[0], dilations[0], activ=activ
+        )
+        self.conv4 = Conv2DBNActiv(
+            nin, nout, 3, 1, dilations[1], dilations[1], activ=activ
+        )
+        self.conv5 = Conv2DBNActiv(
+            nin, nout, 3, 1, dilations[2], dilations[2], activ=activ
+        )
+        self.bottleneck = Conv2DBNActiv(
+            nout * 5, nout, 1, 1, 0, activ=activ
+        )
+        self.dropout = nn.Dropout2d(0.1) if dropout else None
+
+    def forward(self, x):
+        _, _, h, w = x.size()
+        # feat1 = F.interpolate(self.conv1(x), size=(h, w), mode='bilinear', align_corners=True)
+        feat1 = self.conv1(x).repeat(1, 1, h, 1)
+        feat2 = self.conv2(x)
+        feat3 = self.conv3(x)
+        feat4 = self.conv4(x)
+        feat5 = self.conv5(x)
+        out = torch.cat((feat1, feat2, feat3, feat4, feat5), dim=1)
+        out = self.bottleneck(out)
+
+        if self.dropout is not None:
+            out = self.dropout(out)
+
+        return out
+
+
+class LSTMModule(nn.Module):
+
+    def __init__(self, nin_conv, nin_lstm, nout_lstm):
+        super(LSTMModule, self).__init__()
+        self.conv = Conv2DBNActiv(nin_conv, 1, 1, 1, 0)
+        self.lstm = nn.LSTM(
+            input_size=nin_lstm,
+            hidden_size=nout_lstm // 2,
+            bidirectional=True
+        )
+        self.dense = nn.Sequential(
+            nn.Linear(nout_lstm, nin_lstm),
+            nn.BatchNorm1d(nin_lstm),
+            nn.ReLU()
+        )
+
+    def forward(self, x):
+        N, _, nbins, nframes = x.size()
+        h = self.conv(x)[:, 0]  # N, nbins, nframes
+        h = h.permute(2, 0, 1)  # nframes, N, nbins
+        h, _ = self.lstm(h)
+        h = self.dense(h.reshape(-1, h.size()[-1]))  # nframes * N, nbins
+        h = h.reshape(nframes, N, 1, nbins)
+        h = h.permute(1, 2, 3, 0)
+
+        return h

modules/hnsep/vr/nets.py

@@ -0,0 +1,202 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from . import layers
+
+
+class BaseNet(nn.Module):
+
+    def __init__(self, nin, nout, nin_lstm, nout_lstm, dilations=((4, 2), (8, 4), (12, 6)), fixed_length=True):
+        super(BaseNet, self).__init__()
+        self.enc1 = layers.Conv2DBNActiv(nin, nout, 3, 1, 1)
+        self.enc2 = layers.Encoder(nout, nout * 2, 3, 2, 1)
+        self.enc3 = layers.Encoder(nout * 2, nout * 4, 3, 2, 1)
+        self.enc4 = layers.Encoder(nout * 4, nout * 6, 3, 2, 1)
+        self.enc5 = layers.Encoder(nout * 6, nout * 8, 3, 2, 1)
+
+        self.aspp = layers.ASPPModule(nout * 8, nout * 8, dilations, dropout=True)
+
+        self.dec4 = layers.Decoder(nout * (6 + 8), nout * 6, 3, 1, 1)
+        self.dec3 = layers.Decoder(nout * (4 + 6), nout * 4, 3, 1, 1)
+        self.dec2 = layers.Decoder(nout * (2 + 4), nout * 2, 3, 1, 1)
+        self.lstm_dec2 = layers.LSTMModule(nout * 2, nin_lstm, nout_lstm)
+        self.dec1 = layers.Decoder(nout * (1 + 2) + 1, nout * 1, 3, 1, 1)
+        
+        self.fixed_length = fixed_length
+
+    def __call__(self, x):
+        e1 = self.enc1(x)
+        e2 = self.enc2(e1)
+        e3 = self.enc3(e2)
+        e4 = self.enc4(e3)
+        e5 = self.enc5(e4)
+
+        h = self.aspp(e5)
+
+        h = self.dec4(h, e4, fixed_length=self.fixed_length)
+        h = self.dec3(h, e3, fixed_length=self.fixed_length)
+        h = self.dec2(h, e2, fixed_length=self.fixed_length)
+        h = torch.cat([h, self.lstm_dec2(h)], dim=1)
+        h = self.dec1(h, e1, fixed_length=self.fixed_length)
+
+        return h
+
+
+class CascadedNet(nn.Module):
+
+    def __init__(self, n_fft, hop_length, nout=32, nout_lstm=128, is_complex=False, is_mono=False, fixed_length=True):
+        super(CascadedNet, self).__init__()
+        self.n_fft = n_fft
+        self.hop_length = hop_length
+        self.seg_length = 32 * hop_length
+        self.is_complex = is_complex
+        self.is_mono = is_mono
+        self.register_buffer("window", torch.hann_window(n_fft), persistent=False)
+        self.max_bin = n_fft // 2
+        self.output_bin = n_fft // 2 + 1
+        self.nin_lstm = self.max_bin // 2
+        self.offset = 64
+
+        nin = 4 if is_complex else 2
+        if is_mono:
+            nin = nin // 2
+
+        self.stg1_low_band_net = nn.Sequential(
+            BaseNet(nin, nout // 2, self.nin_lstm // 2, nout_lstm, fixed_length=fixed_length),
+            layers.Conv2DBNActiv(nout // 2, nout // 4, 1, 1, 0)
+        )
+        self.stg1_high_band_net = BaseNet(
+            nin, nout // 4, self.nin_lstm // 2, nout_lstm // 2, fixed_length=fixed_length
+        )
+
+        self.stg2_low_band_net = nn.Sequential(
+            BaseNet(nout // 4 + nin, nout, self.nin_lstm // 2, nout_lstm, fixed_length=fixed_length),
+            layers.Conv2DBNActiv(nout, nout // 2, 1, 1, 0)
+        )
+        self.stg2_high_band_net = BaseNet(
+            nout // 4 + nin, nout // 2, self.nin_lstm // 2, nout_lstm // 2, fixed_length=fixed_length
+        )
+
+        self.stg3_full_band_net = BaseNet(
+            3 * nout // 4 + nin, nout, self.nin_lstm, nout_lstm, fixed_length=fixed_length
+        )
+
+        self.out = nn.Conv2d(nout, nin, 1, bias=False)
+        self.aux_out = nn.Conv2d(3 * nout // 4, nin, 1, bias=False)
+
+    def forward(self, x):
+        if self.is_complex:
+            x = torch.cat([x.real, x.imag], dim=1)
+
+        x = x[:, :, :self.max_bin]
+
+        bandw = x.size()[2] // 2
+        l1_in = x[:, :, :bandw]
+        h1_in = x[:, :, bandw:]
+        l1 = self.stg1_low_band_net(l1_in)
+        h1 = self.stg1_high_band_net(h1_in)
+        aux1 = torch.cat([l1, h1], dim=2)
+
+        l2_in = torch.cat([l1_in, l1], dim=1)
+        h2_in = torch.cat([h1_in, h1], dim=1)
+        l2 = self.stg2_low_band_net(l2_in)
+        h2 = self.stg2_high_band_net(h2_in)
+        aux2 = torch.cat([l2, h2], dim=2)
+
+        f3_in = torch.cat([x, aux1, aux2], dim=1)
+        f3 = self.stg3_full_band_net(f3_in)
+
+        if self.is_complex:
+            mask = self.out(f3)
+            if self.is_mono:
+                mask = torch.complex(mask[:, :1], mask[:, 1:])
+            else:
+                mask = torch.complex(mask[:, :2], mask[:, 2:])
+            mask = self.bounded_mask(mask)
+        else:
+            mask = torch.sigmoid(self.out(f3))
+
+        mask = F.pad(
+            input=mask,
+            pad=(0, 0, 0, self.output_bin - mask.size()[2]),
+            mode='replicate'
+        )
+
+        return mask
+
+    def bounded_mask(self, mask, eps=1e-8):
+        mask_mag = torch.abs(mask)
+        mask = torch.tanh(mask_mag) * mask / (mask_mag + eps)
+        return mask
+
+    def predict_mask(self, x):
+        mask = self.forward(x)
+
+        if self.offset > 0:
+            mask = mask[:, :, :, self.offset:-self.offset]
+            assert mask.size()[3] > 0
+
+        return mask
+
+    def predict(self, x):
+        mask = self.forward(x)
+        pred = x * mask
+
+        if self.offset > 0:
+            pred = pred[:, :, :, self.offset:-self.offset]
+            assert pred.size()[3] > 0
+
+        return pred
+
+    def audio2spec(self, x, use_pad=False):
+        B, C, T = x.shape
+        x = x.reshape(B * C, T)
+        if use_pad:
+            T1 = T + self.hop_length
+            T_pad = self.seg_length * ((T1 - 1) // self.seg_length + 1) - T1
+            nl_pad = T_pad // 2 // self.hop_length
+            Tl_pad = nl_pad * self.hop_length
+            x = F.pad(x, (Tl_pad, T_pad - Tl_pad))
+        spec = torch.stft(
+            x,
+            n_fft=self.n_fft,
+            hop_length=self.hop_length,
+            return_complex=True,
+            window=self.window,
+            pad_mode='constant'
+        )
+        spec = spec.reshape(B, C, spec.shape[-2], spec.shape[-1])
+        return spec
+
+    def spec2audio(self, x):
+        B, C, N, T = x.shape
+        x = x.reshape(-1, N, T)
+        x = torch.istft(x, self.n_fft, self.hop_length, window=self.window)
+        x = x.reshape(B, C, -1)
+        return x
+
+    def predict_from_audio(self, x):
+        B, C, T = x.shape
+        x = x.reshape(B * C, T)
+        T1 = T + self.hop_length
+        T_pad = self.seg_length * ((T1 - 1) // self.seg_length + 1) - T1       
+        nl_pad = T_pad // 2 // self.hop_length
+        Tl_pad = nl_pad * self.hop_length       
+        x = F.pad(x, (Tl_pad, T_pad - Tl_pad))
+        spec = torch.stft(
+            x,
+            n_fft=self.n_fft,
+            hop_length=self.hop_length,
+            return_complex=True,
+            window=self.window,
+            pad_mode='constant'
+        )
+        spec = spec.reshape(B, C, spec.shape[-2], spec.shape[-1])
+        mask = self.forward(spec)
+        spec_pred = spec * mask
+        spec_pred = spec_pred.reshape(B * C, spec.shape[-2], spec.shape[-1])
+        x_pred = torch.istft(spec_pred, self.n_fft, self.hop_length, window=self.window)
+        x_pred = x_pred[:, Tl_pad: Tl_pad + T]
+        x_pred = x_pred.reshape(B, C, T)
+        return x_pred

modules/losses/__init__.py

@@ -0,0 +1,3 @@
+from .diff_loss import DiffusionLoss
+from .reflow_loss import RectifiedFlowLoss
+from .dur_loss import DurationLoss

modules/losses/diff_loss.py

@@ -0,0 +1,34 @@
+import torch.nn as nn
+from torch import Tensor
+
+
+class DiffusionLoss(nn.Module):
+    def __init__(self, loss_type):
+        super().__init__()
+        self.loss_type = loss_type
+        if self.loss_type == 'l1':
+            self.loss = nn.L1Loss(reduction='none')
+        elif self.loss_type == 'l2':
+            self.loss = nn.MSELoss(reduction='none')
+        else:
+            raise NotImplementedError()
+
+    @staticmethod
+    def _mask_non_padding(x_recon, noise, non_padding=None):
+        if non_padding is not None:
+            non_padding = non_padding.transpose(1, 2).unsqueeze(1)
+            return x_recon * non_padding, noise * non_padding
+        else:
+            return x_recon, noise
+
+    def _forward(self, x_recon, noise):
+        return self.loss(x_recon, noise)
+
+    def forward(self, x_recon: Tensor, noise: Tensor, non_padding: Tensor = None) -> Tensor:
+        """
+        :param x_recon: [B, 1, M, T]
+        :param noise: [B, 1, M, T]
+        :param non_padding: [B, T, M]
+        """
+        x_recon, noise = self._mask_non_padding(x_recon, noise, non_padding)
+        return self._forward(x_recon, noise).mean()

modules/losses/dur_loss.py

@@ -0,0 +1,56 @@
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+
+class DurationLoss(nn.Module):
+    """
+    Loss module as combination of phone duration loss, word duration loss and sentence duration loss.
+    """
+
+    def __init__(self, offset, loss_type,
+                 lambda_pdur=0.6, lambda_wdur=0.3, lambda_sdur=0.1):
+        super().__init__()
+        self.loss_type = loss_type
+        if self.loss_type == 'mse':
+            self.loss = nn.MSELoss()
+        elif self.loss_type == 'huber':
+            self.loss = nn.HuberLoss()
+        else:
+            raise NotImplementedError()
+        self.offset = offset
+
+        self.lambda_pdur = lambda_pdur
+        self.lambda_wdur = lambda_wdur
+        self.lambda_sdur = lambda_sdur
+
+    def linear2log(self, any_dur):
+        return torch.log(any_dur + self.offset)
+
+    def forward(self, dur_pred: Tensor, dur_gt: Tensor, ph2word: Tensor) -> Tensor:
+        dur_gt = dur_gt.to(dtype=dur_pred.dtype)
+
+        # pdur_loss
+        pdur_loss = self.lambda_pdur * self.loss(self.linear2log(dur_pred), self.linear2log(dur_gt))
+
+        dur_pred = dur_pred.clamp(min=0.)  # clip to avoid NaN loss
+
+        # wdur loss
+        shape = dur_pred.shape[0], ph2word.max() + 1
+        wdur_pred = dur_pred.new_zeros(*shape).scatter_add(
+            1, ph2word, dur_pred
+        )[:, 1:]  # [B, T_ph] => [B, T_w]
+        wdur_gt = dur_gt.new_zeros(*shape).scatter_add(
+            1, ph2word, dur_gt
+        )[:, 1:]  # [B, T_ph] => [B, T_w]
+        wdur_loss = self.lambda_wdur * self.loss(self.linear2log(wdur_pred), self.linear2log(wdur_gt))
+
+        # sdur loss
+        sdur_pred = dur_pred.sum(dim=1)
+        sdur_gt = dur_gt.sum(dim=1)
+        sdur_loss = self.lambda_sdur * self.loss(self.linear2log(sdur_pred), self.linear2log(sdur_gt))
+
+        # combine
+        dur_loss = pdur_loss + wdur_loss + sdur_loss
+
+        return dur_loss

modules/losses/reflow_loss.py

@@ -0,0 +1,50 @@
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+
+class RectifiedFlowLoss(nn.Module):
+    def __init__(self, loss_type, log_norm=True):
+        super().__init__()
+        self.loss_type = loss_type
+        self.log_norm = log_norm
+        if self.loss_type == 'l1':
+            self.loss = nn.L1Loss(reduction='none')
+        elif self.loss_type == 'l2':
+            self.loss = nn.MSELoss(reduction='none')
+        else:
+            raise NotImplementedError()
+
+    @staticmethod
+    def _mask_non_padding(v_pred, v_gt, non_padding=None):
+        if non_padding is not None:
+            non_padding = non_padding.transpose(1, 2).unsqueeze(1)
+            return v_pred * non_padding, v_gt * non_padding
+        else:
+            return v_pred, v_gt
+
+    @staticmethod
+    def get_weights(t):
+        eps = 1e-7
+        t = t.float()
+        t = torch.clip(t, 0 + eps, 1 - eps)
+        weights = 0.398942 / t / (1 - t) * torch.exp(
+            -0.5 * torch.log(t / (1 - t)) ** 2
+        ) + eps
+        return weights[:, None, None, None]
+
+    def _forward(self, v_pred, v_gt, t=None):
+        if self.log_norm:
+            return self.get_weights(t) * self.loss(v_pred, v_gt)
+        else:
+            return self.loss(v_pred, v_gt)
+
+    def forward(self, v_pred: Tensor, v_gt: Tensor, t: Tensor, non_padding: Tensor = None) -> Tensor:
+        """
+        :param v_pred: [B, 1, M, T]
+        :param v_gt: [B, 1, M, T]
+        :param t: [B,]
+        :param non_padding: [B, T, M]
+        """
+        v_pred, v_gt = self._mask_non_padding(v_pred, v_gt, non_padding)
+        return self._forward(v_pred, v_gt, t=t).mean()

modules/metrics/__init__.py

@@ -0,0 +1,2 @@
+from .curve import RawCurveAccuracy, RawCurveR2Score
+from .duration import RhythmCorrectness, PhonemeDurationAccuracy

modules/metrics/curve.py

@@ -0,0 +1,73 @@
+import torch
+import torchmetrics
+from torch import Tensor
+
+
+class RawCurveAccuracy(torchmetrics.Metric):
+    def __init__(self, *, tolerance, **kwargs):
+        super().__init__(**kwargs)
+        self.tolerance = tolerance
+        self.add_state('close', default=torch.tensor(0, dtype=torch.int), dist_reduce_fx='sum')
+        self.add_state('total', default=torch.tensor(0, dtype=torch.int), dist_reduce_fx='sum')
+
+    def update(self, pred: Tensor, target: Tensor, mask=None) -> None:
+        """
+
+        :param pred: predicted curve
+        :param target: reference curve
+        :param mask: valid or non-padding mask
+        """
+        if mask is None:
+            assert pred.shape == target.shape, f'shapes of pred and target mismatch: {pred.shape}, {target.shape}'
+        else:
+            assert pred.shape == target.shape == mask.shape, \
+                f'shapes of pred, target and mask mismatch: {pred.shape}, {target.shape}, {mask.shape}'
+        close = torch.abs(pred - target) <= self.tolerance
+        if mask is not None:
+            close &= mask
+
+        self.close += close.sum()
+        self.total += pred.numel() if mask is None else mask.sum()
+
+    def compute(self) -> Tensor:
+        return self.close / self.total
+
+
+class RawCurveR2Score(torchmetrics.Metric):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.add_state('sum_squared_error', default=torch.tensor(0.0), dist_reduce_fx='sum')
+        self.add_state('sum_error', default=torch.tensor(0.0), dist_reduce_fx='sum')
+        self.add_state('residual', default=torch.tensor(0.0), dist_reduce_fx='sum')
+        self.add_state('total', default=torch.tensor(0), dist_reduce_fx='sum')
+
+    def update(self, pred: Tensor, target: Tensor, mask=None) -> None:
+        """
+
+        :param pred: predicted curve
+        :param target: reference curve
+        :param mask: valid or non-padding mask
+        """
+        if mask is None:
+            assert pred.shape == target.shape, f'shapes of pred and target mismatch: {pred.shape}, {target.shape}'
+        else:
+            assert pred.shape == target.shape == mask.shape, \
+                f'shapes of pred, target and mask mismatch: {pred.shape}, {target.shape}, {mask.shape}'
+            pred = pred[mask]
+            target = target[mask]
+        pred = pred.flatten()
+        target = target.flatten()
+
+        sum_error = torch.sum(target)
+        sum_squared_error = torch.sum(target * target)
+        residual = target - pred
+        rss = torch.sum(residual * residual)
+        total = target.numel() if mask is None else mask.sum()
+
+        self.sum_squared_error += sum_squared_error
+        self.sum_error += sum_error
+        self.residual += rss
+        self.total += total
+
+    def compute(self) -> Tensor:
+        return 1 - self.residual / (self.sum_squared_error - self.sum_error ** 2 / self.total)

modules/metrics/duration.py

@@ -0,0 +1,97 @@
+import torch
+import torchmetrics
+from torch import Tensor
+
+from modules.fastspeech.tts_modules import RhythmRegulator
+
+
+def linguistic_checks(pred, target, ph2word, mask=None):
+    if mask is None:
+        assert pred.shape == target.shape == ph2word.shape, \
+            f'shapes of pred, target and ph2word mismatch: {pred.shape}, {target.shape}, {ph2word.shape}'
+    else:
+        assert pred.shape == target.shape == ph2word.shape == mask.shape, \
+            f'shapes of pred, target and mask mismatch: {pred.shape}, {target.shape}, {ph2word.shape}, {mask.shape}'
+    assert pred.ndim == 2, f'all inputs should be 2D, but got {pred.shape}'
+    assert torch.any(ph2word > 0), 'empty word sequence'
+    assert torch.all(ph2word >= 0), 'unexpected negative word index'
+    assert ph2word.max() <= pred.shape[1], f'word index out of range: {ph2word.max()} > {pred.shape[1]}'
+    assert torch.all(pred >= 0.), f'unexpected negative ph_dur prediction'
+    assert torch.all(target >= 0.), f'unexpected negative ph_dur target'
+
+
+class RhythmCorrectness(torchmetrics.Metric):
+    def __init__(self, *, tolerance, **kwargs):
+        super().__init__(**kwargs)
+        assert 0. < tolerance < 1., 'tolerance should be within (0, 1)'
+        self.tolerance = tolerance
+        self.add_state('correct', default=torch.tensor(0, dtype=torch.int), dist_reduce_fx='sum')
+        self.add_state('total', default=torch.tensor(0, dtype=torch.int), dist_reduce_fx='sum')
+
+    def update(self, pdur_pred: Tensor, pdur_target: Tensor, ph2word: Tensor, mask=None) -> None:
+        """
+
+        :param pdur_pred: predicted ph_dur
+        :param pdur_target: reference ph_dur
+        :param ph2word: word division sequence
+        :param mask: valid or non-padding mask
+        """
+        linguistic_checks(pdur_pred, pdur_target, ph2word, mask=mask)
+
+        shape = pdur_pred.shape[0], ph2word.max() + 1
+        wdur_pred = pdur_pred.new_zeros(*shape).scatter_add(
+            1, ph2word, pdur_pred
+        )[:, 1:]  # [B, T_ph] => [B, T_w]
+        wdur_target = pdur_target.new_zeros(*shape).scatter_add(
+            1, ph2word, pdur_target
+        )[:, 1:]  # [B, T_ph] => [B, T_w]
+        if mask is None:
+            wdur_mask = torch.ones_like(wdur_pred, dtype=torch.bool)
+        else:
+            wdur_mask = mask.new_zeros(*shape).scatter_add(
+                1, ph2word, mask
+            )[:, 1:].bool()  # [B, T_ph] => [B, T_w]
+
+        correct = torch.abs(wdur_pred - wdur_target) <= wdur_target * self.tolerance
+        correct &= wdur_mask
+
+        self.correct += correct.sum()
+        self.total += wdur_mask.sum()
+
+    def compute(self) -> Tensor:
+        return self.correct / self.total
+
+
+class PhonemeDurationAccuracy(torchmetrics.Metric):
+    def __init__(self, *, tolerance, **kwargs):
+        super().__init__(**kwargs)
+        self.tolerance = tolerance
+        self.rr = RhythmRegulator()
+        self.add_state('accurate', default=torch.tensor(0, dtype=torch.int), dist_reduce_fx='sum')
+        self.add_state('total', default=torch.tensor(0, dtype=torch.int), dist_reduce_fx='sum')
+
+    def update(self, pdur_pred: Tensor, pdur_target: Tensor, ph2word: Tensor, mask=None) -> None:
+        """
+
+        :param pdur_pred: predicted ph_dur
+        :param pdur_target: reference ph_dur
+        :param ph2word: word division sequence
+        :param mask: valid or non-padding mask
+        """
+        linguistic_checks(pdur_pred, pdur_target, ph2word, mask=mask)
+
+        shape = pdur_pred.shape[0], ph2word.max() + 1
+        wdur_target = pdur_target.new_zeros(*shape).scatter_add(
+            1, ph2word, pdur_target
+        )[:, 1:]  # [B, T_ph] => [B, T_w]
+        pdur_align = self.rr(pdur_pred, ph2word=ph2word, word_dur=wdur_target)
+
+        accurate = torch.abs(pdur_align - pdur_target) <= pdur_target * self.tolerance
+        if mask is not None:
+            accurate &= mask
+
+        self.accurate += accurate.sum()
+        self.total += pdur_pred.numel() if mask is None else mask.sum()
+
+    def compute(self) -> Tensor:
+        return self.accurate / self.total

modules/nsf_hifigan/env.py

@@ -0,0 +1,32 @@
+class AttrDict(dict):
+    """A dictionary with attribute-style access. It maps attribute access to
+    the real dictionary.  """
+    def __init__(self, *args, **kwargs):
+        dict.__init__(self, *args, **kwargs)
+
+    def __getstate__(self):
+        return self.__dict__.items()
+
+    def __setstate__(self, items):
+        for key, val in items:
+            self.__dict__[key] = val
+
+    def __repr__(self):
+        return "%s(%s)" % (self.__class__.__name__, dict.__repr__(self))
+
+    def __setitem__(self, key, value):
+        return super(AttrDict, self).__setitem__(key, value)
+
+    def __getitem__(self, name):
+        if name not in super(AttrDict, self).keys():
+            return None
+        return super(AttrDict, self).__getitem__(name)
+
+    def __delitem__(self, name):
+        return super(AttrDict, self).__delitem__(name)
+
+    __getattr__ = __getitem__
+    __setattr__ = __setitem__
+
+    def copy(self):
+        return AttrDict(self)

modules/nsf_hifigan/models.py

@@ -0,0 +1,303 @@
+import json
+import pathlib
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from lightning.pytorch.utilities.rank_zero import rank_zero_info
+from torch.nn import Conv1d, ConvTranspose1d
+from torch.nn.utils import weight_norm, remove_weight_norm
+
+from .env import AttrDict
+from .utils import init_weights, get_padding
+
+LRELU_SLOPE = 0.1
+
+
+def load_model(model_path: pathlib.Path):
+    config_file = model_path.with_name('config.json')
+    with open(config_file) as f:
+        data = f.read()
+
+    json_config = json.loads(data)
+    h = AttrDict(json_config)
+
+    generator = Generator(h)
+
+    cp_dict = torch.load(model_path, map_location='cpu')
+    generator.load_state_dict(cp_dict['generator'])
+    generator.eval()
+    generator.remove_weight_norm()
+    del cp_dict
+    return generator, h
+
+
+class ResBlock1(torch.nn.Module):
+    def __init__(self, h, channels, kernel_size=3, dilation=(1, 3, 5)):
+        super(ResBlock1, self).__init__()
+        self.h = h
+        self.convs1 = nn.ModuleList([
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
+                               padding=get_padding(kernel_size, dilation[0]))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
+                               padding=get_padding(kernel_size, dilation[1]))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2],
+                               padding=get_padding(kernel_size, dilation[2])))
+        ])
+        self.convs1.apply(init_weights)
+
+        self.convs2 = nn.ModuleList([
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
+                               padding=get_padding(kernel_size, 1))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
+                               padding=get_padding(kernel_size, 1))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
+                               padding=get_padding(kernel_size, 1)))
+        ])
+        self.convs2.apply(init_weights)
+
+    def forward(self, x):
+        for c1, c2 in zip(self.convs1, self.convs2):
+            xt = F.leaky_relu(x, LRELU_SLOPE)
+            xt = c1(xt)
+            xt = F.leaky_relu(xt, LRELU_SLOPE)
+            xt = c2(xt)
+            x = xt + x
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs1:
+            remove_weight_norm(l)
+        for l in self.convs2:
+            remove_weight_norm(l)
+
+
+class ResBlock2(torch.nn.Module):
+    def __init__(self, h, channels, kernel_size=3, dilation=(1, 3)):
+        super(ResBlock2, self).__init__()
+        self.h = h
+        self.convs = nn.ModuleList([
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
+                               padding=get_padding(kernel_size, dilation[0]))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
+                               padding=get_padding(kernel_size, dilation[1])))
+        ])
+        self.convs.apply(init_weights)
+
+    def forward(self, x):
+        for c in self.convs:
+            xt = F.leaky_relu(x, LRELU_SLOPE)
+            xt = c(xt)
+            x = xt + x
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs:
+            remove_weight_norm(l)
+
+
+class SineGen(torch.nn.Module):
+    """ Definition of sine generator
+    SineGen(samp_rate, harmonic_num = 0,
+            sine_amp = 0.1, noise_std = 0.003,
+            voiced_threshold = 0,
+            flag_for_pulse=False)
+    samp_rate: sampling rate in Hz
+    harmonic_num: number of harmonic overtones (default 0)
+    sine_amp: amplitude of sine-waveform (default 0.1)
+    noise_std: std of Gaussian noise (default 0.003)
+    voiced_threshold: F0 threshold for U/V classification (default 0)
+    flag_for_pulse: this SinGen is used inside PulseGen (default False)
+    Note: when flag_for_pulse is True, the first time step of a voiced
+        segment is always sin(np.pi) or cos(0)
+    """
+
+    def __init__(self, samp_rate, harmonic_num=0,
+                 sine_amp=0.1, noise_std=0.003,
+                 voiced_threshold=0):
+        super(SineGen, self).__init__()
+        self.sine_amp = sine_amp
+        self.noise_std = noise_std
+        self.harmonic_num = harmonic_num
+        self.dim = self.harmonic_num + 1
+        self.sampling_rate = samp_rate
+        self.voiced_threshold = voiced_threshold
+
+    def _f02uv(self, f0):
+        # generate uv signal
+        uv = torch.ones_like(f0)
+        uv = uv * (f0 > self.voiced_threshold)
+        return uv
+
+    def _f02sine(self, f0, upp):
+        """ f0: (batchsize, length, dim)
+            where dim indicates fundamental tone and overtones
+        """
+        rad = f0 / self.sampling_rate * torch.arange(1, upp + 1, device=f0.device)
+        rad2 = torch.fmod(rad[..., -1:].float() + 0.5, 1.0) - 0.5
+        rad_acc = rad2.cumsum(dim=1).fmod(1.0).to(f0)
+        rad += F.pad(rad_acc[:, :-1, :], (0, 0, 1, 0))
+        rad = rad.reshape(f0.shape[0], -1, 1)
+        rad = torch.multiply(rad, torch.arange(1, self.dim + 1, device=f0.device).reshape(1, 1, -1))
+        rand_ini = torch.rand(1, 1, self.dim, device=f0.device)
+        rand_ini[..., 0] = 0
+        rad += rand_ini
+        sines = torch.sin(2 * np.pi * rad)
+        return sines
+
+    @torch.no_grad()
+    def forward(self, f0, upp):
+        """ sine_tensor, uv = forward(f0)
+        input F0: tensor(batchsize=1, length, dim=1)
+                  f0 for unvoiced steps should be 0
+        output sine_tensor: tensor(batchsize=1, length, dim)
+        output uv: tensor(batchsize=1, length, 1)
+        """
+        f0 = f0.unsqueeze(-1)
+        sine_waves = self._f02sine(f0, upp) * self.sine_amp
+        uv = (f0 > self.voiced_threshold).float()
+        uv = F.interpolate(uv.transpose(2, 1), scale_factor=upp, mode='nearest').transpose(2, 1)
+        noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3
+        noise = noise_amp * torch.randn_like(sine_waves)
+        sine_waves = sine_waves * uv + noise
+        return sine_waves
+
+
+class SourceModuleHnNSF(torch.nn.Module):
+    """ SourceModule for hn-nsf
+    SourceModule(sampling_rate, harmonic_num=0, sine_amp=0.1,
+                 add_noise_std=0.003, voiced_threshod=0)
+    sampling_rate: sampling_rate in Hz
+    harmonic_num: number of harmonic above F0 (default: 0)
+    sine_amp: amplitude of sine source signal (default: 0.1)
+    add_noise_std: std of additive Gaussian noise (default: 0.003)
+        note that amplitude of noise in unvoiced is decided
+        by sine_amp
+    voiced_threshold: threhold to set U/V given F0 (default: 0)
+    Sine_source, noise_source = SourceModuleHnNSF(F0_sampled)
+    F0_sampled (batchsize, length, 1)
+    Sine_source (batchsize, length, 1)
+    noise_source (batchsize, length 1)
+    uv (batchsize, length, 1)
+    """
+
+    def __init__(self, sampling_rate, harmonic_num=0, sine_amp=0.1,
+                 add_noise_std=0.003, voiced_threshold=0):
+        super(SourceModuleHnNSF, self).__init__()
+
+        self.sine_amp = sine_amp
+        self.noise_std = add_noise_std
+
+        # to produce sine waveforms
+        self.l_sin_gen = SineGen(sampling_rate, harmonic_num,
+                                 sine_amp, add_noise_std, voiced_threshold)
+
+        # to merge source harmonics into a single excitation
+        self.l_linear = torch.nn.Linear(harmonic_num + 1, 1)
+        self.l_tanh = torch.nn.Tanh()
+
+    def forward(self, x, upp):
+        sine_wavs = self.l_sin_gen(x, upp)
+        sine_merge = self.l_tanh(self.l_linear(sine_wavs))
+        return sine_merge
+
+
+class Generator(torch.nn.Module):
+    def __init__(self, h):
+        super(Generator, self).__init__()
+        self.h = h
+        self.num_kernels = len(h.resblock_kernel_sizes)
+        self.num_upsamples = len(h.upsample_rates)
+        self.mini_nsf = h.mini_nsf
+        self.noise_sigma = h.noise_sigma
+        
+        if h.mini_nsf:
+            self.source_sr = h.sampling_rate / int(np.prod(h.upsample_rates[2: ]))
+            self.upp = int(np.prod(h.upsample_rates[: 2]))
+        else:
+            self.source_sr = h.sampling_rate
+            self.upp = int(np.prod(h.upsample_rates))
+            self.m_source = SourceModuleHnNSF(
+                sampling_rate=h.sampling_rate,
+                harmonic_num=8
+            )
+            self.noise_convs = nn.ModuleList()
+        
+        self.conv_pre = weight_norm(Conv1d(h.num_mels, h.upsample_initial_channel, 7, 1, padding=3))   
+        
+        self.ups = nn.ModuleList()
+        self.resblocks = nn.ModuleList()
+        resblock = ResBlock1 if h.resblock == '1' else ResBlock2
+        ch = h.upsample_initial_channel
+        for i, (u, k) in enumerate(zip(h.upsample_rates, h.upsample_kernel_sizes)):
+            ch //= 2
+            self.ups.append(weight_norm(ConvTranspose1d(2 * ch, ch, k, u, padding=(k - u) // 2)))
+            for j, (k, d) in enumerate(zip(h.resblock_kernel_sizes, h.resblock_dilation_sizes)):
+                self.resblocks.append(resblock(h, ch, k, d))
+            if not h.mini_nsf:
+                if i + 1 < len(h.upsample_rates):  #
+                    stride_f0 = int(np.prod(h.upsample_rates[i + 1:]))
+                    self.noise_convs.append(Conv1d(
+                        1, ch, kernel_size=stride_f0 * 2, stride=stride_f0, padding=stride_f0 // 2))
+                else:
+                    self.noise_convs.append(Conv1d(1, ch, kernel_size=1))
+            elif i == 1:
+                self.source_conv = Conv1d(1, ch, 1)
+                self.source_conv.apply(init_weights)
+
+        self.conv_post = weight_norm(Conv1d(ch, 1, 7, 1, padding=3))
+        
+        self.ups.apply(init_weights)
+        self.conv_post.apply(init_weights)
+        
+    def fastsinegen(self, f0):
+        n = torch.arange(1, self.upp + 1, device=f0.device)
+        s0 = f0.unsqueeze(-1) / self.source_sr
+        ds0 = F.pad(s0[:, 1:, :] - s0[:, :-1, :], (0, 0, 0, 1))
+        rad = s0 * n + 0.5 * ds0 * n * (n - 1) / self.upp
+        rad2 = torch.fmod(rad[..., -1:].float() + 0.5, 1.0) - 0.5
+        rad_acc = rad2.cumsum(dim=1).fmod(1.0).to(f0)
+        rad += F.pad(rad_acc[:, :-1, :], (0, 0, 1, 0))
+        rad = rad.reshape(f0.shape[0], 1, -1)
+        sines = torch.sin(2 * np.pi * rad)
+        return sines
+        
+    def forward(self, x, f0):
+        if self.mini_nsf:
+            har_source = self.fastsinegen(f0)
+        else:
+            har_source = self.m_source(f0, self.upp).transpose(1, 2)
+        x = self.conv_pre(x)
+        if self.noise_sigma is not None and self.noise_sigma > 0:
+            x += self.noise_sigma * torch.randn_like(x)
+        for i in range(self.num_upsamples):
+            x = F.leaky_relu(x, LRELU_SLOPE)
+            x = self.ups[i](x)
+            if not self.mini_nsf:
+                x_source = self.noise_convs[i](har_source)
+                x = x + x_source
+            elif i == 1:
+                x_source = self.source_conv(har_source)
+                x = x + x_source
+            xs = None
+            for j in range(self.num_kernels):
+                if xs is None:
+                    xs = self.resblocks[i * self.num_kernels + j](x)
+                else:
+                    xs += self.resblocks[i * self.num_kernels + j](x)
+            x = xs / self.num_kernels
+        x = F.leaky_relu(x)
+        x = self.conv_post(x)
+        x = torch.tanh(x)
+        return x
+
+    def remove_weight_norm(self):
+        # rank_zero_info('Removing weight norm...')
+        print('Removing weight norm...')
+        for l in self.ups:
+            remove_weight_norm(l)
+        for l in self.resblocks:
+            l.remove_weight_norm()
+        remove_weight_norm(self.conv_pre)
+        remove_weight_norm(self.conv_post)

modules/nsf_hifigan/nvSTFT.py

@@ -0,0 +1,87 @@
+import os
+
+os.environ["LRU_CACHE_CAPACITY"] = "3"
+import torch
+import torch.utils.data
+import numpy as np
+from librosa.filters import mel as librosa_mel_fn
+import torch.nn.functional as F
+
+
+def dynamic_range_compression(x, C=1, clip_val=1e-5):
+    return np.log(np.clip(x, a_min=clip_val, a_max=None) * C)
+
+
+def dynamic_range_decompression(x, C=1):
+    return np.exp(x) / C
+
+
+def dynamic_range_compression_torch(x, C=1, clip_val=1e-5):
+    return torch.log(torch.clamp(x, min=clip_val) * C)
+
+
+def dynamic_range_decompression_torch(x, C=1):
+    return torch.exp(x) / C
+
+
+class STFT:
+    def __init__(
+            self, sr=22050,
+            n_mels=80, n_fft=1024, win_size=1024, hop_length=256,
+            fmin=20, fmax=11025, clip_val=1e-5,
+            device=None
+    ):
+        self.target_sr = sr
+
+        self.n_mels = n_mels
+        self.n_fft = n_fft
+        self.win_size = win_size
+        self.hop_length = hop_length
+        self.fmin = fmin
+        self.fmax = fmax
+        self.clip_val = clip_val
+
+        if device is None:
+            device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        self.device = device
+
+        mel_basis = librosa_mel_fn(sr=sr, n_fft=n_fft, n_mels=n_mels, fmin=fmin, fmax=fmax)
+        self.mel_basis = torch.from_numpy(mel_basis).float().to(device)
+
+    def get_mel(self, y, keyshift=0, speed=1, center=False):
+
+        factor = 2 ** (keyshift / 12)
+        n_fft_new = int(np.round(self.n_fft * factor))
+        win_size_new = int(np.round(self.win_size * factor))
+        hop_length_new = int(np.round(self.hop_length * speed))
+
+        if torch.min(y) < -1.:
+            print('min value is ', torch.min(y))
+        if torch.max(y) > 1.:
+            print('max value is ', torch.max(y))
+
+        window = torch.hann_window(win_size_new, device=self.device)
+
+        y = torch.nn.functional.pad(y.unsqueeze(1), (
+            (win_size_new - hop_length_new) // 2,
+            (win_size_new - hop_length_new + 1) // 2
+        ), mode='reflect')
+        y = y.squeeze(1)
+
+        spec = torch.stft(
+            y, n_fft_new, hop_length=hop_length_new,
+            win_length=win_size_new, window=window,
+            center=center, pad_mode='reflect',
+            normalized=False, onesided=True, return_complex=True
+        ).abs()
+        if keyshift != 0:
+            size = self.n_fft // 2 + 1
+            resize = spec.size(1)
+            if resize < size:
+                spec = F.pad(spec, (0, 0, 0, size - resize))
+            spec = spec[:, :size, :] * self.win_size / win_size_new
+
+        spec = torch.matmul(self.mel_basis, spec)
+        spec = dynamic_range_compression_torch(spec, clip_val=self.clip_val)
+
+        return spec

modules/nsf_hifigan/utils.py

@@ -0,0 +1,13 @@
+import matplotlib
+
+matplotlib.use("Agg")
+
+
+def init_weights(m, mean=0.0, std=0.01):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        m.weight.data.normal_(mean, std)
+
+
+def get_padding(kernel_size, dilation=1):
+    return int((kernel_size*dilation - dilation)/2)

modules/optimizer/chained_optimizer.py

@@ -0,0 +1,132 @@
+import torch
+from torch import Tensor
+from torch.optim import Optimizer
+from torch.optim.optimizer import ParamsT
+from dataclasses import dataclass
+from typing import Any, Dict, List, Type, Callable, Optional
+
+
+@dataclass
+class OptimizerSpec:
+    """Spec for creating an optimizer that is part of a `ChainedOptimizer`."""
+
+    class_type: Type[Optimizer]
+    init_args: Dict[str, Any]
+    param_filter: Optional[Callable[[Tensor], bool]]
+
+
+class ChainedOptimizer(Optimizer):
+    """
+    A wrapper around multiple optimizers that allows for chaining them together.
+    The optimizers are applied in the order they are passed in the constructor.
+    Each optimizer is responsible for updating a subset of the parameters, which
+    is determined by the `param_filter` function. If no optimizer is found for a
+    parameter group, an exception is raised.
+    """
+
+    def __init__(
+        self,
+        params: ParamsT,
+        optimizer_specs: List[OptimizerSpec],
+        lr: float,
+        weight_decay: float = 0.0,
+        optimizer_selection_callback: Optional[Callable[[Tensor, int], None]] = None,
+        **common_kwargs,
+    ):
+        self.optimizer_specs = optimizer_specs
+        self.optimizer_selection_callback = optimizer_selection_callback
+        self.optimizers: List[Optimizer] = []
+        defaults = dict(lr=lr, weight_decay=weight_decay)
+        super().__init__(params, defaults)
+
+        # Split the params for each optimizer
+        params_for_optimizers = [[] for _ in optimizer_specs]
+        for param_group in self.param_groups:
+            params = param_group["params"]
+            indices = param_group["optimizer_and_param_group_indices"] = set()
+            for param in params:
+                assert isinstance(param, Tensor), f"Expected a Tensor, got {type(param)}"
+                found_optimizer = False
+                for index, spec in enumerate(optimizer_specs):
+                    if spec.param_filter is None or spec.param_filter(param):
+                        if self.optimizer_selection_callback is not None:
+                            self.optimizer_selection_callback(param, index)
+                        params_for_optimizers[index].append(param)
+                        indices.add((index, 0))
+                        found_optimizer = True
+                        break
+                if not found_optimizer:
+                    raise ValueError("No valid optimizer found for the given parameter")
+
+        # Initialize the optimizers
+        for spec, selected_params in zip(optimizer_specs, params_for_optimizers):
+            optimizer_args = {
+                'lr': lr,
+                'weight_decay': weight_decay,
+            }
+            optimizer_args.update(common_kwargs)
+            optimizer_args.update(spec.init_args)
+            optimizer = spec.class_type(selected_params, **optimizer_args)
+            self.optimizers.append(optimizer)
+
+    def state_dict(self) -> Dict[str, Any]:
+        return {
+            "optimizers": [opt.state_dict() for opt in self.optimizers],
+            **super().state_dict(),
+        }
+
+    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
+        optimizers = state_dict.pop("optimizers")
+        super().load_state_dict(state_dict)
+        for i in range(len(self.optimizers)):
+            self.optimizers[i].load_state_dict(optimizers[i])
+
+    def zero_grad(self, set_to_none: bool = True) -> None:
+        for opt in self.optimizers:
+            opt.zero_grad(set_to_none=set_to_none)
+
+    def _copy_lr_to_optimizers(self) -> None:
+        for param_group in self.param_groups:
+            indices = param_group["optimizer_and_param_group_indices"]
+            for optimizer_idx, param_group_idx in indices:
+                self.optimizers[optimizer_idx].param_groups[param_group_idx]["lr"] = param_group["lr"]
+
+    def step(self, closure=None) -> None:
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+        self._copy_lr_to_optimizers()
+        for opt in self.optimizers:
+            opt.step(closure=None)
+        return loss
+
+    def add_param_group(self, param_group: Dict[str, Any]) -> None:
+        super().add_param_group(param_group)
+
+        # If optimizer has not been initialized, skip adding the param groups
+        if not self.optimizers:
+            return
+
+        # Split the params for each optimizer
+        params_for_optimizers = [[] for _ in self.optimizer_specs]
+        params = param_group["params"]
+        indices = param_group["optimizer_and_param_group_indices"] = set()
+        for param in params:
+            assert isinstance(param, Tensor), f"Expected a Tensor, got {type(param)}"
+            found_optimizer = False
+            for index, spec in enumerate(self.optimizer_specs):
+                if spec.param_filter is None or spec.param_filter(param):
+                    if self.optimizer_selection_callback is not None:
+                        self.optimizer_selection_callback(param, index)
+                    params_for_optimizers[index].append(param)
+                    indices.add((index, len(self.optimizers[index].param_groups)))
+                    found_optimizer = True
+                    break
+            if not found_optimizer:
+                raise ValueError("No valid optimizer found for the given parameter group")
+
+        # Add the selected param group to the optimizers
+        for optimizer, selected_params in zip(self.optimizers, params_for_optimizers):
+            if selected_params:
+                optimizer.add_param_group({"params": selected_params})

modules/optimizer/muon.py

@@ -0,0 +1,200 @@
+import collections
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+from torch.nn import Parameter
+from typing import List
+from .chained_optimizer import ChainedOptimizer, OptimizerSpec
+
+from modules.commons.common_layers import AdamWLinear, AdamWConv1d
+
+
+def zeropower_via_newtonschulz5(G: Tensor, steps: int) -> Tensor:
+    """
+    Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a
+    quintic iteration whose coefficients are selected to maximize the slope at zero. For the purpose
+    of minimizing steps, it turns out to be empirically effective to keep increasing the slope at
+    zero even beyond the point where the iteration no longer converges all the way to one everywhere
+    on the interval. This iteration therefore does not produce UV^T but rather something like US'V^T
+    where S' is diagonal with S_{ii}' ~ Uniform(0.5, 1.5), which turns out not to hurt model
+    performance at all relative to UV^T, where USV^T = G is the SVD.
+    """
+    assert G.ndim == 3 # batched Muon implementation by @scottjmaddox, and put into practice in the record by @YouJiacheng
+    a, b, c = (3.4445, -4.7750,  2.0315)
+    
+    X = G.to(torch.float32)
+
+    # Ensure spectral norm is at most 1
+    X = F.normalize(X, p=2.0, dim=(-2, -1), eps=1e-7)
+    
+    X = X.to(torch.float16)
+    
+    # Perform the NS iterations
+    if X.size(-2) < X.size(-1):
+        for _ in range(steps):
+            A = torch.bmm(X, X.mT)
+            A = torch.baddbmm(A, A, A, beta=b, alpha=c)
+            X = torch.baddbmm(X, A, X, beta=a, alpha=1)
+    else:
+        for _ in range(steps):
+            A = torch.bmm(X.mT, X)
+            A = torch.baddbmm(A, A, A, beta=b, alpha=c)
+            X = torch.baddbmm(X, X, A, beta=a, alpha=1)
+            
+    return X
+
+
+def gram_newton_schulz(G: Tensor, steps: int) -> Tensor:
+    """
+    Refer to: 
+    Gram Newton-Schulz: A Fast, Hardware-Aware Newton-Schulz Algorithm for Muon
+    Authors: Jack Zhang, Noah Amsel, Berlin Chen, Tri Dao
+    Blogpost: https://dao-ailab.github.io/blog/2026/gram-newton-schulz/
+    
+    Gram Newton-Schulz iteration to compute the orthogonalization of G.
+    Mathematically identical to standard Newton-Schulz but computes iterating 
+    on the smaller NxN Gram matrix to save up to 50% FLOPs.
+    """
+    assert G.ndim == 3
+    reset_iterations = [2]
+    original_shape = G.shape
+    dtype = G.dtype
+
+    X = G.to(torch.float32)
+    X = F.normalize(X, p=2.0, dim=(-2, -1), eps=1e-7)
+    should_transpose = X.size(-2) > X.size(-1)
+    if should_transpose:
+        X = X.mT
+    X = X.to(torch.float16)
+
+    a, b, c = (3.4445, -4.7750,  2.0315)
+
+    if X.size(-2) != X.size(-1):
+        R = torch.bmm(X, X.mT)
+        Q = None
+        for i in range(steps):
+            if i in reset_iterations and i != 0:
+                X = torch.bmm(Q, X)
+                R = torch.bmm(X, X.mT)
+                Q = None
+            Z = torch.baddbmm(R, R, R, beta=b, alpha=c)
+            if i != 0 and i not in reset_iterations:
+                Q = torch.baddbmm(Q, Q, Z, beta=a, alpha=1.0)
+            else:
+                Q = Z.clone()
+                Q.diagonal(dim1=-2, dim2=-1).add_(a)
+            if i < steps - 1 and (i + 1) not in reset_iterations:
+                RZ = torch.baddbmm(R, R, Z, beta=a, alpha=1.0)
+                R = torch.baddbmm(RZ, Z, RZ, beta=a, alpha=1.0)
+        X = torch.bmm(Q, X) if not should_transpose else torch.bmm(X.mT, Q)
+    else:
+        for _ in range(steps):
+            A = torch.bmm(X, X.mT)
+            B = torch.baddbmm(A, A, A, beta=b, alpha=c)
+            X = torch.baddbmm(X, B, X, beta=a, alpha=1.0)
+
+    return X.to(dtype).view(original_shape)
+
+
+class Muon(torch.optim.Optimizer):
+    """
+    Muon - MomentUm Orthogonalized by Newton-schulz
+
+    https://kellerjordan.github.io/posts/muon/
+
+    Muon internally runs standard SGD-momentum, and then performs an orthogonalization post-
+    processing step, in which each 2D parameter's update is replaced with the nearest orthogonal
+    matrix. To efficiently orthogonalize each update, we use a Newton-Schulz iteration, which has
+    the advantage that it can be stably run in float16 on the GPU.
+
+    Some warnings:
+    - This optimizer should not be used for the embedding layer, the final fully connected layer,
+    or any {0,1}-D parameters; those should all be optimized by a standard method (e.g., AdamW).
+    - To use it with 4D convolutional filters, it works well to just flatten their last 3 dimensions.
+
+    Arguments:
+        lr: The learning rate used by the internal SGD.
+        momentum: The momentum used by the internal SGD.
+        nesterov: Whether to use Nesterov-style momentum in the internal SGD. (recommended)
+        ns_steps: The number of Newton-Schulz iteration steps to use.
+    """
+
+    def __init__(self, params, lr=5e-4, weight_decay=0.1, momentum=0.95, nesterov=True, ns_steps=5):
+        defaults = dict(lr=lr, weight_decay=weight_decay, momentum=momentum, nesterov=nesterov, ns_steps=ns_steps)
+        super().__init__(params, defaults)
+    
+    @torch.no_grad()
+    def step(self, closure=None):
+        for group in self.param_groups:
+            shape_groups = {}
+            for p in filter(lambda p: p.grad is not None, group["params"]):
+                g = p.grad
+                state = self.state[p]
+                if "momentum_buffer" not in state:
+                    state["momentum_buffer"] = torch.zeros_like(g)
+                key = (p.shape, p.device, p.dtype)
+                if key not in shape_groups:
+                    shape_groups[key] = {"params": [], "grads": [], "buffers": []}
+                shape_groups[key]["params"].append(p)
+                shape_groups[key]["grads"].append(g)
+                shape_groups[key]["buffers"].append(state["momentum_buffer"])
+            for key in shape_groups:
+                group_data = shape_groups[key]
+                p, g, buf, m = group_data["params"], group_data["grads"], group_data["buffers"], group["momentum"]
+                torch._foreach_lerp_(buf, g, 1-m)
+                if group["nesterov"]:
+                    torch._foreach_lerp_(g, buf, m)
+                    g = torch.stack(g)
+                else:
+                    g = torch.stack(buf)
+                original_shape = g.shape
+                if g.ndim >= 4:  # for the case of conv filters
+                    g = g.view(g.size(0), g.size(1), -1)
+                g = gram_newton_schulz(g, steps=group["ns_steps"])
+                
+                if group["weight_decay"] > 0:
+                    torch._foreach_mul_(p, 1 - group["lr"] * group["weight_decay"])
+                torch._foreach_add_(p, g.view(original_shape).unbind(0), alpha=-group["lr"] * max(g[0].size()) ** 0.5)
+
+
+def get_params_for_muon(model) -> List[Parameter]:
+    """
+    Filter parameters of a module into two groups: those that can be optimized by Muon,
+    and those that should be optimized by a standard optimizer.
+    Args:
+        module: The module to filter parameters for.
+    Returns:
+        A list of parameters that should be optimized with muon.
+    """
+    excluded_module_classes = (nn.Embedding, AdamWLinear, AdamWConv1d)
+    muon_params = []
+    # BFS through all submodules and exclude parameters from certain module types
+    queue = collections.deque([model])
+    while queue:
+        module = queue.popleft()
+        if isinstance(module, excluded_module_classes):
+            continue
+        for param in module.parameters(recurse=False):
+            if not param.requires_grad:
+                continue
+            if param.ndim >= 2:
+                muon_params.append(param)
+        queue.extend(list(module.children()))
+    return muon_params
+
+
+class Muon_AdamW(ChainedOptimizer):
+    def __init__(self, model, lr=0.0005, weight_decay=0.0, muon_args=None, adamw_args=None, verbose=False):
+        muon_args = {} if muon_args is None else muon_args
+        adamw_args = {} if adamw_args is None else adamw_args
+        muon_params_id_set = set(id(p) for p in get_params_for_muon(model))
+        spec_muon = OptimizerSpec(Muon, muon_args, lambda param: id(param) in muon_params_id_set)
+        spec_adamw = OptimizerSpec(torch.optim.AdamW, adamw_args, None)
+        specs = [spec_muon, spec_adamw]
+        callback = None
+        if verbose:
+            callback = lambda p, spec_idx: print(
+            f"Adding param {p.shape} to optimizer{spec_idx} {str(specs[spec_idx].class_type)}"
+        )
+        super().__init__(model.parameters(), specs, lr=lr, weight_decay=weight_decay, optimizer_selection_callback=callback)

modules/pe/__init__.py

@@ -0,0 +1,18 @@
+from utils import hparams
+
+from .pm import ParselmouthPE
+from .pw import HarvestPE
+from .rmvpe import RMVPE
+
+
+def initialize_pe():
+    pe = hparams['pe']
+    pe_ckpt = hparams['pe_ckpt']
+    if pe == 'parselmouth':
+        return ParselmouthPE()
+    elif pe == 'rmvpe':
+        return RMVPE(pe_ckpt)
+    elif pe == 'harvest':
+        return HarvestPE()
+    else:
+        raise ValueError(f" [x] Unknown f0 extractor: {pe}")

modules/pe/pm.py

@@ -0,0 +1,15 @@
+from basics.base_pe import BasePE
+from utils.binarizer_utils import get_pitch_parselmouth
+
+
+class ParselmouthPE(BasePE):
+    def get_pitch(
+            self,waveform, samplerate, length,
+            *, hop_size, f0_min=65, f0_max=1100,
+            speed=1, interp_uv=False
+    ):
+        return get_pitch_parselmouth(
+            waveform, samplerate=samplerate, length=length,
+            hop_size=hop_size, f0_min=f0_min, f0_max=f0_max,
+            speed=speed, interp_uv=interp_uv
+        )

modules/pe/pw.py

@@ -0,0 +1,29 @@
+from basics.base_pe import BasePE
+import numpy as np
+import pyworld as pw
+from utils.pitch_utils import interp_f0
+
+
+class HarvestPE(BasePE):
+    def get_pitch(
+            self, waveform, samplerate, length,
+            *, hop_size, f0_min=65, f0_max=1100,
+            speed=1, interp_uv=False
+    ):
+        hop_size = int(np.round(hop_size * speed))
+        time_step = 1000 * hop_size / samplerate
+
+        f0, _ = pw.harvest(
+            waveform.astype(np.float64), samplerate,
+            f0_floor=f0_min, f0_ceil=f0_max, frame_period=time_step
+        )
+        f0 = f0.astype(np.float32)
+
+        if f0.size < length:
+            f0 = np.pad(f0, (0, length - f0.size))
+        f0 = f0[:length]
+        uv = f0 == 0
+
+        if interp_uv:
+            f0, uv = interp_f0(f0, uv)
+        return f0, uv

modules/pe/rmvpe/__init__.py

@@ -0,0 +1,5 @@
+from .constants import *
+from .model import E2E0
+from .utils import to_local_average_f0, to_viterbi_f0
+from .inference import RMVPE
+from .spec import MelSpectrogram

modules/pe/rmvpe/constants.py

@@ -0,0 +1,9 @@
+SAMPLE_RATE = 16000
+
+N_CLASS = 360
+
+N_MELS = 128
+MEL_FMIN = 30
+MEL_FMAX = 8000
+WINDOW_LENGTH = 1024
+CONST = 1997.3794084376191

modules/pe/rmvpe/deepunet.py

@@ -0,0 +1,173 @@
+import torch
+import torch.nn as nn
+from .constants import N_MELS
+
+
+class ConvBlockRes(nn.Module):
+    def __init__(self, in_channels, out_channels, momentum=0.01):
+        super(ConvBlockRes, self).__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_channels=in_channels,
+                      out_channels=out_channels,
+                      kernel_size=(3, 3),
+                      stride=(1, 1),
+                      padding=(1, 1),
+                      bias=False),
+            nn.BatchNorm2d(out_channels, momentum=momentum),
+            nn.ReLU(),
+
+            nn.Conv2d(in_channels=out_channels,
+                      out_channels=out_channels,
+                      kernel_size=(3, 3),
+                      stride=(1, 1),
+                      padding=(1, 1),
+                      bias=False),
+            nn.BatchNorm2d(out_channels, momentum=momentum),
+            nn.ReLU(),
+        )
+        if in_channels != out_channels:
+            self.shortcut = nn.Conv2d(in_channels, out_channels, (1, 1))
+            self.is_shortcut = True
+        else:
+            self.is_shortcut = False
+
+    def forward(self, x):
+        if self.is_shortcut:
+            return self.conv(x) + self.shortcut(x)
+        else:
+            return self.conv(x) + x
+
+
+class ResEncoderBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, n_blocks=1, momentum=0.01):
+        super(ResEncoderBlock, self).__init__()
+        self.n_blocks = n_blocks
+        self.conv = nn.ModuleList()
+        self.conv.append(ConvBlockRes(in_channels, out_channels, momentum))
+        for i in range(n_blocks - 1):
+            self.conv.append(ConvBlockRes(out_channels, out_channels, momentum))
+        self.kernel_size = kernel_size
+        if self.kernel_size is not None:
+            self.pool = nn.AvgPool2d(kernel_size=kernel_size)
+
+    def forward(self, x):
+        for i in range(self.n_blocks):
+            x = self.conv[i](x)
+        if self.kernel_size is not None:
+            return x, self.pool(x)
+        else:
+            return x
+
+
+class ResDecoderBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, stride, n_blocks=1, momentum=0.01):
+        super(ResDecoderBlock, self).__init__()
+        out_padding = (0, 1) if stride == (1, 2) else (1, 1)
+        self.n_blocks = n_blocks
+        self.conv1 = nn.Sequential(
+            nn.ConvTranspose2d(in_channels=in_channels,
+                               out_channels=out_channels,
+                               kernel_size=(3, 3),
+                               stride=stride,
+                               padding=(1, 1),
+                               output_padding=out_padding,
+                               bias=False),
+            nn.BatchNorm2d(out_channels, momentum=momentum),
+            nn.ReLU(),
+        )
+        self.conv2 = nn.ModuleList()
+        self.conv2.append(ConvBlockRes(out_channels * 2, out_channels, momentum))
+        for i in range(n_blocks-1):
+            self.conv2.append(ConvBlockRes(out_channels, out_channels, momentum))
+
+    def forward(self, x, concat_tensor):
+        x = self.conv1(x)
+        x = torch.cat((x, concat_tensor), dim=1)
+        for i in range(self.n_blocks):
+            x = self.conv2[i](x)
+        return x
+
+
+class Encoder(nn.Module):
+    def __init__(self, in_channels, in_size, n_encoders, kernel_size, n_blocks, out_channels=16, momentum=0.01):
+        super(Encoder, self).__init__()
+        self.n_encoders = n_encoders
+        self.bn = nn.BatchNorm2d(in_channels, momentum=momentum)
+        self.layers = nn.ModuleList()
+        self.latent_channels = []
+        for i in range(self.n_encoders):
+            self.layers.append(ResEncoderBlock(in_channels, out_channels, kernel_size, n_blocks, momentum=momentum))
+            self.latent_channels.append([out_channels, in_size])
+            in_channels = out_channels
+            out_channels *= 2
+            in_size //= 2
+        self.out_size = in_size
+        self.out_channel = out_channels
+
+    def forward(self, x):
+        concat_tensors = []
+        x = self.bn(x)
+        for i in range(self.n_encoders):
+            _, x = self.layers[i](x)
+            concat_tensors.append(_)
+        return x, concat_tensors
+
+
+class Intermediate(nn.Module):
+    def __init__(self, in_channels, out_channels, n_inters, n_blocks, momentum=0.01):
+        super(Intermediate, self).__init__()
+        self.n_inters = n_inters
+        self.layers = nn.ModuleList()
+        self.layers.append(ResEncoderBlock(in_channels, out_channels, None, n_blocks, momentum))
+        for i in range(self.n_inters-1):
+            self.layers.append(ResEncoderBlock(out_channels, out_channels, None, n_blocks, momentum))
+
+    def forward(self, x):
+        for i in range(self.n_inters):
+            x = self.layers[i](x)
+        return x
+
+
+class Decoder(nn.Module):
+    def __init__(self, in_channels, n_decoders, stride, n_blocks, momentum=0.01):
+        super(Decoder, self).__init__()
+        self.layers = nn.ModuleList()
+        self.n_decoders = n_decoders
+        for i in range(self.n_decoders):
+            out_channels = in_channels // 2
+            self.layers.append(ResDecoderBlock(in_channels, out_channels, stride, n_blocks, momentum))
+            in_channels = out_channels
+
+    def forward(self, x, concat_tensors):
+        for i in range(self.n_decoders):
+            x = self.layers[i](x, concat_tensors[-1-i])
+        return x
+
+
+class TimbreFilter(nn.Module):
+    def __init__(self, latent_rep_channels):
+        super(TimbreFilter, self).__init__()
+        self.layers = nn.ModuleList()
+        for latent_rep in latent_rep_channels:
+            self.layers.append(ConvBlockRes(latent_rep[0], latent_rep[0]))
+
+    def forward(self, x_tensors):
+        out_tensors = []
+        for i, layer in enumerate(self.layers):
+            out_tensors.append(layer(x_tensors[i]))
+        return out_tensors
+
+
+class DeepUnet0(nn.Module):
+    def __init__(self, kernel_size, n_blocks, en_de_layers=5, inter_layers=4, in_channels=1, en_out_channels=16):
+        super(DeepUnet0, self).__init__()
+        self.encoder = Encoder(in_channels, N_MELS, en_de_layers, kernel_size, n_blocks, en_out_channels)
+        self.intermediate = Intermediate(self.encoder.out_channel // 2, self.encoder.out_channel, inter_layers, n_blocks)
+        self.tf = TimbreFilter(self.encoder.latent_channels)
+        self.decoder = Decoder(self.encoder.out_channel, en_de_layers, kernel_size, n_blocks)
+
+    def forward(self, x):
+        x, concat_tensors = self.encoder(x)
+        x = self.intermediate(x)
+        x = self.decoder(x, concat_tensors)
+        return x

modules/pe/rmvpe/inference.py

@@ -0,0 +1,78 @@
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torchaudio.transforms import Resample
+
+from basics.base_pe import BasePE
+from utils.infer_utils import resample_align_curve
+from utils.pitch_utils import interp_f0
+from .constants import *
+from .model import E2E0
+from .spec import MelSpectrogram
+from .utils import to_local_average_f0, to_viterbi_f0
+
+
+class RMVPE(BasePE):
+    def __init__(self, model_path, hop_length=160):
+        self.resample_kernel = {}
+        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        self.model = E2E0(4, 1, (2, 2)).eval().to(self.device)
+        ckpt = torch.load(model_path, map_location=self.device)
+        self.model.load_state_dict(ckpt['model'], strict=False)
+        self.hop_length = hop_length
+        self.seg_length = 32 * hop_length
+        self.mel_extractor = MelSpectrogram(
+            N_MELS, SAMPLE_RATE, WINDOW_LENGTH, hop_length, None, MEL_FMIN, MEL_FMAX
+        ).to(self.device)
+
+    @torch.no_grad()
+    def mel2hidden(self, mel):
+        n_frames = mel.shape[-1]
+        mel = F.pad(mel, (0, 32 * ((n_frames - 1) // 32 + 1) - n_frames), mode='reflect')
+        hidden = self.model(mel)
+        return hidden[:, :n_frames]
+
+    def decode(self, hidden, thred=0.03, use_viterbi=False):
+        if use_viterbi:
+            f0 = to_viterbi_f0(hidden, thred=thred)
+        else:
+            f0 = to_local_average_f0(hidden, thred=thred)
+        return f0
+
+    def infer_from_audio(self, audio, sample_rate=16000, thred=0.03, use_viterbi=False):
+        audio = torch.from_numpy(audio).float().unsqueeze(0).to(self.device)
+        if sample_rate == 16000:
+            audio_res = audio
+        else:
+            key_str = str(sample_rate)
+            if key_str not in self.resample_kernel:
+                self.resample_kernel[key_str] = Resample(sample_rate, 16000, lowpass_filter_width=128)
+            self.resample_kernel[key_str] = self.resample_kernel[key_str].to(self.device)
+            audio_res = self.resample_kernel[key_str](audio)
+        B, T = audio_res.shape
+        n_frames = T // self.hop_length + 1 
+        T1 = T + self.hop_length
+        T_pad = self.seg_length * ((T1 - 1) // self.seg_length + 1) - T1
+        audio_res = F.pad(audio_res, (0, T_pad))
+        mel = self.mel_extractor(audio_res, center=True)
+        with torch.no_grad():
+            hidden = self.model(mel)
+        f0 = self.decode(hidden[:, :n_frames], thred=thred, use_viterbi=use_viterbi)
+        return f0
+
+    def get_pitch(
+            self, waveform, samplerate, length,
+            *, hop_size, f0_min=65, f0_max=1100,
+            speed=1, interp_uv=False
+    ):
+        f0 = self.infer_from_audio(waveform, sample_rate=samplerate)
+        uv = f0 == 0
+        f0, uv = interp_f0(f0, uv)
+
+        hop_size = int(np.round(hop_size * speed))
+        time_step = hop_size / samplerate
+        f0_res = resample_align_curve(f0, 0.01, time_step, length)
+        uv_res = resample_align_curve(uv.astype(np.float32), 0.01, time_step, length) > 0.5
+        if not interp_uv:
+            f0_res[uv_res] = 0
+        return f0_res, uv_res

modules/pe/rmvpe/model.py

@@ -0,0 +1,32 @@
+from torch import nn
+
+from .constants import *
+from .deepunet import DeepUnet0
+from .seq import BiGRU
+
+
+class E2E0(nn.Module):
+    def __init__(self, n_blocks, n_gru, kernel_size, en_de_layers=5, inter_layers=4, in_channels=1,
+                 en_out_channels=16):
+        super(E2E0, self).__init__()
+        self.unet = DeepUnet0(kernel_size, n_blocks, en_de_layers, inter_layers, in_channels, en_out_channels)
+        self.cnn = nn.Conv2d(en_out_channels, 3, (3, 3), padding=(1, 1))
+        if n_gru:
+            self.fc = nn.Sequential(
+                BiGRU(3 * N_MELS, 256, n_gru),
+                nn.Linear(512, N_CLASS),
+                nn.Dropout(0.25),
+                nn.Sigmoid()
+            )
+        else:
+            self.fc = nn.Sequential(
+                nn.Linear(3 * N_MELS, N_CLASS),
+                nn.Dropout(0.25),
+                nn.Sigmoid()
+            )
+
+    def forward(self, mel):
+        mel = mel.transpose(-1, -2).unsqueeze(1)
+        x = self.cnn(self.unet(mel)).transpose(1, 2).flatten(-2)
+        x = self.fc(x)
+        return x

modules/pe/rmvpe/seq.py

@@ -0,0 +1,10 @@
+import torch.nn as nn
+
+
+class BiGRU(nn.Module):
+    def __init__(self, input_features, hidden_features, num_layers):
+        super(BiGRU, self).__init__()
+        self.gru = nn.GRU(input_features, hidden_features, num_layers=num_layers, batch_first=True, bidirectional=True)
+
+    def forward(self, x):
+        return self.gru(x)[0]

modules/pe/rmvpe/spec.py

@@ -0,0 +1,68 @@
+import torch
+import numpy as np
+import torch.nn.functional as F
+from librosa.filters import mel
+
+
+class MelSpectrogram(torch.nn.Module):
+    def __init__(
+            self,
+            n_mel_channels,
+            sampling_rate,
+            win_length,
+            hop_length,
+            n_fft=None,
+            mel_fmin=0,
+            mel_fmax=None,
+            clamp=1e-5
+    ):
+        super().__init__()
+        n_fft = win_length if n_fft is None else n_fft
+        self.hann_window = {}
+        mel_basis = mel(
+            sr=sampling_rate,
+            n_fft=n_fft,
+            n_mels=n_mel_channels,
+            fmin=mel_fmin,
+            fmax=mel_fmax,
+            htk=True)
+        mel_basis = torch.from_numpy(mel_basis).float()
+        self.register_buffer("mel_basis", mel_basis)
+        self.n_fft = win_length if n_fft is None else n_fft
+        self.hop_length = hop_length
+        self.win_length = win_length
+        self.sampling_rate = sampling_rate
+        self.n_mel_channels = n_mel_channels
+        self.clamp = clamp
+
+    def forward(self, audio, keyshift=0, speed=1, center=True):
+        factor = 2 ** (keyshift / 12)
+        n_fft_new = int(np.round(self.n_fft * factor))
+        win_length_new = int(np.round(self.win_length * factor))
+        hop_length_new = int(np.round(self.hop_length * speed))
+
+        keyshift_key = str(keyshift) + '_' + str(audio.device)
+        if keyshift_key not in self.hann_window:
+            self.hann_window[keyshift_key] = torch.hann_window(win_length_new).to(audio.device)
+
+        fft = torch.stft(
+            audio,
+            n_fft=n_fft_new,
+            hop_length=hop_length_new,
+            win_length=win_length_new,
+            window=self.hann_window[keyshift_key],
+            center=center,
+            return_complex=True
+        )
+        magnitude = fft.abs()
+
+        if keyshift != 0:
+            size = self.n_fft // 2 + 1
+            resize = magnitude.size(1)
+            if resize < size:
+                magnitude = F.pad(magnitude, (0, 0, 0, size - resize))
+            magnitude = magnitude[:, :size, :] * self.win_length / win_length_new
+
+        mel_output = torch.matmul(self.mel_basis, magnitude)
+        log_mel_spec = torch.log(torch.clamp(mel_output, min=self.clamp))
+        return log_mel_spec

modules/pe/rmvpe/utils.py

@@ -0,0 +1,43 @@
+import librosa
+import numpy as np
+import torch
+
+from .constants import *
+
+
+def to_local_average_f0(hidden, center=None, thred=0.03):
+    idx = torch.arange(N_CLASS, device=hidden.device)[None, None, :]  # [B=1, T=1, N]
+    idx_cents = idx * 20 + CONST  # [B=1, N]
+    if center is None:
+        center = torch.argmax(hidden, dim=2, keepdim=True)  # [B, T, 1]
+    start = torch.clip(center - 4, min=0)  # [B, T, 1]
+    end = torch.clip(center + 5, max=N_CLASS)  # [B, T, 1]
+    idx_mask = (idx >= start) & (idx < end)  # [B, T, N]
+    weights = hidden * idx_mask  # [B, T, N]
+    product_sum = torch.sum(weights * idx_cents, dim=2)  # [B, T]
+    weight_sum = torch.sum(weights, dim=2)  # [B, T]
+    cents = product_sum / (weight_sum + (weight_sum == 0))  # avoid dividing by zero, [B, T]
+    f0 = 10 * 2 ** (cents / 1200)
+    uv = hidden.max(dim=2)[0] < thred  # [B, T]
+    f0 = f0 * ~uv
+    return f0.squeeze(0).cpu().numpy()
+
+
+def to_viterbi_f0(hidden, thred=0.03):
+    # Create viterbi transition matrix
+    if not hasattr(to_viterbi_f0, 'transition'):
+        xx, yy = np.meshgrid(range(N_CLASS), range(N_CLASS))
+        transition = np.maximum(30 - abs(xx - yy), 0)
+        transition = transition / transition.sum(axis=1, keepdims=True)
+        to_viterbi_f0.transition = transition
+
+    # Convert to probability
+    prob = hidden.squeeze(0).cpu().numpy()
+    prob = prob.T
+    prob = prob / prob.sum(axis=0)
+
+    # Perform viterbi decoding
+    path = librosa.sequence.viterbi(prob, to_viterbi_f0.transition).astype(np.int64)
+    center = torch.from_numpy(path).unsqueeze(0).unsqueeze(-1).to(hidden.device)
+
+    return to_local_average_f0(hidden, center=center, thred=thred)

modules/toplevel.py

@@ -0,0 +1,366 @@
+from typing import Dict
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+import modules.compat as compat
+from basics.base_module import CategorizedModule
+from modules.aux_decoder import AuxDecoderAdaptor
+from modules.commons.common_layers import (
+    NormalInitEmbedding as Embedding,
+    SinusoidalPosEmb, AdamWLinear,
+)
+from modules.core import (
+    GaussianDiffusion, PitchDiffusion, MultiVarianceDiffusion,
+    RectifiedFlow, PitchRectifiedFlow, MultiVarianceRectifiedFlow
+)
+from modules.fastspeech.acoustic_encoder import FastSpeech2Acoustic
+from modules.fastspeech.param_adaptor import ParameterAdaptorModule
+from modules.fastspeech.tts_modules import RhythmRegulator, LengthRegulator, StretchRegulator
+from modules.fastspeech.variance_encoder import FastSpeech2Variance, MelodyEncoder
+from utils.hparams import hparams
+
+
+class ShallowDiffusionOutput:
+    def __init__(self, *, aux_out=None, diff_out=None):
+        self.aux_out = aux_out
+        self.diff_out = diff_out
+
+
+class DiffSingerAcoustic(CategorizedModule, ParameterAdaptorModule):
+    @property
+    def category(self):
+        return 'acoustic'
+
+    def __init__(self, vocab_size, out_dims):
+        CategorizedModule.__init__(self)
+        ParameterAdaptorModule.__init__(self)
+        self.fs2 = FastSpeech2Acoustic(
+            vocab_size=vocab_size
+        )
+
+        self.use_shallow_diffusion = hparams.get('use_shallow_diffusion', False)
+        self.shallow_args = hparams.get('shallow_diffusion_args', {})
+        if self.use_shallow_diffusion:
+            self.train_aux_decoder = self.shallow_args['train_aux_decoder']
+            self.train_diffusion = self.shallow_args['train_diffusion']
+            self.aux_decoder_grad = self.shallow_args['aux_decoder_grad']
+            self.aux_decoder = AuxDecoderAdaptor(
+                in_dims=hparams['hidden_size'], out_dims=out_dims, num_feats=1,
+                spec_min=hparams['spec_min'], spec_max=hparams['spec_max'],
+                aux_decoder_arch=self.shallow_args['aux_decoder_arch'],
+                aux_decoder_args=self.shallow_args['aux_decoder_args']
+            )
+        self.diffusion_type = hparams.get('diffusion_type', 'ddpm')
+        self.backbone_type = compat.get_backbone_type(hparams)
+        self.backbone_args = compat.get_backbone_args(hparams, self.backbone_type)
+        if self.diffusion_type == 'ddpm':
+            self.diffusion = GaussianDiffusion(
+                out_dims=out_dims,
+                num_feats=1,
+                timesteps=hparams['timesteps'],
+                k_step=hparams['K_step'],
+                backbone_type=self.backbone_type,
+                backbone_args=self.backbone_args,
+                spec_min=hparams['spec_min'],
+                spec_max=hparams['spec_max']
+            )
+        elif self.diffusion_type == 'reflow':
+            self.diffusion = RectifiedFlow(
+                out_dims=out_dims,
+                num_feats=1,
+                t_start=hparams['T_start'],
+                time_scale_factor=hparams['time_scale_factor'],
+                backbone_type=self.backbone_type,
+                backbone_args=self.backbone_args,
+                spec_min=hparams['spec_min'],
+                spec_max=hparams['spec_max']
+            )
+        else:
+            raise NotImplementedError(self.diffusion_type)
+
+    def forward(
+            self, txt_tokens, mel2ph, f0, key_shift=None, speed=None,
+            spk_embed_id=None, languages=None, gt_mel=None, infer=True, **kwargs
+    ) -> ShallowDiffusionOutput:
+        condition = self.fs2(
+            txt_tokens, mel2ph, f0, key_shift=key_shift, speed=speed,
+            spk_embed_id=spk_embed_id, languages=languages,
+            **kwargs
+        )
+        if infer:
+            if self.use_shallow_diffusion:
+                aux_mel_pred = self.aux_decoder(condition, infer=True)
+                aux_mel_pred *= ((mel2ph > 0).float()[:, :, None])
+                if gt_mel is not None and self.shallow_args['val_gt_start']:
+                    src_mel = gt_mel
+                else:
+                    src_mel = aux_mel_pred
+            else:
+                aux_mel_pred = src_mel = None
+            mel_pred = self.diffusion(condition, src_spec=src_mel, infer=True)
+            mel_pred *= ((mel2ph > 0).float()[:, :, None])
+            return ShallowDiffusionOutput(aux_out=aux_mel_pred, diff_out=mel_pred)
+        else:
+            if self.use_shallow_diffusion:
+                if self.train_aux_decoder:
+                    aux_cond = condition * self.aux_decoder_grad + condition.detach() * (1 - self.aux_decoder_grad)
+                    aux_out = self.aux_decoder(aux_cond, infer=False)
+                else:
+                    aux_out = None
+                if self.train_diffusion:
+                    diff_out = self.diffusion(condition, gt_spec=gt_mel, infer=False)
+                else:
+                    diff_out = None
+                return ShallowDiffusionOutput(aux_out=aux_out, diff_out=diff_out)
+
+            else:
+                aux_out = None
+                diff_out = self.diffusion(condition, gt_spec=gt_mel, infer=False)
+                return ShallowDiffusionOutput(aux_out=aux_out, diff_out=diff_out)
+
+
+class DiffSingerVariance(CategorizedModule, ParameterAdaptorModule):
+    @property
+    def category(self):
+        return 'variance'
+
+    def __init__(self, vocab_size):
+        CategorizedModule.__init__(self)
+        ParameterAdaptorModule.__init__(self)
+        self.predict_dur = hparams['predict_dur']
+        self.predict_pitch = hparams['predict_pitch']
+
+        self.use_stretch_embed = hparams.get('use_stretch_embed', None)
+        assert self.use_stretch_embed is not None, "You may be loading an old version of the model checkpoint, which is incompatible with the new version due to some bug fixes. It is recommended to roll back to the old version (commit id: 6df0ee977c3728f14cb79c2db8b19df30b23a0bf)"
+        if self.use_stretch_embed and (self.predict_pitch or self.predict_variances):
+            self.sr = StretchRegulator()
+            self.stretch_embed = nn.Sequential(
+                SinusoidalPosEmb(hparams['hidden_size']),
+                nn.Linear(hparams['hidden_size'], hparams['hidden_size'] * 4),
+                nn.GELU(),
+                nn.Linear(hparams['hidden_size'] * 4, hparams['hidden_size']),
+            )
+            self.stretch_embed_rnn = nn.GRU(hparams['hidden_size'], hparams['hidden_size'], 1, batch_first=True)
+
+        self.use_spk_id = hparams['use_spk_id']
+        if self.use_spk_id:
+            self.spk_embed = Embedding(hparams['num_spk'], hparams['hidden_size'])
+
+        self.fs2 = FastSpeech2Variance(
+            vocab_size=vocab_size
+        )
+        self.rr = RhythmRegulator()
+        self.lr = LengthRegulator()
+        self.diffusion_type = hparams.get('diffusion_type', 'ddpm')
+        if self.predict_pitch:
+            self.use_melody_encoder = hparams.get('use_melody_encoder', False)
+            if self.use_melody_encoder:
+                self.melody_encoder = MelodyEncoder(enc_hparams=hparams['melody_encoder_args'])
+                self.delta_pitch_embed = AdamWLinear(1, hparams['hidden_size'])
+            else:
+                self.base_pitch_embed = AdamWLinear(1, hparams['hidden_size'])
+
+            self.pitch_retake_embed = Embedding(2, hparams['hidden_size'])
+            pitch_hparams = hparams['pitch_prediction_args']
+            self.pitch_backbone_type = compat.get_backbone_type(hparams, nested_config=pitch_hparams)
+            self.pitch_backbone_args = compat.get_backbone_args(pitch_hparams, backbone_type=self.pitch_backbone_type)
+            if self.diffusion_type == 'ddpm':
+                self.pitch_predictor = PitchDiffusion(
+                    vmin=pitch_hparams['pitd_norm_min'],
+                    vmax=pitch_hparams['pitd_norm_max'],
+                    cmin=pitch_hparams['pitd_clip_min'],
+                    cmax=pitch_hparams['pitd_clip_max'],
+                    repeat_bins=pitch_hparams['repeat_bins'],
+                    timesteps=hparams['timesteps'],
+                    k_step=hparams['K_step'],
+                    backbone_type=self.pitch_backbone_type,
+                    backbone_args=self.pitch_backbone_args
+                )
+            elif self.diffusion_type == 'reflow':
+                self.pitch_predictor = PitchRectifiedFlow(
+                    vmin=pitch_hparams['pitd_norm_min'],
+                    vmax=pitch_hparams['pitd_norm_max'],
+                    cmin=pitch_hparams['pitd_clip_min'],
+                    cmax=pitch_hparams['pitd_clip_max'],
+                    repeat_bins=pitch_hparams['repeat_bins'],
+                    time_scale_factor=hparams['time_scale_factor'],
+                    backbone_type=self.pitch_backbone_type,
+                    backbone_args=self.pitch_backbone_args
+                )
+            else:
+                raise ValueError(f"Invalid diffusion type: {self.diffusion_type}")
+
+        if self.predict_variances:
+            self.pitch_embed = AdamWLinear(1, hparams['hidden_size'])
+            self.variance_embeds = nn.ModuleDict({
+                v_name: AdamWLinear(1, hparams['hidden_size'])
+                for v_name in self.variance_prediction_list
+            })
+
+            if self.diffusion_type == 'ddpm':
+                self.variance_predictor = self.build_adaptor(cls=MultiVarianceDiffusion)
+            elif self.diffusion_type == 'reflow':
+                self.variance_predictor = self.build_adaptor(cls=MultiVarianceRectifiedFlow)
+            else:
+                raise NotImplementedError(self.diffusion_type)
+
+        self.use_variance_scaling = hparams.get('use_variance_scaling', False)
+        self.custom_variance_scaling_factor = {
+            'energy': 1. / 96,
+            'breathiness': 1. / 96,
+            'voicing': 1. / 96,
+            'tension': 0.1,
+            'key_shift': 1. / 12,
+            'speed': 1.
+        }
+        self.default_variance_scaling_factor = {
+            'energy': 1.,
+            'breathiness': 1.,
+            'voicing': 1.,
+            'tension': 1.,
+            'key_shift': 1.,
+            'speed': 1.
+        }
+        if self.use_variance_scaling:
+            self.variance_retake_scaling = self.custom_variance_scaling_factor
+        else:
+            self.variance_retake_scaling = self.default_variance_scaling_factor
+
+    def forward(
+            self, txt_tokens, midi, ph2word, ph_dur=None, word_dur=None, mel2ph=None,
+            note_midi=None, note_rest=None, note_dur=None, note_glide=None, mel2note=None,
+            base_pitch=None, pitch=None, pitch_expr=None, pitch_retake=None,
+            variance_retake: Dict[str, Tensor] = None,
+            spk_id=None, languages=None,
+            infer=True, **kwargs
+    ):
+        if self.use_spk_id:
+            ph_spk_mix_embed = kwargs.get('ph_spk_mix_embed')
+            spk_mix_embed = kwargs.get('spk_mix_embed')
+            if ph_spk_mix_embed is not None and spk_mix_embed is not None:
+                ph_spk_embed = ph_spk_mix_embed
+                spk_embed = spk_mix_embed
+            else:
+                ph_spk_embed = spk_embed = self.spk_embed(spk_id)[:, None, :]  # [B,] => [B, T=1, H]
+        else:
+            ph_spk_embed = spk_embed = None
+
+        encoder_out, dur_pred_out = self.fs2(
+            txt_tokens, midi=midi, ph2word=ph2word,
+            ph_dur=ph_dur, word_dur=word_dur,
+            spk_embed=ph_spk_embed, languages=languages,
+            infer=infer
+        )
+
+        if not self.predict_pitch and not self.predict_variances:
+            return dur_pred_out, None, ({} if infer else None)
+
+        if mel2ph is None and word_dur is not None:  # inference from file
+            dur_pred_align = self.rr(dur_pred_out, ph2word, word_dur)
+            mel2ph = self.lr(dur_pred_align)
+            mel2ph = F.pad(mel2ph, [0, base_pitch.shape[1] - mel2ph.shape[1]])
+
+        encoder_out = F.pad(encoder_out, [0, 0, 1, 0])
+        mel2ph_ = mel2ph[..., None].repeat([1, 1, hparams['hidden_size']])
+        condition = torch.gather(encoder_out, 1, mel2ph_)
+
+        if self.use_stretch_embed:
+            stretch = torch.round(1000 * self.sr(mel2ph, ph_dur))
+            if self.training and stretch.numel() > 1000:
+                # construct a phoneme stretching index lookup table with a total of 1001 indexes (0~1000)
+                table = self.stretch_embed(torch.arange(0, 1001, device=stretch.device))
+                stretch_embed = torch.index_select(table, 0, stretch.view(-1).long()).view_as(condition)
+            else:
+                stretch_embed = self.stretch_embed(stretch)
+            condition += stretch_embed
+            self.stretch_embed_rnn.flatten_parameters()
+            stretch_embed_rnn_out, _ = self.stretch_embed_rnn(condition)
+            condition = condition + stretch_embed_rnn_out
+            
+        if self.use_spk_id:
+            condition += spk_embed
+
+        if self.predict_pitch:
+            if self.use_melody_encoder:
+                melody_encoder_out = self.melody_encoder(
+                    note_midi, note_rest, note_dur,
+                    glide=note_glide
+                )
+                melody_encoder_out = F.pad(melody_encoder_out, [0, 0, 1, 0])
+                mel2note_ = mel2note[..., None].repeat([1, 1, hparams['hidden_size']])
+                melody_condition = torch.gather(melody_encoder_out, 1, mel2note_)
+                pitch_cond = condition + melody_condition
+            else:
+                pitch_cond = condition.clone()  # preserve the original tensor to avoid further inplace operations
+
+            retake_unset = pitch_retake is None
+            if retake_unset:
+                pitch_retake = torch.ones_like(mel2ph, dtype=torch.bool)
+
+            if pitch_expr is None:
+                pitch_retake_embed = self.pitch_retake_embed(pitch_retake.long())
+            else:
+                retake_true_embed = self.pitch_retake_embed(
+                    torch.ones(1, 1, dtype=torch.long, device=txt_tokens.device)
+                )  # [B=1, T=1] => [B=1, T=1, H]
+                retake_false_embed = self.pitch_retake_embed(
+                    torch.zeros(1, 1, dtype=torch.long, device=txt_tokens.device)
+                )  # [B=1, T=1] => [B=1, T=1, H]
+                pitch_expr = (pitch_expr * pitch_retake)[:, :, None]  # [B, T, 1]
+                pitch_retake_embed = pitch_expr * retake_true_embed + (1. - pitch_expr) * retake_false_embed
+
+            pitch_cond += pitch_retake_embed
+            if self.use_melody_encoder:
+                if retake_unset:  # generate from scratch
+                    delta_pitch_in = torch.zeros_like(base_pitch)
+                else:
+                    delta_pitch_in = (pitch - base_pitch) * ~pitch_retake
+                if self.use_variance_scaling:
+                    pitch_cond += self.delta_pitch_embed(delta_pitch_in[:, :, None] / 12)
+                else:
+                    pitch_cond += self.delta_pitch_embed(delta_pitch_in[:, :, None])
+            else:
+                if not retake_unset:  # retake
+                    base_pitch = base_pitch * pitch_retake + pitch * ~pitch_retake
+                if self.use_variance_scaling:
+                    pitch_cond += self.base_pitch_embed(base_pitch[:, :, None] / 128)
+                else:
+                    pitch_cond += self.base_pitch_embed(base_pitch[:, :, None])
+
+            if infer:
+                pitch_pred_out = self.pitch_predictor(pitch_cond, infer=True)
+            else:
+                pitch_pred_out = self.pitch_predictor(pitch_cond, pitch - base_pitch, infer=False)
+        else:
+            pitch_pred_out = None
+
+        if not self.predict_variances:
+            return dur_pred_out, pitch_pred_out, ({} if infer else None)
+
+        if pitch is None:
+            pitch = base_pitch + pitch_pred_out
+        if self.use_variance_scaling:
+            var_cond = condition + self.pitch_embed(pitch[:, :, None] / 12)
+        else:
+            var_cond = condition + self.pitch_embed(pitch[:, :, None])
+
+        variance_inputs = self.collect_variance_inputs(**kwargs)
+
+        if variance_retake is not None:
+            variance_embeds = [
+                self.variance_embeds[v_name](v_input[:, :, None] * self.variance_retake_scaling[v_name]) * ~variance_retake[v_name][:, :, None]
+                for v_name, v_input in zip(self.variance_prediction_list, variance_inputs)
+            ]
+            var_cond += torch.stack(variance_embeds, dim=-1).sum(-1)
+
+        variance_outputs = self.variance_predictor(var_cond, variance_inputs, infer=infer)
+
+        if infer:
+            variances_pred_out = self.collect_variance_outputs(variance_outputs)
+        else:
+            variances_pred_out = variance_outputs
+
+        return dur_pred_out, pitch_pred_out, variances_pred_out

modules/vocoders/__init__.py

@@ -0,0 +1,2 @@
+from modules.vocoders import ddsp
+from modules.vocoders import nsf_hifigan

modules/vocoders/ddsp.py

@@ -0,0 +1,120 @@
+import pathlib
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+import yaml
+from librosa.filters import mel as librosa_mel_fn
+
+from basics.base_vocoder import BaseVocoder
+from modules.vocoders.registry import register_vocoder
+from utils.hparams import hparams
+
+
+class DotDict(dict):
+    def __getattr__(*args):
+        val = dict.get(*args)
+        return DotDict(val) if type(val) is dict else val
+
+    __setattr__ = dict.__setitem__
+    __delattr__ = dict.__delitem__
+
+
+def load_model(model_path: pathlib.Path, device='cpu'):
+    config_file = model_path.with_name('config.yaml')
+    with open(config_file, "r") as config:
+        args = yaml.safe_load(config)
+    args = DotDict(args)
+
+    # load model
+    print(' [Loading] ' + str(model_path))
+    model = torch.jit.load(model_path, map_location=torch.device(device))
+    model.eval()
+
+    return model, args
+
+
+@register_vocoder
+class DDSP(BaseVocoder):
+    def __init__(self, device='cpu'):
+        self.device = device
+        model_path = pathlib.Path(hparams['vocoder_ckpt'])
+        assert model_path.exists(), 'DDSP model file is not found!'
+        self.model, self.args = load_model(model_path, device=self.device)
+
+    def to_device(self, device):
+        pass
+
+    def get_device(self):
+        return self.device
+
+    def spec2wav_torch(self, mel, f0):  # mel: [B, T, bins] f0: [B, T]
+        if self.args.data.sampling_rate != hparams['audio_sample_rate']:
+            print('Mismatch parameters: hparams[\'audio_sample_rate\']=', hparams['audio_sample_rate'], '!=',
+                  self.args.data.sampling_rate, '(vocoder)')
+        if self.args.data.n_mels != hparams['audio_num_mel_bins']:
+            print('Mismatch parameters: hparams[\'audio_num_mel_bins\']=', hparams['audio_num_mel_bins'], '!=',
+                  self.args.data.n_mels, '(vocoder)')
+        if self.args.data.n_fft != hparams['fft_size']:
+            print('Mismatch parameters: hparams[\'fft_size\']=', hparams['fft_size'], '!=', self.args.data.n_fft,
+                  '(vocoder)')
+        if self.args.data.win_length != hparams['win_size']:
+            print('Mismatch parameters: hparams[\'win_size\']=', hparams['win_size'], '!=', self.args.data.win_length,
+                  '(vocoder)')
+        if self.args.data.block_size != hparams['hop_size']:
+            print('Mismatch parameters: hparams[\'hop_size\']=', hparams['hop_size'], '!=', self.args.data.block_size,
+                  '(vocoder)')
+        if self.args.data.mel_fmin != hparams['fmin']:
+            print('Mismatch parameters: hparams[\'fmin\']=', hparams['fmin'], '!=', self.args.data.mel_fmin,
+                  '(vocoder)')
+        if self.args.data.mel_fmax != hparams['fmax']:
+            print('Mismatch parameters: hparams[\'fmax\']=', hparams['fmax'], '!=', self.args.data.mel_fmax,
+                  '(vocoder)')
+        with torch.no_grad():
+            mel = mel.to(self.device)
+            mel_base = hparams.get('mel_base', 10)
+            if mel_base != 'e':
+                assert mel_base in [10, '10'], "mel_base must be 'e', '10' or 10."
+            else:
+                # log mel to log10 mel
+                mel = 0.434294 * mel
+            f0 = f0.unsqueeze(-1).to(self.device)
+            signal, _, (s_h, s_n) = self.model(mel, f0)
+            signal = signal.view(-1)
+        return signal
+
+    def spec2wav(self, mel, f0):
+        if self.args.data.sampling_rate != hparams['audio_sample_rate']:
+            print('Mismatch parameters: hparams[\'audio_sample_rate\']=', hparams['audio_sample_rate'], '!=',
+                  self.args.data.sampling_rate, '(vocoder)')
+        if self.args.data.n_mels != hparams['audio_num_mel_bins']:
+            print('Mismatch parameters: hparams[\'audio_num_mel_bins\']=', hparams['audio_num_mel_bins'], '!=',
+                  self.args.data.n_mels, '(vocoder)')
+        if self.args.data.n_fft != hparams['fft_size']:
+            print('Mismatch parameters: hparams[\'fft_size\']=', hparams['fft_size'], '!=', self.args.data.n_fft,
+                  '(vocoder)')
+        if self.args.data.win_length != hparams['win_size']:
+            print('Mismatch parameters: hparams[\'win_size\']=', hparams['win_size'], '!=', self.args.data.win_length,
+                  '(vocoder)')
+        if self.args.data.block_size != hparams['hop_size']:
+            print('Mismatch parameters: hparams[\'hop_size\']=', hparams['hop_size'], '!=', self.args.data.block_size,
+                  '(vocoder)')
+        if self.args.data.mel_fmin != hparams['fmin']:
+            print('Mismatch parameters: hparams[\'fmin\']=', hparams['fmin'], '!=', self.args.data.mel_fmin,
+                  '(vocoder)')
+        if self.args.data.mel_fmax != hparams['fmax']:
+            print('Mismatch parameters: hparams[\'fmax\']=', hparams['fmax'], '!=', self.args.data.mel_fmax,
+                  '(vocoder)')
+        with torch.no_grad():
+            mel = torch.FloatTensor(mel).unsqueeze(0).to(self.device)
+            mel_base = hparams.get('mel_base', 10)
+            if mel_base != 'e':
+                assert mel_base in [10, '10'], "mel_base must be 'e', '10' or 10."
+            else:
+                # log mel to log10 mel
+                mel = 0.434294 * mel
+            f0 = torch.FloatTensor(f0).unsqueeze(0).unsqueeze(-1).to(self.device)
+            signal, _, (s_h, s_n) = self.model(mel, f0)
+            signal = signal.view(-1)
+        wav_out = signal.cpu().numpy()
+        return wav_out

modules/vocoders/nsf_hifigan.py

@@ -0,0 +1,104 @@
+import pathlib
+
+import torch
+
+try:
+    from lightning.pytorch.utilities.rank_zero import rank_zero_info
+except ModuleNotFoundError:
+    rank_zero_info = print
+
+from modules.nsf_hifigan.models import load_model
+from basics.base_vocoder import BaseVocoder
+from modules.vocoders.registry import register_vocoder
+from utils.hparams import hparams
+
+
+@register_vocoder
+class NsfHifiGAN(BaseVocoder):
+    def __init__(self):
+        model_path = pathlib.Path(hparams['vocoder_ckpt'])
+        if not model_path.exists():
+            raise FileNotFoundError(
+                f'NSF-HiFiGAN vocoder model is not found at \'{model_path}\'. '
+                'Please follow instructions in docs/BestPractices.md#vocoders to get one.'
+            )
+        rank_zero_info(f'| Load HifiGAN: {model_path}')
+        self.model, self.h = load_model(model_path)
+
+    @property
+    def device(self):
+        return next(self.model.parameters()).device
+
+    def to_device(self, device):
+        self.model.to(device)
+
+    def get_device(self):
+        return self.device
+
+    def spec2wav_torch(self, mel, **kwargs):  # mel: [B, T, bins]
+        if self.h.sampling_rate != hparams['audio_sample_rate']:
+            print('Mismatch parameters: hparams[\'audio_sample_rate\']=', hparams['audio_sample_rate'], '!=',
+                  self.h.sampling_rate, '(vocoder)')
+        if self.h.num_mels != hparams['audio_num_mel_bins']:
+            print('Mismatch parameters: hparams[\'audio_num_mel_bins\']=', hparams['audio_num_mel_bins'], '!=',
+                  self.h.num_mels, '(vocoder)')
+        if self.h.n_fft != hparams['fft_size']:
+            print('Mismatch parameters: hparams[\'fft_size\']=', hparams['fft_size'], '!=', self.h.n_fft, '(vocoder)')
+        if self.h.win_size != hparams['win_size']:
+            print('Mismatch parameters: hparams[\'win_size\']=', hparams['win_size'], '!=', self.h.win_size,
+                  '(vocoder)')
+        if self.h.hop_size != hparams['hop_size']:
+            print('Mismatch parameters: hparams[\'hop_size\']=', hparams['hop_size'], '!=', self.h.hop_size,
+                  '(vocoder)')
+        if self.h.fmin != hparams['fmin']:
+            print('Mismatch parameters: hparams[\'fmin\']=', hparams['fmin'], '!=', self.h.fmin, '(vocoder)')
+        if self.h.fmax != hparams['fmax']:
+            print('Mismatch parameters: hparams[\'fmax\']=', hparams['fmax'], '!=', self.h.fmax, '(vocoder)')
+        with torch.no_grad():
+            c = mel.transpose(2, 1)  # [B, T, bins]
+            mel_base = hparams.get('mel_base', 10)
+            if mel_base != 'e':
+                assert mel_base in [10, '10'], "mel_base must be 'e', '10' or 10."
+                # log10 to log mel
+                c = 2.30259 * c
+            f0 = kwargs.get('f0')  # [B, T]
+            if f0 is not None:
+                y = self.model(c, f0).view(-1)
+            else:
+                y = self.model(c).view(-1)
+        return y
+
+    def spec2wav(self, mel, **kwargs):
+        if self.h.sampling_rate != hparams['audio_sample_rate']:
+            print('Mismatch parameters: hparams[\'audio_sample_rate\']=', hparams['audio_sample_rate'], '!=',
+                  self.h.sampling_rate, '(vocoder)')
+        if self.h.num_mels != hparams['audio_num_mel_bins']:
+            print('Mismatch parameters: hparams[\'audio_num_mel_bins\']=', hparams['audio_num_mel_bins'], '!=',
+                  self.h.num_mels, '(vocoder)')
+        if self.h.n_fft != hparams['fft_size']:
+            print('Mismatch parameters: hparams[\'fft_size\']=', hparams['fft_size'], '!=', self.h.n_fft, '(vocoder)')
+        if self.h.win_size != hparams['win_size']:
+            print('Mismatch parameters: hparams[\'win_size\']=', hparams['win_size'], '!=', self.h.win_size,
+                  '(vocoder)')
+        if self.h.hop_size != hparams['hop_size']:
+            print('Mismatch parameters: hparams[\'hop_size\']=', hparams['hop_size'], '!=', self.h.hop_size,
+                  '(vocoder)')
+        if self.h.fmin != hparams['fmin']:
+            print('Mismatch parameters: hparams[\'fmin\']=', hparams['fmin'], '!=', self.h.fmin, '(vocoder)')
+        if self.h.fmax != hparams['fmax']:
+            print('Mismatch parameters: hparams[\'fmax\']=', hparams['fmax'], '!=', self.h.fmax, '(vocoder)')
+        with torch.no_grad():
+            c = torch.FloatTensor(mel).unsqueeze(0).transpose(2, 1).to(self.device)
+            mel_base = hparams.get('mel_base', 10)
+            if mel_base != 'e':
+                assert mel_base in [10, '10'], "mel_base must be 'e', '10' or 10."
+                # log10 to log mel
+                c = 2.30259 * c
+            f0 = kwargs.get('f0')
+            if f0 is not None:
+                f0 = torch.FloatTensor(f0[None, :]).to(self.device)
+                y = self.model(c, f0).view(-1)
+            else:
+                y = self.model(c).view(-1)
+        wav_out = y.cpu().numpy()
+        return wav_out

modules/vocoders/registry.py

@@ -0,0 +1,21 @@
+import importlib
+
+
+VOCODERS = {}
+
+
+def register_vocoder(cls):
+    VOCODERS[cls.__name__.lower()] = cls
+    VOCODERS[cls.__name__] = cls
+    return cls
+
+
+def get_vocoder_cls(hparams):
+    if hparams['vocoder'] in VOCODERS:
+        return VOCODERS[hparams['vocoder']]
+    else:
+        vocoder_cls = hparams['vocoder']
+        pkg = ".".join(vocoder_cls.split(".")[:-1])
+        cls_name = vocoder_cls.split(".")[-1]
+        vocoder_cls = getattr(importlib.import_module(pkg), cls_name)
+        return vocoder_cls