Upload 39 files

requirements.txt

@@ -1,16 +1,15 @@
 joblib>=1.1.0
-numba==0.56.4
+numba
 numpy==1.23.5
 scipy
 librosa==0.9.1
-llvmlite==0.39.0
-fairseq==0.12.2
-faiss-cpu==1.7.3
-gradio==3.34.0
+llvmlite
+fairseq
+faiss-cpu
+gradio
 Cython
 pydub>=0.25.1
 soundfile>=0.12.1
-ffmpeg-python>=0.2.0
 tensorboardX
 Jinja2>=3.1.2
 json5
@@ -41,8 +40,8 @@ httpx
 onnxruntime; sys_platform == 'darwin'
 onnxruntime-gpu; sys_platform != 'darwin'
 torchcrepe==0.0.20
-fastapi==0.88
+fastapi
 torchfcpe
-ffmpy==0.3.1
 python-dotenv>=1.0.0
 av
+pybase16384

rvc/__init__.py

@@ -0,0 +1,4 @@
+from . import ipex
+import sys
+
+del sys.modules["rvc.ipex"]

rvc/f0/__init__.py

@@ -0,0 +1,10 @@
+from .f0 import F0Predictor
+
+from .crepe import CRePE
+from .dio import Dio
+from .fcpe import FCPE
+from .harvest import Harvest
+from .pm import PM
+from .rmvpe import RMVPE
+
+__all__ = ["F0Predictor", "CRePE", "Dio", "FCPE", "Harvest", "PM", "RMVPE"]

rvc/f0/crepe.py

@@ -0,0 +1,56 @@
+from typing import Any, Optional, Union
+
+import numpy as np
+import torch
+import torchcrepe
+
+from .f0 import F0Predictor
+
+
+class CRePE(F0Predictor):
+    def __init__(
+        self,
+        hop_length=512,
+        f0_min=50,
+        f0_max=1100,
+        sampling_rate=44100,
+        device="cpu",
+    ):
+        if "privateuseone" in str(device):
+            device = "cpu"
+        super().__init__(
+            hop_length,
+            f0_min,
+            f0_max,
+            sampling_rate,
+            device,
+        )
+
+    def compute_f0(
+        self,
+        wav: np.ndarray,
+        p_len: Optional[int] = None,
+        filter_radius: Optional[Union[int, float]] = None,
+    ):
+        if p_len is None:
+            p_len = wav.shape[0] // self.hop_length
+        if not torch.is_tensor(wav):
+            wav = torch.from_numpy(wav)
+        # Pick a batch size that doesn't cause memory errors on your gpu
+        batch_size = 512
+        # Compute pitch using device 'device'
+        f0, pd = torchcrepe.predict(
+            wav.float().to(self.device).unsqueeze(dim=0),
+            self.sampling_rate,
+            self.hop_length,
+            self.f0_min,
+            self.f0_max,
+            batch_size=batch_size,
+            device=self.device,
+            return_periodicity=True,
+        )
+        pd = torchcrepe.filter.median(pd, 3)
+        f0 = torchcrepe.filter.mean(f0, 3)
+        f0[pd < 0.1] = 0
+        f0 = f0[0].cpu().numpy()
+        return self._interpolate_f0(self._resize_f0(f0, p_len))[0]

rvc/f0/deepunet.py

@@ -0,0 +1,217 @@
+from typing import List, Tuple, Union
+
+import torch
+import torch.nn as nn
+
+
+class ConvBlockRes(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        momentum: float = 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(),
+        )
+        # self.shortcut:Optional[nn.Module] = None
+        if in_channels != out_channels:
+            self.shortcut = nn.Conv2d(in_channels, out_channels, (1, 1))
+
+    def forward(self, x: torch.Tensor):
+        if not hasattr(self, "shortcut"):
+            return self.conv(x) + x
+        else:
+            return self.conv(x) + self.shortcut(x)
+
+
+class Encoder(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        in_size: int,
+        n_encoders: int,
+        kernel_size: Tuple[int, int],
+        n_blocks: int,
+        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()
+        for _ in range(self.n_encoders):
+            self.layers.append(
+                ResEncoderBlock(
+                    in_channels, out_channels, kernel_size, n_blocks, momentum=momentum
+                )
+            )
+            in_channels = out_channels
+            out_channels *= 2
+            in_size //= 2
+        self.out_size = in_size
+        self.out_channel = out_channels
+
+    def __call__(self, x: torch.Tensor) -> Tuple[torch.Tensor, List[torch.Tensor]]:
+        return super().__call__(x)
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, List[torch.Tensor]]:
+        concat_tensors: List[torch.Tensor] = []
+        x = self.bn(x)
+        for layer in self.layers:
+            t, x = layer(x)
+            concat_tensors.append(t)
+        return x, concat_tensors
+
+
+class ResEncoderBlock(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: Tuple[int, int],
+        n_blocks=1,
+        momentum=0.01,
+    ):
+        super(ResEncoderBlock, self).__init__()
+        self.n_blocks = n_blocks
+        self.kernel_size = kernel_size
+
+        self.conv = nn.ModuleList()
+        self.conv.append(ConvBlockRes(in_channels, out_channels, momentum))
+        for _ in range(n_blocks - 1):
+            self.conv.append(ConvBlockRes(out_channels, out_channels, momentum))
+
+        if self.kernel_size is not None:
+            self.pool = nn.AvgPool2d(kernel_size=kernel_size)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        for conv in self.conv:
+            x = conv(x)
+        if self.kernel_size is not None:
+            return x, self.pool(x)
+        return x
+
+
+class Intermediate(nn.Module):
+    def __init__(self, in_channels, out_channels, n_inters, n_blocks, momentum=0.01):
+        super(Intermediate, self).__init__()
+
+        self.layers = nn.ModuleList()
+        self.layers.append(
+            ResEncoderBlock(in_channels, out_channels, None, n_blocks, momentum)
+        )
+        for _ in range(n_inters - 1):
+            self.layers.append(
+                ResEncoderBlock(out_channels, out_channels, None, n_blocks, momentum)
+            )
+
+    def forward(self, x):
+        for layer in self.layers:
+            x = layer(x)
+        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.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 _ 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 conv2 in self.conv2:
+            x = conv2(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 _ 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: torch.Tensor, concat_tensors: List[torch.Tensor]):
+        for i, layer in enumerate(self.layers):
+            x = layer(x, concat_tensors[-1 - i])
+        return x
+
+
+class DeepUnet(nn.Module):
+    def __init__(
+        self,
+        kernel_size: Tuple[int, int],
+        n_blocks: int,
+        en_de_layers=5,
+        inter_layers=4,
+        in_channels=1,
+        en_out_channels=16,
+    ):
+        super(DeepUnet, self).__init__()
+        self.encoder = Encoder(
+            in_channels, 128, 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.decoder = Decoder(
+            self.encoder.out_channel, en_de_layers, kernel_size, n_blocks
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x, concat_tensors = self.encoder(x)
+        x = self.intermediate(x)
+        x = self.decoder(x, concat_tensors)
+        return x

rvc/f0/dio.py

@@ -0,0 +1,31 @@
+from typing import Any, Optional, Union
+
+import numpy as np
+import pyworld
+
+from .f0 import F0Predictor
+
+
+class Dio(F0Predictor):
+    def __init__(self, hop_length=512, f0_min=50, f0_max=1100, sampling_rate=44100):
+        super().__init__(hop_length, f0_min, f0_max, sampling_rate)
+
+    def compute_f0(
+        self,
+        wav: np.ndarray,
+        p_len: Optional[int] = None,
+        filter_radius: Optional[Union[int, float]] = None,
+    ):
+        if p_len is None:
+            p_len = wav.shape[0] // self.hop_length
+        f0, t = pyworld.dio(
+            wav.astype(np.double),
+            fs=self.sampling_rate,
+            f0_floor=self.f0_min,
+            f0_ceil=self.f0_max,
+            frame_period=1000 * self.hop_length / self.sampling_rate,
+        )
+        f0 = pyworld.stonemask(wav.astype(np.double), f0, t, self.sampling_rate)
+        for index, pitch in enumerate(f0):
+            f0[index] = round(pitch, 1)
+        return self._interpolate_f0(self._resize_f0(f0, p_len))[0]

rvc/f0/e2e.py

@@ -0,0 +1,67 @@
+from typing import Tuple
+
+import torch.nn as nn
+
+from .deepunet import DeepUnet
+
+
+class E2E(nn.Module):
+    def __init__(
+        self,
+        n_blocks: int,
+        n_gru: int,
+        kernel_size: Tuple[int, int],
+        en_de_layers=5,
+        inter_layers=4,
+        in_channels=1,
+        en_out_channels=16,
+    ):
+        super(E2E, self).__init__()
+
+        self.unet = DeepUnet(
+            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(
+                self.BiGRU(3 * 128, 256, n_gru),
+                nn.Linear(512, 360),
+                nn.Dropout(0.25),
+                nn.Sigmoid(),
+            )
+        else:
+            self.fc = nn.Sequential(
+                nn.Linear(3 * nn.N_MELS, nn.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
+
+    class BiGRU(nn.Module):
+        def __init__(
+            self,
+            input_features: int,
+            hidden_features: int,
+            num_layers: int,
+        ):
+            super().__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]

rvc/f0/f0.py

@@ -0,0 +1,78 @@
+from typing import Optional, Union
+
+import torch
+import numpy as np
+
+
+class F0Predictor(object):
+    def __init__(
+        self,
+        hop_length=512,
+        f0_min=50,
+        f0_max=1100,
+        sampling_rate=44100,
+        device: Optional[str] = None,
+    ):
+        self.hop_length = hop_length
+        self.f0_min = f0_min
+        self.f0_max = f0_max
+        self.sampling_rate = sampling_rate
+        if device is None:
+            device = "cuda:0" if torch.cuda.is_available() else "cpu"
+        self.device = device
+
+    def compute_f0(
+        self,
+        wav: np.ndarray,
+        p_len: Optional[int] = None,
+        filter_radius: Optional[Union[int, float]] = None,
+    ): ...
+
+    def _interpolate_f0(self, f0: np.ndarray):
+        """
+        对F0进行插值处理
+        """
+
+        data = np.reshape(f0, (f0.size, 1))
+
+        vuv_vector = np.zeros((data.size, 1), dtype=np.float32)
+        vuv_vector[data > 0.0] = 1.0
+        vuv_vector[data <= 0.0] = 0.0
+
+        ip_data = data
+
+        frame_number = data.size
+        last_value = 0.0
+        for i in range(frame_number):
+            if data[i] <= 0.0:
+                j = i + 1
+                for j in range(i + 1, frame_number):
+                    if data[j] > 0.0:
+                        break
+                if j < frame_number - 1:
+                    if last_value > 0.0:
+                        step = (data[j] - data[i - 1]) / float(j - i)
+                        for k in range(i, j):
+                            ip_data[k] = data[i - 1] + step * (k - i + 1)
+                    else:
+                        for k in range(i, j):
+                            ip_data[k] = data[j]
+                else:
+                    for k in range(i, frame_number):
+                        ip_data[k] = last_value
+            else:
+                ip_data[i] = data[i]  # 这里可能存在一个没有必要的拷贝
+                last_value = data[i]
+
+        return ip_data[:, 0], vuv_vector[:, 0]
+
+    def _resize_f0(self, x: np.ndarray, target_len: int):
+        source = np.array(x)
+        source[source < 0.001] = np.nan
+        target = np.interp(
+            np.arange(0, len(source) * target_len, len(source)) / target_len,
+            np.arange(0, len(source)),
+            source,
+        )
+        res = np.nan_to_num(target)
+        return res

rvc/f0/fcpe.py

@@ -0,0 +1,53 @@
+from typing import Optional, Union
+
+import numpy as np
+import torch
+
+from .f0 import F0Predictor
+
+
+class FCPE(F0Predictor):
+    def __init__(
+        self,
+        hop_length=512,
+        f0_min=50,
+        f0_max=1100,
+        sampling_rate=44100,
+        device="cpu",
+    ):
+        super().__init__(
+            hop_length,
+            f0_min,
+            f0_max,
+            sampling_rate,
+            device,
+        )
+
+        from torchfcpe import (
+            spawn_bundled_infer_model,
+        )  # must be imported at here, or it will cause fairseq crash on training
+
+        self.model = spawn_bundled_infer_model(self.device)
+
+    def compute_f0(
+        self,
+        wav: np.ndarray,
+        p_len: Optional[int] = None,
+        filter_radius: Optional[Union[int, float]] = 0.006,
+    ):
+        if p_len is None:
+            p_len = wav.shape[0] // self.hop_length
+        if not torch.is_tensor(wav):
+            wav = torch.from_numpy(wav)
+        f0 = (
+            self.model.infer(
+                wav.float().to(self.device).unsqueeze(0),
+                sr=self.sampling_rate,
+                decoder_mode="local_argmax",
+                threshold=filter_radius,
+            )
+            .squeeze()
+            .cpu()
+            .numpy()
+        )
+        return self._interpolate_f0(self._resize_f0(f0, p_len))[0]

rvc/f0/harvest.py

@@ -0,0 +1,32 @@
+from typing import Any, Optional, Union
+
+import numpy as np
+import pyworld
+from scipy import signal
+
+from .f0 import F0Predictor
+
+
+class Harvest(F0Predictor):
+    def __init__(self, hop_length=512, f0_min=50, f0_max=1100, sampling_rate=44100):
+        super().__init__(hop_length, f0_min, f0_max, sampling_rate)
+
+    def compute_f0(
+        self,
+        wav: np.ndarray,
+        p_len: Optional[int] = None,
+        filter_radius: Optional[Union[int, float]] = None,
+    ):
+        if p_len is None:
+            p_len = wav.shape[0] // self.hop_length
+        f0, t = pyworld.harvest(
+            wav.astype(np.double),
+            fs=self.sampling_rate,
+            f0_ceil=self.f0_max,
+            f0_floor=self.f0_min,
+            frame_period=1000 * self.hop_length / self.sampling_rate,
+        )
+        f0 = pyworld.stonemask(wav.astype(np.double), f0, t, self.sampling_rate)
+        if filter_radius is not None and filter_radius > 2:
+            f0 = signal.medfilt(f0, filter_radius)
+        return self._interpolate_f0(self._resize_f0(f0, p_len))[0]

rvc/f0/mel.py

@@ -0,0 +1,71 @@
+from typing import Optional
+
+import torch
+import numpy as np
+from librosa.filters import mel
+
+from .stft import STFT
+
+
+class MelSpectrogram(torch.nn.Module):
+    def __init__(
+        self,
+        is_half: bool,
+        n_mel_channels: int,
+        sampling_rate: int,
+        win_length: int,
+        hop_length: int,
+        n_fft: Optional[int] = None,
+        mel_fmin: int = 0,
+        mel_fmax: int = None,
+        clamp: float = 1e-5,
+        device=torch.device("cpu"),
+    ):
+        super().__init__()
+        if n_fft is None:
+            n_fft = win_length
+        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 = n_fft
+        self.hop_length = hop_length
+        self.win_length = win_length
+        self.clamp = clamp
+        self.is_half = is_half
+
+        self.stft = STFT(
+            filter_length=n_fft,
+            hop_length=hop_length,
+            win_length=win_length,
+            window="hann",
+            use_torch_stft="privateuseone" not in str(device),
+        ).to(device)
+
+    def forward(
+        self,
+        audio: torch.Tensor,
+        keyshift=0,
+        speed=1,
+        center=True,
+    ):
+        factor = 2 ** (keyshift / 12)
+        win_length_new = int(np.round(self.win_length * factor))
+        magnitude = self.stft(audio, keyshift, speed, center)
+        if keyshift != 0:
+            size = self.n_fft // 2 + 1
+            resize = magnitude.size(1)
+            if resize < size:
+                magnitude = torch.nn.functional.pad(magnitude, (0, 0, 0, size - resize))
+            magnitude = magnitude[:, :size, :] * self.win_length / win_length_new
+        mel_output = torch.matmul(self.mel_basis, magnitude)
+        if self.is_half:
+            mel_output = mel_output.half()
+        log_mel_spec = torch.log(torch.clamp(mel_output, min=self.clamp))
+        return log_mel_spec

rvc/f0/models.py

@@ -0,0 +1,16 @@
+import torch
+
+
+def get_rmvpe(
+    model_path="assets/rmvpe/rmvpe.pt", device=torch.device("cpu"), is_half=False
+):
+    from rvc.f0.e2e import E2E
+
+    model = E2E(4, 1, (2, 2))
+    ckpt = torch.load(model_path, map_location=device)
+    model.load_state_dict(ckpt)
+    model.eval()
+    if is_half:
+        model = model.half()
+    model = model.to(device)
+    return model

rvc/f0/pm.py

@@ -0,0 +1,39 @@
+from typing import Any, Optional
+
+import numpy as np
+import parselmouth
+
+from .f0 import F0Predictor
+
+
+class PM(F0Predictor):
+    def __init__(self, hop_length=512, f0_min=50, f0_max=1100, sampling_rate=44100):
+        super().__init__(hop_length, f0_min, f0_max, sampling_rate)
+
+    def compute_f0(
+        self,
+        wav: np.ndarray,
+        p_len: Optional[int] = None,
+        filter_radius: Optional[int] = None,
+    ):
+        x = wav
+        if p_len is None:
+            p_len = x.shape[0] // self.hop_length
+        else:
+            assert abs(p_len - x.shape[0] // self.hop_length) < 4, "pad length error"
+        time_step = self.hop_length / self.sampling_rate * 1000
+        f0 = (
+            parselmouth.Sound(x, self.sampling_rate)
+            .to_pitch_ac(
+                time_step=time_step / 1000,
+                voicing_threshold=0.6,
+                pitch_floor=self.f0_min,
+                pitch_ceiling=self.f0_max,
+            )
+            .selected_array["frequency"]
+        )
+
+        pad_size = (p_len - len(f0) + 1) // 2
+        if pad_size > 0 or p_len - len(f0) - pad_size > 0:
+            f0 = np.pad(f0, [[pad_size, p_len - len(f0) - pad_size]], mode="constant")
+        return self._interpolate_f0(f0)[0]

rvc/f0/rmvpe.py

@@ -0,0 +1,162 @@
+from io import BytesIO
+import os
+from typing import Any, Optional, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from rvc.jit import load_inputs, get_jit_model, export_jit_model, save_pickle
+
+from .mel import MelSpectrogram
+from .f0 import F0Predictor
+from .models import get_rmvpe
+
+
+def rmvpe_jit_export(
+    model_path: str,
+    mode: str = "script",
+    inputs_path: str = None,
+    save_path: str = None,
+    device=torch.device("cpu"),
+    is_half=False,
+):
+    if not save_path:
+        save_path = model_path.rstrip(".pth")
+        save_path += ".half.jit" if is_half else ".jit"
+    if "cuda" in str(device) and ":" not in str(device):
+        device = torch.device("cuda:0")
+
+    model = get_rmvpe(model_path, device, is_half)
+    inputs = None
+    if mode == "trace":
+        inputs = load_inputs(inputs_path, device, is_half)
+    ckpt = export_jit_model(model, mode, inputs, device, is_half)
+    ckpt["device"] = str(device)
+    save_pickle(ckpt, save_path)
+    return ckpt
+
+
+class RMVPE(F0Predictor):
+    def __init__(
+        self,
+        model_path: str,
+        is_half: bool,
+        device: str,
+        use_jit=False,
+    ):
+        hop_length = 160
+        f0_min = 30
+        f0_max = 8000
+        sampling_rate = 16000
+
+        super().__init__(
+            hop_length,
+            f0_min,
+            f0_max,
+            sampling_rate,
+            device,
+        )
+
+        self.is_half = is_half
+        cents_mapping = 20 * np.arange(360) + 1997.3794084376191
+        self.cents_mapping = np.pad(cents_mapping, (4, 4))  # 368
+
+        self.mel_extractor = MelSpectrogram(
+            is_half=is_half,
+            n_mel_channels=128,
+            sampling_rate=sampling_rate,
+            win_length=1024,
+            hop_length=hop_length,
+            mel_fmin=f0_min,
+            mel_fmax=f0_max,
+            device=self.device,
+        ).to(self.device)
+
+        if "privateuseone" in str(self.device):
+            import onnxruntime as ort
+
+            self.model = ort.InferenceSession(
+                "%s/rmvpe.onnx" % os.environ["rmvpe_root"],
+                providers=["DmlExecutionProvider"],
+            )
+        else:
+
+            def rmvpe_jit_model():
+                ckpt = get_jit_model(model_path, is_half, self.device, rmvpe_jit_export)
+                model = torch.jit.load(BytesIO(ckpt["model"]), map_location=self.device)
+                model = model.to(self.device)
+                return model
+
+            if use_jit and not (is_half and "cpu" in str(self.device)):
+                self.model = rmvpe_jit_model()
+            else:
+                self.model = get_rmvpe(model_path, self.device, is_half)
+
+    def compute_f0(
+        self,
+        wav: np.ndarray,
+        p_len: Optional[int] = None,
+        filter_radius: Optional[Union[int, float]] = None,
+    ):
+        if p_len is None:
+            p_len = wav.shape[0] // self.hop_length
+        if not torch.is_tensor(wav):
+            wav = torch.from_numpy(wav)
+        mel = self.mel_extractor(wav.float().to(self.device).unsqueeze(0), center=True)
+        hidden = self._mel2hidden(mel)
+        if "privateuseone" not in str(self.device):
+            hidden = hidden.squeeze(0).cpu().numpy()
+        else:
+            hidden = hidden[0]
+        if self.is_half == True:
+            hidden = hidden.astype("float32")
+
+        f0 = self._decode(hidden, thred=filter_radius)
+
+        return self._interpolate_f0(self._resize_f0(f0, p_len))[0]
+
+    def _to_local_average_cents(self, salience, threshold=0.05):
+        center = np.argmax(salience, axis=1)  # 帧长#index
+        salience = np.pad(salience, ((0, 0), (4, 4)))  # 帧长,368
+        center += 4
+        todo_salience = []
+        todo_cents_mapping = []
+        starts = center - 4
+        ends = center + 5
+        for idx in range(salience.shape[0]):
+            todo_salience.append(salience[:, starts[idx] : ends[idx]][idx])
+            todo_cents_mapping.append(self.cents_mapping[starts[idx] : ends[idx]])
+        todo_salience = np.array(todo_salience)  # 帧长，9
+        todo_cents_mapping = np.array(todo_cents_mapping)  # 帧长，9
+        product_sum = np.sum(todo_salience * todo_cents_mapping, 1)
+        weight_sum = np.sum(todo_salience, 1)  # 帧长
+        devided = product_sum / weight_sum  # 帧长
+        maxx = np.max(salience, axis=1)  # 帧长
+        devided[maxx <= threshold] = 0
+        return devided
+
+    def _mel2hidden(self, mel):
+        with torch.no_grad():
+            n_frames = mel.shape[-1]
+            n_pad = 32 * ((n_frames - 1) // 32 + 1) - n_frames
+            if n_pad > 0:
+                mel = F.pad(mel, (0, n_pad), mode="constant")
+            if "privateuseone" in str(self.device):
+                onnx_input_name = self.model.get_inputs()[0].name
+                onnx_outputs_names = self.model.get_outputs()[0].name
+                hidden = self.model.run(
+                    [onnx_outputs_names],
+                    input_feed={onnx_input_name: mel.cpu().numpy()},
+                )[0]
+            else:
+                mel = mel.half() if self.is_half else mel.float()
+                hidden = self.model(mel)
+            return hidden[:, :n_frames]
+
+    def _decode(self, hidden, thred=0.03):
+        cents_pred = self._to_local_average_cents(hidden, threshold=thred)
+        f0 = 10 * (2 ** (cents_pred / 1200))
+        f0[f0 == 10] = 0
+        # f0 = np.array([10 * (2 ** (cent_pred / 1200)) if cent_pred else 0 for cent_pred in cents_pred])
+        return f0

rvc/f0/stft.py

@@ -0,0 +1,194 @@
+from typing import Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from librosa.util import pad_center
+from scipy.signal import get_window
+
+
+class STFT(torch.nn.Module):
+    def __init__(
+        self,
+        filter_length=1024,
+        hop_length=512,
+        win_length: Optional[int] = None,
+        window="hann",
+        use_torch_stft=True,
+    ):
+        """
+        This module implements an STFT using 1D convolution and 1D transpose convolutions.
+        This is a bit tricky so there are some cases that probably won't work as working
+        out the same sizes before and after in all overlap add setups is tough. Right now,
+        this code should work with hop lengths that are half the filter length (50% overlap
+        between frames).
+
+        Keyword Arguments:
+            filter_length {int} -- Length of filters used (default: {1024})
+            hop_length {int} -- Hop length of STFT (restrict to 50% overlap between frames) (default: {512})
+            win_length {[type]} -- Length of the window function applied to each frame (if not specified, it
+                equals the filter length). (default: {None})
+            window {str} -- Type of window to use (options are bartlett, hann, hamming, blackman, blackmanharris)
+                (default: {'hann'})
+        """
+        super(STFT, self).__init__()
+        self.filter_length = filter_length
+        self.hop_length = hop_length
+        self.pad_amount = int(self.filter_length / 2)
+        self.win_length = win_length
+        self.hann_window = {}
+        self.use_torch_stft = use_torch_stft
+
+        if use_torch_stft:
+            return
+
+        fourier_basis = np.fft.fft(np.eye(self.filter_length))
+
+        cutoff = int((self.filter_length / 2 + 1))
+        fourier_basis = np.vstack(
+            [np.real(fourier_basis[:cutoff, :]), np.imag(fourier_basis[:cutoff, :])]
+        )
+        forward_basis = torch.FloatTensor(fourier_basis)
+        inverse_basis = torch.FloatTensor(np.linalg.pinv(fourier_basis))
+
+        if win_length is None or not win_length:
+            win_length = filter_length
+        assert filter_length >= win_length
+
+        # get window and zero center pad it to filter_length
+        fft_window = get_window(window, win_length, fftbins=True)
+        fft_window = pad_center(fft_window, size=filter_length)
+        fft_window = torch.from_numpy(fft_window).float()
+
+        # window the bases
+        forward_basis *= fft_window
+        inverse_basis = (inverse_basis.T * fft_window).T
+
+        self.register_buffer("forward_basis", forward_basis.float())
+        self.register_buffer("inverse_basis", inverse_basis.float())
+        self.register_buffer("fft_window", fft_window.float())
+
+    def __call__(
+        self,
+        input_data: torch.Tensor,
+        keyshift: int = 0,
+        speed: int = 1,
+        center: bool = True,
+    ) -> torch.Tensor:
+        return super().__call__(input_data, keyshift, speed, center)
+
+    def transform(
+        self,
+        input_data: torch.Tensor,
+        return_phase=False,
+    ) -> Tuple[Union[Tuple[torch.Tensor, torch.Tensor], torch.Tensor]]:
+        """Take input data (audio) to STFT domain.
+
+        Arguments:
+            input_data {tensor} -- Tensor of floats, with shape (num_batch, num_samples)
+
+        Returns:
+            magnitude {tensor} -- Magnitude of STFT with shape (num_batch,
+                num_frequencies, num_frames)
+            phase {tensor} -- Phase of STFT with shape (num_batch,
+                num_frequencies, num_frames)
+        """
+        input_data = F.pad(
+            input_data,
+            (self.pad_amount, self.pad_amount),
+            mode="reflect",
+        )
+        forward_transform = input_data.unfold(
+            1, self.filter_length, self.hop_length
+        ).permute(0, 2, 1)
+        forward_transform = torch.matmul(self.forward_basis, forward_transform)
+        cutoff = int((self.filter_length / 2) + 1)
+        real_part = forward_transform[:, :cutoff, :]
+        imag_part = forward_transform[:, cutoff:, :]
+        magnitude = torch.sqrt(real_part**2 + imag_part**2)
+        if return_phase:
+            phase = torch.atan2(imag_part.data, real_part.data)
+            return magnitude, phase
+        else:
+            return magnitude
+
+    def inverse(
+        self,
+        magnitude: torch.Tensor,
+        phase: torch.Tensor,
+    ) -> torch.Tensor:
+        """Call the inverse STFT (iSTFT), given magnitude and phase tensors produced
+        by the ```transform``` function.
+
+        Arguments:
+            magnitude {tensor} -- Magnitude of STFT with shape (num_batch,
+                num_frequencies, num_frames)
+            phase {tensor} -- Phase of STFT with shape (num_batch,
+                num_frequencies, num_frames)
+
+        Returns:
+            inverse_transform {tensor} -- Reconstructed audio given magnitude and phase. Of
+                shape (num_batch, num_samples)
+        """
+        cat = torch.cat(
+            [magnitude * torch.cos(phase), magnitude * torch.sin(phase)], dim=1
+        )
+        fold = torch.nn.Fold(
+            output_size=(1, (cat.size(-1) - 1) * self.hop_length + self.filter_length),
+            kernel_size=(1, self.filter_length),
+            stride=(1, self.hop_length),
+        )
+        inverse_transform = torch.matmul(self.inverse_basis, cat)
+        inverse_transform: torch.Tensor = fold(inverse_transform)[
+            :, 0, 0, self.pad_amount : -self.pad_amount
+        ]
+        window_square_sum = (
+            self.fft_window.pow(2).repeat(cat.size(-1), 1).T.unsqueeze(0)
+        )
+        window_square_sum = fold(window_square_sum)[
+            :, 0, 0, self.pad_amount : -self.pad_amount
+        ]
+        inverse_transform /= window_square_sum
+        return inverse_transform
+
+    def forward(
+        self,
+        input_data: torch.Tensor,
+        keyshift: int = 0,
+        speed: int = 1,
+        center: bool = True,
+    ) -> torch.Tensor:
+        factor = 2 ** (keyshift / 12)
+        n_fft_new = int(np.round(self.filter_length * factor))
+        win_length_new = int(np.round(self.win_length * factor))
+        hop_length_new = int(np.round(self.hop_length * speed))
+        if self.use_torch_stft:
+            keyshift_key = str(keyshift) + "_" + str(input_data.device)
+            if keyshift_key not in self.hann_window:
+                self.hann_window[keyshift_key] = torch.hann_window(
+                    self.win_length,
+                ).to(input_data.device)
+            fft = torch.stft(
+                input_data,
+                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,
+            )
+            return torch.sqrt(fft.real.pow(2) + fft.imag.pow(2))
+        return self.transform(input_data)
+        """Take input data (audio) to STFT domain and then back to audio.
+
+        Arguments:
+            input_data {tensor} -- Tensor of floats, with shape (num_batch, num_samples)
+
+        Returns:
+            reconstruction {tensor} -- Reconstructed audio given magnitude and phase. Of
+                shape (num_batch, num_samples)
+        reconstruction = self.inverse(
+            self.transform(input_data, return_phase=True),
+        )
+        return reconstruction
+        """

rvc/hubert.py

@@ -0,0 +1,339 @@
+import math
+import random
+from typing import Optional, Tuple
+
+from fairseq.checkpoint_utils import load_model_ensemble_and_task
+from fairseq.utils import index_put
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+
+# @torch.jit.script
+def pad_to_multiple(x, multiple, dim=-1, value=0):
+    # Inspired from https://github.com/lucidrains/local-attention/blob/master/local_attention/local_attention.py#L41
+    if x is None:
+        return None, 0
+    tsz = x.size(dim)
+    m = tsz / multiple
+    remainder = math.ceil(m) * multiple - tsz
+    if int(tsz % multiple) == 0:
+        return x, 0
+    pad_offset = (0,) * (-1 - dim) * 2
+
+    return F.pad(x, (*pad_offset, 0, remainder), value=value), remainder
+
+
+def extract_features(
+    self,
+    x,
+    padding_mask=None,
+    tgt_layer=None,
+    min_layer=0,
+):
+    if padding_mask is not None:
+        x = index_put(x, padding_mask, 0)
+
+    x_conv = self.pos_conv(x.transpose(1, 2))
+    x_conv = x_conv.transpose(1, 2)
+    x = x + x_conv
+
+    if not self.layer_norm_first:
+        x = self.layer_norm(x)
+
+    # pad to the sequence length dimension
+    x, pad_length = pad_to_multiple(x, self.required_seq_len_multiple, dim=-2, value=0)
+    if pad_length > 0 and padding_mask is None:
+        padding_mask = x.new_zeros((x.size(0), x.size(1)), dtype=torch.bool)
+        padding_mask[:, -pad_length:] = True
+    else:
+        padding_mask, _ = pad_to_multiple(
+            padding_mask, self.required_seq_len_multiple, dim=-1, value=True
+        )
+    x = F.dropout(x, p=self.dropout, training=self.training)
+
+    # B x T x C -> T x B x C
+    x = x.transpose(0, 1)
+
+    layer_results = []
+    r = None
+    for i, layer in enumerate(self.layers):
+        dropout_probability = np.random.random() if self.layerdrop > 0 else 1
+        if not self.training or (dropout_probability > self.layerdrop):
+            x, (z, lr) = layer(
+                x, self_attn_padding_mask=padding_mask, need_weights=False
+            )
+            if i >= min_layer:
+                layer_results.append((x, z, lr))
+        if i == tgt_layer:
+            r = x
+            break
+
+    if r is not None:
+        x = r
+
+    # T x B x C -> B x T x C
+    x = x.transpose(0, 1)
+
+    # undo paddding
+    if pad_length > 0:
+        x = x[:, :-pad_length]
+
+        def undo_pad(a, b, c):
+            return (
+                a[:-pad_length],
+                b[:-pad_length] if b is not None else b,
+                c[:-pad_length],
+            )
+
+        layer_results = [undo_pad(*u) for u in layer_results]
+
+    return x, layer_results
+
+
+def compute_mask_indices(
+    shape: Tuple[int, int],
+    padding_mask: Optional[torch.Tensor],
+    mask_prob: float,
+    mask_length: int,
+    mask_type: str = "static",
+    mask_other: float = 0.0,
+    min_masks: int = 0,
+    no_overlap: bool = False,
+    min_space: int = 0,
+    require_same_masks: bool = True,
+    mask_dropout: float = 0.0,
+) -> torch.Tensor:
+    """
+    Computes random mask spans for a given shape
+
+    Args:
+        shape: the the shape for which to compute masks.
+            should be of size 2 where first element is batch size and 2nd is timesteps
+        padding_mask: optional padding mask of the same size as shape, which will prevent masking padded elements
+        mask_prob: probability for each token to be chosen as start of the span to be masked. this will be multiplied by
+            number of timesteps divided by length of mask span to mask approximately this percentage of all elements.
+            however due to overlaps, the actual number will be smaller (unless no_overlap is True)
+        mask_type: how to compute mask lengths
+            static = fixed size
+            uniform = sample from uniform distribution [mask_other, mask_length*2]
+            normal = sample from normal distribution with mean mask_length and stdev mask_other. mask is min 1 element
+            poisson = sample from possion distribution with lambda = mask length
+        min_masks: minimum number of masked spans
+        no_overlap: if false, will switch to an alternative recursive algorithm that prevents spans from overlapping
+        min_space: only used if no_overlap is True, this is how many elements to keep unmasked between spans
+        require_same_masks: if true, will randomly drop out masks until same amount of masks remains in each sample
+        mask_dropout: randomly dropout this percentage of masks in each example
+    """
+
+    bsz, all_sz = shape
+    mask = torch.full((bsz, all_sz), False)
+
+    all_num_mask = int(
+        # add a random number for probabilistic rounding
+        mask_prob * all_sz / float(mask_length)
+        + torch.rand([1]).item()
+    )
+
+    all_num_mask = max(min_masks, all_num_mask)
+
+    mask_idcs = []
+    for i in range(bsz):
+        if padding_mask is not None:
+            sz = all_sz - padding_mask[i].long().sum().item()
+            num_mask = int(mask_prob * sz / float(mask_length) + np.random.rand())
+            num_mask = max(min_masks, num_mask)
+        else:
+            sz = all_sz
+            num_mask = all_num_mask
+
+        if mask_type == "static":
+            lengths = torch.full([num_mask], mask_length)
+        elif mask_type == "uniform":
+            lengths = torch.randint(mask_other, mask_length * 2 + 1, size=[num_mask])
+        elif mask_type == "normal":
+            lengths = torch.normal(mask_length, mask_other, size=[num_mask])
+            lengths = [max(1, int(round(x))) for x in lengths]
+        else:
+            raise Exception("unknown mask selection " + mask_type)
+
+        if sum(lengths) == 0:
+            lengths[0] = min(mask_length, sz - 1)
+
+        if no_overlap:
+            mask_idc = []
+
+            def arrange(s, e, length, keep_length):
+                span_start = torch.randint(low=s, high=e - length, size=[1]).item()
+                mask_idc.extend(span_start + i for i in range(length))
+
+                new_parts = []
+                if span_start - s - min_space >= keep_length:
+                    new_parts.append((s, span_start - min_space + 1))
+                if e - span_start - length - min_space > keep_length:
+                    new_parts.append((span_start + length + min_space, e))
+                return new_parts
+
+            parts = [(0, sz)]
+            min_length = min(lengths)
+            for length in sorted(lengths, reverse=True):
+                t = [e - s if e - s >= length + min_space else 0 for s, e in parts]
+                lens = torch.asarray(t, dtype=torch.int)
+                l_sum = torch.sum(lens)
+                if l_sum == 0:
+                    break
+                probs = lens / torch.sum(lens)
+                c = torch.multinomial(probs.float(), len(parts)).item()
+                s, e = parts.pop(c)
+                parts.extend(arrange(s, e, length, min_length))
+            mask_idc = torch.asarray(mask_idc)
+        else:
+            min_len = min(lengths)
+            if sz - min_len <= num_mask:
+                min_len = sz - num_mask - 1
+            mask_idc = torch.asarray(
+                random.sample([i for i in range(sz - min_len)], num_mask)
+            )
+            mask_idc = torch.asarray(
+                [
+                    mask_idc[j] + offset
+                    for j in range(len(mask_idc))
+                    for offset in range(lengths[j])
+                ]
+            )
+
+        mask_idcs.append(torch.unique(mask_idc[mask_idc < sz]))
+
+    min_len = min([len(m) for m in mask_idcs])
+    for i, mask_idc in enumerate(mask_idcs):
+        if isinstance(mask_idc, torch.Tensor):
+            mask_idc = torch.asarray(mask_idc, dtype=torch.float)
+        if len(mask_idc) > min_len and require_same_masks:
+            mask_idc = torch.asarray(
+                random.sample([i for i in range(mask_idc)], min_len)
+            )
+        if mask_dropout > 0:
+            num_holes = int(round(len(mask_idc) * mask_dropout))
+            mask_idc = torch.asarray(
+                random.sample([i for i in range(mask_idc)], len(mask_idc) - num_holes)
+            )
+
+        mask[i, mask_idc.int()] = True
+
+    return mask
+
+
+def apply_mask(self, x, padding_mask, target_list):
+    B, T, C = x.shape
+    torch.zeros_like(x)
+    if self.mask_prob > 0:
+        mask_indices = compute_mask_indices(
+            (B, T),
+            padding_mask,
+            self.mask_prob,
+            self.mask_length,
+            self.mask_selection,
+            self.mask_other,
+            min_masks=2,
+            no_overlap=self.no_mask_overlap,
+            min_space=self.mask_min_space,
+        )
+        mask_indices = mask_indices.to(x.device)
+        x[mask_indices] = self.mask_emb
+    else:
+        mask_indices = None
+
+    if self.mask_channel_prob > 0:
+        mask_channel_indices = compute_mask_indices(
+            (B, C),
+            None,
+            self.mask_channel_prob,
+            self.mask_channel_length,
+            self.mask_channel_selection,
+            self.mask_channel_other,
+            no_overlap=self.no_mask_channel_overlap,
+            min_space=self.mask_channel_min_space,
+        )
+        mask_channel_indices = (
+            mask_channel_indices.to(x.device).unsqueeze(1).expand(-1, T, -1)
+        )
+        x[mask_channel_indices] = 0
+
+    return x, mask_indices
+
+
+def get_hubert(model_path="assets/hubert/hubert_base.pt", device=torch.device("cpu")):
+    models, _, _ = load_model_ensemble_and_task(
+        [model_path],
+        suffix="",
+    )
+    hubert_model = models[0]
+    hubert_model = hubert_model.to(device)
+
+    def _apply_mask(x, padding_mask, target_list):
+        return apply_mask(hubert_model, x, padding_mask, target_list)
+
+    hubert_model.apply_mask = _apply_mask
+
+    def _extract_features(
+        x,
+        padding_mask=None,
+        tgt_layer=None,
+        min_layer=0,
+    ):
+        return extract_features(
+            hubert_model.encoder,
+            x,
+            padding_mask=padding_mask,
+            tgt_layer=tgt_layer,
+            min_layer=min_layer,
+        )
+
+    hubert_model.encoder.extract_features = _extract_features
+
+    hubert_model._forward = hubert_model.forward
+
+    def hubert_extract_features(
+        self,
+        source: torch.Tensor,
+        padding_mask: Optional[torch.Tensor] = None,
+        mask: bool = False,
+        ret_conv: bool = False,
+        output_layer: Optional[int] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        res = self._forward(
+            source,
+            padding_mask=padding_mask,
+            mask=mask,
+            features_only=True,
+            output_layer=output_layer,
+        )
+        feature = res["features"] if ret_conv else res["x"]
+        return feature, res["padding_mask"]
+
+    def _hubert_extract_features(
+        source: torch.Tensor,
+        padding_mask: Optional[torch.Tensor] = None,
+        mask: bool = False,
+        ret_conv: bool = False,
+        output_layer: Optional[int] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        return hubert_extract_features(
+            hubert_model, source, padding_mask, mask, ret_conv, output_layer
+        )
+
+    hubert_model.extract_features = _hubert_extract_features
+
+    def infer(source, padding_mask, output_layer: torch.Tensor):
+        output_layer = output_layer.item()
+        logits = hubert_model.extract_features(
+            source=source, padding_mask=padding_mask, output_layer=output_layer
+        )
+        feats = hubert_model.final_proj(logits[0]) if output_layer == 9 else logits[0]
+        return feats
+
+    hubert_model.infer = infer
+    # hubert_model.forward=infer
+    # hubert_model.forward
+
+    return hubert_model

rvc/ipex/__init__.py

@@ -0,0 +1,10 @@
+try:
+    import torch
+
+    if torch.xpu.is_available():
+        from .init import ipex_init
+
+        ipex_init()
+        from .gradscaler import gradscaler_init
+except Exception:  # pylint: disable=broad-exception-caught
+    pass

rvc/ipex/attention.py

@@ -0,0 +1,218 @@
+import torch
+import intel_extension_for_pytorch as ipex  # pylint: disable=import-error, unused-import
+
+# pylint: disable=protected-access, missing-function-docstring, line-too-long
+
+original_torch_bmm = torch.bmm
+
+
+def torch_bmm(input, mat2, *, out=None):
+    if input.dtype != mat2.dtype:
+        mat2 = mat2.to(input.dtype)
+
+    # ARC GPUs can't allocate more than 4GB to a single block, Slice it:
+    batch_size_attention, input_tokens, mat2_shape = (
+        input.shape[0],
+        input.shape[1],
+        mat2.shape[2],
+    )
+    block_multiply = input.element_size()
+    slice_block_size = input_tokens * mat2_shape / 1024 / 1024 * block_multiply
+    block_size = batch_size_attention * slice_block_size
+
+    split_slice_size = batch_size_attention
+    if block_size > 4:
+        do_split = True
+        # Find something divisible with the input_tokens
+        while (split_slice_size * slice_block_size) > 4:
+            split_slice_size = split_slice_size // 2
+            if split_slice_size <= 1:
+                split_slice_size = 1
+                break
+    else:
+        do_split = False
+
+    split_2_slice_size = input_tokens
+    if split_slice_size * slice_block_size > 4:
+        slice_block_size2 = split_slice_size * mat2_shape / 1024 / 1024 * block_multiply
+        do_split_2 = True
+        # Find something divisible with the input_tokens
+        while (split_2_slice_size * slice_block_size2) > 4:
+            split_2_slice_size = split_2_slice_size // 2
+            if split_2_slice_size <= 1:
+                split_2_slice_size = 1
+                break
+    else:
+        do_split_2 = False
+
+    if do_split:
+        hidden_states = torch.zeros(
+            input.shape[0],
+            input.shape[1],
+            mat2.shape[2],
+            device=input.device,
+            dtype=input.dtype,
+        )
+        for i in range(batch_size_attention // split_slice_size):
+            start_idx = i * split_slice_size
+            end_idx = (i + 1) * split_slice_size
+            if do_split_2:
+                for i2 in range(
+                    input_tokens // split_2_slice_size
+                ):  # pylint: disable=invalid-name
+                    start_idx_2 = i2 * split_2_slice_size
+                    end_idx_2 = (i2 + 1) * split_2_slice_size
+                    hidden_states[start_idx:end_idx, start_idx_2:end_idx_2] = (
+                        original_torch_bmm(
+                            input[start_idx:end_idx, start_idx_2:end_idx_2],
+                            mat2[start_idx:end_idx, start_idx_2:end_idx_2],
+                            out=out,
+                        )
+                    )
+            else:
+                hidden_states[start_idx:end_idx] = original_torch_bmm(
+                    input[start_idx:end_idx], mat2[start_idx:end_idx], out=out
+                )
+    else:
+        return original_torch_bmm(input, mat2, out=out)
+    return hidden_states
+
+
+original_scaled_dot_product_attention = torch.nn.functional.scaled_dot_product_attention
+
+
+def scaled_dot_product_attention(
+    query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False
+):
+    # ARC GPUs can't allocate more than 4GB to a single block, Slice it:
+    if len(query.shape) == 3:
+        batch_size_attention, query_tokens, shape_four = query.shape
+        shape_one = 1
+        no_shape_one = True
+    else:
+        shape_one, batch_size_attention, query_tokens, shape_four = query.shape
+        no_shape_one = False
+
+    block_multiply = query.element_size()
+    slice_block_size = (
+        shape_one * query_tokens * shape_four / 1024 / 1024 * block_multiply
+    )
+    block_size = batch_size_attention * slice_block_size
+
+    split_slice_size = batch_size_attention
+    if block_size > 4:
+        do_split = True
+        # Find something divisible with the shape_one
+        while (split_slice_size * slice_block_size) > 4:
+            split_slice_size = split_slice_size // 2
+            if split_slice_size <= 1:
+                split_slice_size = 1
+                break
+    else:
+        do_split = False
+
+    split_2_slice_size = query_tokens
+    if split_slice_size * slice_block_size > 4:
+        slice_block_size2 = (
+            shape_one * split_slice_size * shape_four / 1024 / 1024 * block_multiply
+        )
+        do_split_2 = True
+        # Find something divisible with the batch_size_attention
+        while (split_2_slice_size * slice_block_size2) > 4:
+            split_2_slice_size = split_2_slice_size // 2
+            if split_2_slice_size <= 1:
+                split_2_slice_size = 1
+                break
+    else:
+        do_split_2 = False
+
+    if do_split:
+        hidden_states = torch.zeros(query.shape, device=query.device, dtype=query.dtype)
+        for i in range(batch_size_attention // split_slice_size):
+            start_idx = i * split_slice_size
+            end_idx = (i + 1) * split_slice_size
+            if do_split_2:
+                for i2 in range(
+                    query_tokens // split_2_slice_size
+                ):  # pylint: disable=invalid-name
+                    start_idx_2 = i2 * split_2_slice_size
+                    end_idx_2 = (i2 + 1) * split_2_slice_size
+                    if no_shape_one:
+                        hidden_states[start_idx:end_idx, start_idx_2:end_idx_2] = (
+                            original_scaled_dot_product_attention(
+                                query[start_idx:end_idx, start_idx_2:end_idx_2],
+                                key[start_idx:end_idx, start_idx_2:end_idx_2],
+                                value[start_idx:end_idx, start_idx_2:end_idx_2],
+                                attn_mask=(
+                                    attn_mask[start_idx:end_idx, start_idx_2:end_idx_2]
+                                    if attn_mask is not None
+                                    else attn_mask
+                                ),
+                                dropout_p=dropout_p,
+                                is_causal=is_causal,
+                            )
+                        )
+                    else:
+                        hidden_states[:, start_idx:end_idx, start_idx_2:end_idx_2] = (
+                            original_scaled_dot_product_attention(
+                                query[:, start_idx:end_idx, start_idx_2:end_idx_2],
+                                key[:, start_idx:end_idx, start_idx_2:end_idx_2],
+                                value[:, start_idx:end_idx, start_idx_2:end_idx_2],
+                                attn_mask=(
+                                    attn_mask[
+                                        :, start_idx:end_idx, start_idx_2:end_idx_2
+                                    ]
+                                    if attn_mask is not None
+                                    else attn_mask
+                                ),
+                                dropout_p=dropout_p,
+                                is_causal=is_causal,
+                            )
+                        )
+            else:
+                if no_shape_one:
+                    hidden_states[start_idx:end_idx] = (
+                        original_scaled_dot_product_attention(
+                            query[start_idx:end_idx],
+                            key[start_idx:end_idx],
+                            value[start_idx:end_idx],
+                            attn_mask=(
+                                attn_mask[start_idx:end_idx]
+                                if attn_mask is not None
+                                else attn_mask
+                            ),
+                            dropout_p=dropout_p,
+                            is_causal=is_causal,
+                        )
+                    )
+                else:
+                    hidden_states[:, start_idx:end_idx] = (
+                        original_scaled_dot_product_attention(
+                            query[:, start_idx:end_idx],
+                            key[:, start_idx:end_idx],
+                            value[:, start_idx:end_idx],
+                            attn_mask=(
+                                attn_mask[:, start_idx:end_idx]
+                                if attn_mask is not None
+                                else attn_mask
+                            ),
+                            dropout_p=dropout_p,
+                            is_causal=is_causal,
+                        )
+                    )
+    else:
+        return original_scaled_dot_product_attention(
+            query,
+            key,
+            value,
+            attn_mask=attn_mask,
+            dropout_p=dropout_p,
+            is_causal=is_causal,
+        )
+    return hidden_states
+
+
+def attention_init():
+    # ARC GPUs can't allocate more than 4GB to a single block:
+    torch.bmm = torch_bmm
+    torch.nn.functional.scaled_dot_product_attention = scaled_dot_product_attention

rvc/ipex/gradscaler.py

@@ -0,0 +1,188 @@
+from collections import defaultdict
+
+import torch
+import intel_extension_for_pytorch as ipex  # pylint: disable=import-error, unused-import
+import intel_extension_for_pytorch._C as core  # pylint: disable=import-error, unused-import
+
+# pylint: disable=protected-access, missing-function-docstring, line-too-long
+
+OptState = ipex.cpu.autocast._grad_scaler.OptState
+_MultiDeviceReplicator = ipex.cpu.autocast._grad_scaler._MultiDeviceReplicator
+_refresh_per_optimizer_state = (
+    ipex.cpu.autocast._grad_scaler._refresh_per_optimizer_state
+)
+
+
+def _unscale_grads_(
+    self, optimizer, inv_scale, found_inf, allow_fp16
+):  # pylint: disable=unused-argument
+    per_device_inv_scale = _MultiDeviceReplicator(inv_scale)
+    per_device_found_inf = _MultiDeviceReplicator(found_inf)
+
+    # To set up _amp_foreach_non_finite_check_and_unscale_, split grads by device and dtype.
+    # There could be hundreds of grads, so we'd like to iterate through them just once.
+    # However, we don't know their devices or dtypes in advance.
+
+    # https://stackoverflow.com/questions/5029934/defaultdict-of-defaultdict
+    # Google says mypy struggles with defaultdicts type annotations.
+    per_device_and_dtype_grads = defaultdict(lambda: defaultdict(list))  # type: ignore[var-annotated]
+    # sync grad to master weight
+    if hasattr(optimizer, "sync_grad"):
+        optimizer.sync_grad()
+    with torch.no_grad():
+        for group in optimizer.param_groups:
+            for param in group["params"]:
+                if param.grad is None:
+                    continue
+                if (not allow_fp16) and param.grad.dtype == torch.float16:
+                    raise ValueError("Attempting to unscale FP16 gradients.")
+                if param.grad.is_sparse:
+                    # is_coalesced() == False means the sparse grad has values with duplicate indices.
+                    # coalesce() deduplicates indices and adds all values that have the same index.
+                    # For scaled fp16 values, there's a good chance coalescing will cause overflow,
+                    # so we should check the coalesced _values().
+                    if param.grad.dtype is torch.float16:
+                        param.grad = param.grad.coalesce()
+                    to_unscale = param.grad._values()
+                else:
+                    to_unscale = param.grad
+
+                # -: is there a way to split by device and dtype without appending in the inner loop?
+                to_unscale = to_unscale.to("cpu")
+                per_device_and_dtype_grads[to_unscale.device][to_unscale.dtype].append(
+                    to_unscale
+                )
+
+        for _, per_dtype_grads in per_device_and_dtype_grads.items():
+            for grads in per_dtype_grads.values():
+                core._amp_foreach_non_finite_check_and_unscale_(
+                    grads,
+                    per_device_found_inf.get("cpu"),
+                    per_device_inv_scale.get("cpu"),
+                )
+
+    return per_device_found_inf._per_device_tensors
+
+
+def unscale_(self, optimizer):
+    """
+    Divides ("unscales") the optimizer's gradient tensors by the scale factor.
+    :meth:`unscale_` is optional, serving cases where you need to
+    :ref:`modify or inspect gradients<working-with-unscaled-gradients>`
+    between the backward pass(es) and :meth:`step`.
+    If :meth:`unscale_` is not called explicitly,  gradients will be unscaled  automatically during :meth:`step`.
+    Simple example, using :meth:`unscale_` to enable clipping of unscaled gradients::
+        ...
+        scaler.scale(loss).backward()
+        scaler.unscale_(optimizer)
+        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
+        scaler.step(optimizer)
+        scaler.update()
+    Args:
+        optimizer (torch.optim.Optimizer):  Optimizer that owns the gradients to be unscaled.
+    .. warning::
+        :meth:`unscale_` should only be called once per optimizer per :meth:`step` call,
+        and only after all gradients for that optimizer's assigned parameters have been accumulated.
+        Calling :meth:`unscale_` twice for a given optimizer between each :meth:`step` triggers a RuntimeError.
+    .. warning::
+        :meth:`unscale_` may unscale sparse gradients out of place, replacing the ``.grad`` attribute.
+    """
+    if not self._enabled:
+        return
+
+    self._check_scale_growth_tracker("unscale_")
+
+    optimizer_state = self._per_optimizer_states[id(optimizer)]
+
+    if optimizer_state["stage"] is OptState.UNSCALED:  # pylint: disable=no-else-raise
+        raise RuntimeError(
+            "unscale_() has already been called on this optimizer since the last update()."
+        )
+    elif optimizer_state["stage"] is OptState.STEPPED:
+        raise RuntimeError("unscale_() is being called after step().")
+
+    # FP32 division can be imprecise for certain compile options, so we carry out the reciprocal in FP64.
+    assert self._scale is not None
+    inv_scale = (
+        self._scale.to("cpu").double().reciprocal().float().to(self._scale.device)
+    )
+    found_inf = torch.full((1,), 0.0, dtype=torch.float32, device=self._scale.device)
+
+    optimizer_state["found_inf_per_device"] = self._unscale_grads_(
+        optimizer, inv_scale, found_inf, False
+    )
+    optimizer_state["stage"] = OptState.UNSCALED
+
+
+def update(self, new_scale=None):
+    """
+    Updates the scale factor.
+    If any optimizer steps were skipped the scale is multiplied by ``backoff_factor``
+    to reduce it. If ``growth_interval`` unskipped iterations occurred consecutively,
+    the scale is multiplied by ``growth_factor`` to increase it.
+    Passing ``new_scale`` sets the new scale value manually. (``new_scale`` is not
+    used directly, it's used to fill GradScaler's internal scale tensor. So if
+    ``new_scale`` was a tensor, later in-place changes to that tensor will not further
+    affect the scale GradScaler uses internally.)
+    Args:
+        new_scale (float or :class:`torch.FloatTensor`, optional, default=None):  New scale factor.
+    .. warning::
+        :meth:`update` should only be called at the end of the iteration, after ``scaler.step(optimizer)`` has
+        been invoked for all optimizers used this iteration.
+    """
+    if not self._enabled:
+        return
+
+    _scale, _growth_tracker = self._check_scale_growth_tracker("update")
+
+    if new_scale is not None:
+        # Accept a new user-defined scale.
+        if isinstance(new_scale, float):
+            self._scale.fill_(new_scale)  # type: ignore[union-attr]
+        else:
+            reason = "new_scale should be a float or a 1-element torch.FloatTensor with requires_grad=False."
+            assert isinstance(new_scale, torch.FloatTensor), reason  # type: ignore[attr-defined]
+            assert new_scale.numel() == 1, reason
+            assert new_scale.requires_grad is False, reason
+            self._scale.copy_(new_scale)  # type: ignore[union-attr]
+    else:
+        # Consume shared inf/nan data collected from optimizers to update the scale.
+        # If all found_inf tensors are on the same device as self._scale, this operation is asynchronous.
+        found_infs = [
+            found_inf.to(device="cpu", non_blocking=True)
+            for state in self._per_optimizer_states.values()
+            for found_inf in state["found_inf_per_device"].values()
+        ]
+
+        assert len(found_infs) > 0, "No inf checks were recorded prior to update."
+
+        found_inf_combined = found_infs[0]
+        if len(found_infs) > 1:
+            for i in range(1, len(found_infs)):
+                found_inf_combined += found_infs[i]
+
+        to_device = _scale.device
+        _scale = _scale.to("cpu")
+        _growth_tracker = _growth_tracker.to("cpu")
+
+        core._amp_update_scale_(
+            _scale,
+            _growth_tracker,
+            found_inf_combined,
+            self._growth_factor,
+            self._backoff_factor,
+            self._growth_interval,
+        )
+
+        _scale = _scale.to(to_device)
+        _growth_tracker = _growth_tracker.to(to_device)
+    # To prepare for next iteration, clear the data collected from optimizers this iteration.
+    self._per_optimizer_states = defaultdict(_refresh_per_optimizer_state)
+
+
+def gradscaler_init():
+    torch.xpu.amp.GradScaler = ipex.cpu.autocast._grad_scaler.GradScaler
+    torch.xpu.amp.GradScaler._unscale_grads_ = _unscale_grads_
+    torch.xpu.amp.GradScaler.unscale_ = unscale_
+    torch.xpu.amp.GradScaler.update = update
+    return torch.xpu.amp.GradScaler

rvc/ipex/hijacks.py

@@ -0,0 +1,366 @@
+import contextlib
+import importlib
+
+import torch
+import intel_extension_for_pytorch as ipex  # pylint: disable=import-error, unused-import
+
+# pylint: disable=protected-access, missing-function-docstring, line-too-long, unnecessary-lambda, no-else-return
+
+
+class CondFunc:  # pylint: disable=missing-class-docstring
+    def __new__(cls, orig_func, sub_func, cond_func):
+        self = super(CondFunc, cls).__new__(cls)
+        if isinstance(orig_func, str):
+            func_path = orig_func.split(".")
+            for i in range(len(func_path) - 1, -1, -1):
+                try:
+                    resolved_obj = importlib.import_module(".".join(func_path[:i]))
+                    break
+                except ImportError:
+                    pass
+            for attr_name in func_path[i:-1]:
+                resolved_obj = getattr(resolved_obj, attr_name)
+            orig_func = getattr(resolved_obj, func_path[-1])
+            setattr(
+                resolved_obj,
+                func_path[-1],
+                lambda *args, **kwargs: self(*args, **kwargs),
+            )
+        self.__init__(orig_func, sub_func, cond_func)
+        return lambda *args, **kwargs: self(*args, **kwargs)
+
+    def __init__(self, orig_func, sub_func, cond_func):
+        self.__orig_func = orig_func
+        self.__sub_func = sub_func
+        self.__cond_func = cond_func
+
+    def __call__(self, *args, **kwargs):
+        if not self.__cond_func or self.__cond_func(self.__orig_func, *args, **kwargs):
+            return self.__sub_func(self.__orig_func, *args, **kwargs)
+        else:
+            return self.__orig_func(*args, **kwargs)
+
+
+_utils = torch.utils.data._utils
+
+
+def _shutdown_workers(self):
+    if (
+        torch.utils.data._utils is None
+        or torch.utils.data._utils.python_exit_status is True
+        or torch.utils.data._utils.python_exit_status is None
+    ):
+        return
+    if hasattr(self, "_shutdown") and not self._shutdown:
+        self._shutdown = True
+        try:
+            if hasattr(self, "_pin_memory_thread"):
+                self._pin_memory_thread_done_event.set()
+                self._worker_result_queue.put((None, None))
+                self._pin_memory_thread.join()
+                self._worker_result_queue.cancel_join_thread()
+                self._worker_result_queue.close()
+            self._workers_done_event.set()
+            for worker_id in range(len(self._workers)):
+                if self._persistent_workers or self._workers_status[worker_id]:
+                    self._mark_worker_as_unavailable(worker_id, shutdown=True)
+            for w in self._workers:  # pylint: disable=invalid-name
+                w.join(timeout=torch.utils.data._utils.MP_STATUS_CHECK_INTERVAL)
+            for q in self._index_queues:  # pylint: disable=invalid-name
+                q.cancel_join_thread()
+                q.close()
+        finally:
+            if self._worker_pids_set:
+                torch.utils.data._utils.signal_handling._remove_worker_pids(id(self))
+                self._worker_pids_set = False
+            for w in self._workers:  # pylint: disable=invalid-name
+                if w.is_alive():
+                    w.terminate()
+
+
+class DummyDataParallel(
+    torch.nn.Module
+):  # pylint: disable=missing-class-docstring, unused-argument, too-few-public-methods
+    def __new__(
+        cls, module, device_ids=None, output_device=None, dim=0
+    ):  # pylint: disable=unused-argument
+        if isinstance(device_ids, list) and len(device_ids) > 1:
+            print("IPEX backend doesn't support DataParallel on multiple XPU devices")
+        return module.to("xpu")
+
+
+def return_null_context(*args, **kwargs):  # pylint: disable=unused-argument
+    return contextlib.nullcontext()
+
+
+def check_device(device):
+    return bool(
+        (isinstance(device, torch.device) and device.type == "cuda")
+        or (isinstance(device, str) and "cuda" in device)
+        or isinstance(device, int)
+    )
+
+
+def return_xpu(device):
+    return (
+        f"xpu:{device[-1]}"
+        if isinstance(device, str) and ":" in device
+        else (
+            f"xpu:{device}"
+            if isinstance(device, int)
+            else torch.device("xpu") if isinstance(device, torch.device) else "xpu"
+        )
+    )
+
+
+def ipex_no_cuda(orig_func, *args, **kwargs):
+    torch.cuda.is_available = lambda: False
+    orig_func(*args, **kwargs)
+    torch.cuda.is_available = torch.xpu.is_available
+
+
+original_autocast = torch.autocast
+
+
+def ipex_autocast(*args, **kwargs):
+    if len(args) > 0 and args[0] == "cuda":
+        return original_autocast("xpu", *args[1:], **kwargs)
+    else:
+        return original_autocast(*args, **kwargs)
+
+
+original_torch_cat = torch.cat
+
+
+def torch_cat(tensor, *args, **kwargs):
+    if len(tensor) == 3 and (
+        tensor[0].dtype != tensor[1].dtype or tensor[2].dtype != tensor[1].dtype
+    ):
+        return original_torch_cat(
+            [tensor[0].to(tensor[1].dtype), tensor[1], tensor[2].to(tensor[1].dtype)],
+            *args,
+            **kwargs,
+        )
+    else:
+        return original_torch_cat(tensor, *args, **kwargs)
+
+
+original_interpolate = torch.nn.functional.interpolate
+
+
+def interpolate(
+    tensor,
+    size=None,
+    scale_factor=None,
+    mode="nearest",
+    align_corners=None,
+    recompute_scale_factor=None,
+    antialias=False,
+):  # pylint: disable=too-many-arguments
+    if antialias or align_corners is not None:
+        return_device = tensor.device
+        return_dtype = tensor.dtype
+        return original_interpolate(
+            tensor.to("cpu", dtype=torch.float32),
+            size=size,
+            scale_factor=scale_factor,
+            mode=mode,
+            align_corners=align_corners,
+            recompute_scale_factor=recompute_scale_factor,
+            antialias=antialias,
+        ).to(return_device, dtype=return_dtype)
+    else:
+        return original_interpolate(
+            tensor,
+            size=size,
+            scale_factor=scale_factor,
+            mode=mode,
+            align_corners=align_corners,
+            recompute_scale_factor=recompute_scale_factor,
+            antialias=antialias,
+        )
+
+
+original_linalg_solve = torch.linalg.solve
+
+
+def linalg_solve(A, B, *args, **kwargs):  # pylint: disable=invalid-name
+    if A.device != torch.device("cpu") or B.device != torch.device("cpu"):
+        return_device = A.device
+        return original_linalg_solve(A.to("cpu"), B.to("cpu"), *args, **kwargs).to(
+            return_device
+        )
+    else:
+        return original_linalg_solve(A, B, *args, **kwargs)
+
+
+def ipex_hijacks():
+    CondFunc(
+        "torch.Tensor.to",
+        lambda orig_func, self, device=None, *args, **kwargs: orig_func(
+            self, return_xpu(device), *args, **kwargs
+        ),
+        lambda orig_func, self, device=None, *args, **kwargs: check_device(device),
+    )
+    CondFunc(
+        "torch.Tensor.cuda",
+        lambda orig_func, self, device=None, *args, **kwargs: orig_func(
+            self, return_xpu(device), *args, **kwargs
+        ),
+        lambda orig_func, self, device=None, *args, **kwargs: check_device(device),
+    )
+    CondFunc(
+        "torch.empty",
+        lambda orig_func, *args, device=None, **kwargs: orig_func(
+            *args, device=return_xpu(device), **kwargs
+        ),
+        lambda orig_func, *args, device=None, **kwargs: check_device(device),
+    )
+    CondFunc(
+        "torch.load",
+        lambda orig_func, *args, map_location=None, **kwargs: orig_func(
+            *args, return_xpu(map_location), **kwargs
+        ),
+        lambda orig_func, *args, map_location=None, **kwargs: map_location is None
+        or check_device(map_location),
+    )
+    CondFunc(
+        "torch.randn",
+        lambda orig_func, *args, device=None, **kwargs: orig_func(
+            *args, device=return_xpu(device), **kwargs
+        ),
+        lambda orig_func, *args, device=None, **kwargs: check_device(device),
+    )
+    CondFunc(
+        "torch.ones",
+        lambda orig_func, *args, device=None, **kwargs: orig_func(
+            *args, device=return_xpu(device), **kwargs
+        ),
+        lambda orig_func, *args, device=None, **kwargs: check_device(device),
+    )
+    CondFunc(
+        "torch.zeros",
+        lambda orig_func, *args, device=None, **kwargs: orig_func(
+            *args, device=return_xpu(device), **kwargs
+        ),
+        lambda orig_func, *args, device=None, **kwargs: check_device(device),
+    )
+    CondFunc(
+        "torch.tensor",
+        lambda orig_func, *args, device=None, **kwargs: orig_func(
+            *args, device=return_xpu(device), **kwargs
+        ),
+        lambda orig_func, *args, device=None, **kwargs: check_device(device),
+    )
+    CondFunc(
+        "torch.linspace",
+        lambda orig_func, *args, device=None, **kwargs: orig_func(
+            *args, device=return_xpu(device), **kwargs
+        ),
+        lambda orig_func, *args, device=None, **kwargs: check_device(device),
+    )
+
+    CondFunc(
+        "torch.Generator",
+        lambda orig_func, device=None: torch.xpu.Generator(device),
+        lambda orig_func, device=None: device is not None
+        and device != torch.device("cpu")
+        and device != "cpu",
+    )
+
+    CondFunc(
+        "torch.batch_norm",
+        lambda orig_func, input, weight, bias, *args, **kwargs: orig_func(
+            input,
+            (
+                weight
+                if weight is not None
+                else torch.ones(input.size()[1], device=input.device)
+            ),
+            (
+                bias
+                if bias is not None
+                else torch.zeros(input.size()[1], device=input.device)
+            ),
+            *args,
+            **kwargs,
+        ),
+        lambda orig_func, input, *args, **kwargs: input.device != torch.device("cpu"),
+    )
+    CondFunc(
+        "torch.instance_norm",
+        lambda orig_func, input, weight, bias, *args, **kwargs: orig_func(
+            input,
+            (
+                weight
+                if weight is not None
+                else torch.ones(input.size()[1], device=input.device)
+            ),
+            (
+                bias
+                if bias is not None
+                else torch.zeros(input.size()[1], device=input.device)
+            ),
+            *args,
+            **kwargs,
+        ),
+        lambda orig_func, input, *args, **kwargs: input.device != torch.device("cpu"),
+    )
+
+    # Functions with dtype errors:
+    CondFunc(
+        "torch.nn.modules.GroupNorm.forward",
+        lambda orig_func, self, input: orig_func(
+            self, input.to(self.weight.data.dtype)
+        ),
+        lambda orig_func, self, input: input.dtype != self.weight.data.dtype,
+    )
+    CondFunc(
+        "torch.nn.modules.linear.Linear.forward",
+        lambda orig_func, self, input: orig_func(
+            self, input.to(self.weight.data.dtype)
+        ),
+        lambda orig_func, self, input: input.dtype != self.weight.data.dtype,
+    )
+    CondFunc(
+        "torch.nn.modules.conv.Conv2d.forward",
+        lambda orig_func, self, input: orig_func(
+            self, input.to(self.weight.data.dtype)
+        ),
+        lambda orig_func, self, input: input.dtype != self.weight.data.dtype,
+    )
+    CondFunc(
+        "torch.nn.functional.layer_norm",
+        lambda orig_func, input, normalized_shape=None, weight=None, *args, **kwargs: orig_func(
+            input.to(weight.data.dtype), normalized_shape, weight, *args, **kwargs
+        ),
+        lambda orig_func, input, normalized_shape=None, weight=None, *args, **kwargs: weight
+        is not None
+        and input.dtype != weight.data.dtype,
+    )
+
+    # Diffusers Float64 (ARC GPUs doesn't support double or Float64):
+    if not torch.xpu.has_fp64_dtype():
+        CondFunc(
+            "torch.from_numpy",
+            lambda orig_func, ndarray: orig_func(ndarray.astype("float32")),
+            lambda orig_func, ndarray: ndarray.dtype == float,
+        )
+
+    # Broken functions when torch.cuda.is_available is True:
+    CondFunc(
+        "torch.utils.data.dataloader._BaseDataLoaderIter.__init__",
+        lambda orig_func, *args, **kwargs: ipex_no_cuda(orig_func, *args, **kwargs),
+        lambda orig_func, *args, **kwargs: True,
+    )
+
+    # Functions that make compile mad with CondFunc:
+    torch.utils.data.dataloader._MultiProcessingDataLoaderIter._shutdown_workers = (
+        _shutdown_workers
+    )
+    torch.nn.DataParallel = DummyDataParallel
+    torch.autocast = ipex_autocast
+    torch.cat = torch_cat
+    torch.linalg.solve = linalg_solve
+    torch.nn.functional.interpolate = interpolate
+    torch.backends.cuda.sdp_kernel = return_null_context

rvc/ipex/init.py

@@ -0,0 +1,192 @@
+import os
+import sys
+import contextlib
+
+import torch
+import intel_extension_for_pytorch as ipex  # pylint: disable=import-error, unused-import
+
+from .hijacks import ipex_hijacks
+from .attention import attention_init
+
+# pylint: disable=protected-access, missing-function-docstring, line-too-long
+
+
+def ipex_init():  # pylint: disable=too-many-statements
+    try:
+        # Replace cuda with xpu:
+        torch.cuda.current_device = torch.xpu.current_device
+        torch.cuda.current_stream = torch.xpu.current_stream
+        torch.cuda.device = torch.xpu.device
+        torch.cuda.device_count = torch.xpu.device_count
+        torch.cuda.device_of = torch.xpu.device_of
+        torch.cuda.get_device_name = torch.xpu.get_device_name
+        torch.cuda.get_device_properties = torch.xpu.get_device_properties
+        torch.cuda.init = torch.xpu.init
+        torch.cuda.is_available = torch.xpu.is_available
+        torch.cuda.is_initialized = torch.xpu.is_initialized
+        torch.cuda.is_current_stream_capturing = lambda: False
+        torch.cuda.set_device = torch.xpu.set_device
+        torch.cuda.stream = torch.xpu.stream
+        torch.cuda.synchronize = torch.xpu.synchronize
+        torch.cuda.Event = torch.xpu.Event
+        torch.cuda.Stream = torch.xpu.Stream
+        torch.cuda.FloatTensor = torch.xpu.FloatTensor
+        torch.Tensor.cuda = torch.Tensor.xpu
+        torch.Tensor.is_cuda = torch.Tensor.is_xpu
+        torch.cuda._initialization_lock = torch.xpu.lazy_init._initialization_lock
+        torch.cuda._initialized = torch.xpu.lazy_init._initialized
+        torch.cuda._lazy_seed_tracker = torch.xpu.lazy_init._lazy_seed_tracker
+        torch.cuda._queued_calls = torch.xpu.lazy_init._queued_calls
+        torch.cuda._tls = torch.xpu.lazy_init._tls
+        torch.cuda.threading = torch.xpu.lazy_init.threading
+        torch.cuda.traceback = torch.xpu.lazy_init.traceback
+        torch.cuda.Optional = torch.xpu.Optional
+        torch.cuda.__cached__ = torch.xpu.__cached__
+        torch.cuda.__loader__ = torch.xpu.__loader__
+        torch.cuda.ComplexFloatStorage = torch.xpu.ComplexFloatStorage
+        torch.cuda.Tuple = torch.xpu.Tuple
+        torch.cuda.streams = torch.xpu.streams
+        torch.cuda._lazy_new = torch.xpu._lazy_new
+        torch.cuda.FloatStorage = torch.xpu.FloatStorage
+        torch.cuda.Any = torch.xpu.Any
+        torch.cuda.__doc__ = torch.xpu.__doc__
+        torch.cuda.default_generators = torch.xpu.default_generators
+        torch.cuda.HalfTensor = torch.xpu.HalfTensor
+        torch.cuda._get_device_index = torch.xpu._get_device_index
+        torch.cuda.__path__ = torch.xpu.__path__
+        torch.cuda.Device = torch.xpu.Device
+        torch.cuda.IntTensor = torch.xpu.IntTensor
+        torch.cuda.ByteStorage = torch.xpu.ByteStorage
+        torch.cuda.set_stream = torch.xpu.set_stream
+        torch.cuda.BoolStorage = torch.xpu.BoolStorage
+        torch.cuda.os = torch.xpu.os
+        torch.cuda.torch = torch.xpu.torch
+        torch.cuda.BFloat16Storage = torch.xpu.BFloat16Storage
+        torch.cuda.Union = torch.xpu.Union
+        torch.cuda.DoubleTensor = torch.xpu.DoubleTensor
+        torch.cuda.ShortTensor = torch.xpu.ShortTensor
+        torch.cuda.LongTensor = torch.xpu.LongTensor
+        torch.cuda.IntStorage = torch.xpu.IntStorage
+        torch.cuda.LongStorage = torch.xpu.LongStorage
+        torch.cuda.__annotations__ = torch.xpu.__annotations__
+        torch.cuda.__package__ = torch.xpu.__package__
+        torch.cuda.__builtins__ = torch.xpu.__builtins__
+        torch.cuda.CharTensor = torch.xpu.CharTensor
+        torch.cuda.List = torch.xpu.List
+        torch.cuda._lazy_init = torch.xpu._lazy_init
+        torch.cuda.BFloat16Tensor = torch.xpu.BFloat16Tensor
+        torch.cuda.DoubleStorage = torch.xpu.DoubleStorage
+        torch.cuda.ByteTensor = torch.xpu.ByteTensor
+        torch.cuda.StreamContext = torch.xpu.StreamContext
+        torch.cuda.ComplexDoubleStorage = torch.xpu.ComplexDoubleStorage
+        torch.cuda.ShortStorage = torch.xpu.ShortStorage
+        torch.cuda._lazy_call = torch.xpu._lazy_call
+        torch.cuda.HalfStorage = torch.xpu.HalfStorage
+        torch.cuda.random = torch.xpu.random
+        torch.cuda._device = torch.xpu._device
+        torch.cuda.classproperty = torch.xpu.classproperty
+        torch.cuda.__name__ = torch.xpu.__name__
+        torch.cuda._device_t = torch.xpu._device_t
+        torch.cuda.warnings = torch.xpu.warnings
+        torch.cuda.__spec__ = torch.xpu.__spec__
+        torch.cuda.BoolTensor = torch.xpu.BoolTensor
+        torch.cuda.CharStorage = torch.xpu.CharStorage
+        torch.cuda.__file__ = torch.xpu.__file__
+        torch.cuda._is_in_bad_fork = torch.xpu.lazy_init._is_in_bad_fork
+        # torch.cuda.is_current_stream_capturing = torch.xpu.is_current_stream_capturing
+
+        # Memory:
+        torch.cuda.memory = torch.xpu.memory
+        if "linux" in sys.platform and "WSL2" in os.popen("uname -a").read():
+            torch.xpu.empty_cache = lambda: None
+        torch.cuda.empty_cache = torch.xpu.empty_cache
+        torch.cuda.memory_stats = torch.xpu.memory_stats
+        torch.cuda.memory_summary = torch.xpu.memory_summary
+        torch.cuda.memory_snapshot = torch.xpu.memory_snapshot
+        torch.cuda.memory_allocated = torch.xpu.memory_allocated
+        torch.cuda.max_memory_allocated = torch.xpu.max_memory_allocated
+        torch.cuda.memory_reserved = torch.xpu.memory_reserved
+        torch.cuda.memory_cached = torch.xpu.memory_reserved
+        torch.cuda.max_memory_reserved = torch.xpu.max_memory_reserved
+        torch.cuda.max_memory_cached = torch.xpu.max_memory_reserved
+        torch.cuda.reset_peak_memory_stats = torch.xpu.reset_peak_memory_stats
+        torch.cuda.reset_max_memory_cached = torch.xpu.reset_peak_memory_stats
+        torch.cuda.reset_max_memory_allocated = torch.xpu.reset_peak_memory_stats
+        torch.cuda.memory_stats_as_nested_dict = torch.xpu.memory_stats_as_nested_dict
+        torch.cuda.reset_accumulated_memory_stats = (
+            torch.xpu.reset_accumulated_memory_stats
+        )
+
+        # RNG:
+        torch.cuda.get_rng_state = torch.xpu.get_rng_state
+        torch.cuda.get_rng_state_all = torch.xpu.get_rng_state_all
+        torch.cuda.set_rng_state = torch.xpu.set_rng_state
+        torch.cuda.set_rng_state_all = torch.xpu.set_rng_state_all
+        torch.cuda.manual_seed = torch.xpu.manual_seed
+        torch.cuda.manual_seed_all = torch.xpu.manual_seed_all
+        torch.cuda.seed = torch.xpu.seed
+        torch.cuda.seed_all = torch.xpu.seed_all
+        torch.cuda.initial_seed = torch.xpu.initial_seed
+
+        # AMP:
+        torch.cuda.amp = torch.xpu.amp
+        if not hasattr(torch.cuda.amp, "common"):
+            torch.cuda.amp.common = contextlib.nullcontext()
+        torch.cuda.amp.common.amp_definitely_not_available = lambda: False
+        try:
+            torch.cuda.amp.GradScaler = torch.xpu.amp.GradScaler
+        except Exception:  # pylint: disable=broad-exception-caught
+            try:
+                from .gradscaler import (
+                    gradscaler_init,
+                )  # pylint: disable=import-outside-toplevel, import-error
+
+                gradscaler_init()
+                torch.cuda.amp.GradScaler = torch.xpu.amp.GradScaler
+            except Exception:  # pylint: disable=broad-exception-caught
+                torch.cuda.amp.GradScaler = ipex.cpu.autocast._grad_scaler.GradScaler
+
+        # C
+        torch._C._cuda_getCurrentRawStream = ipex._C._getCurrentStream
+        ipex._C._DeviceProperties.major = 2023
+        ipex._C._DeviceProperties.minor = 2
+
+        # Fix functions with ipex:
+        torch.cuda.mem_get_info = lambda device=None: [
+            (
+                torch.xpu.get_device_properties(device).total_memory
+                - torch.xpu.memory_allocated(device)
+            ),
+            torch.xpu.get_device_properties(device).total_memory,
+        ]
+        torch._utils._get_available_device_type = lambda: "xpu"
+        torch.has_cuda = True
+        torch.cuda.has_half = True
+        torch.cuda.is_bf16_supported = lambda *args, **kwargs: True
+        torch.cuda.is_fp16_supported = lambda *args, **kwargs: True
+        torch.version.cuda = "11.7"
+        torch.cuda.get_device_capability = lambda *args, **kwargs: [11, 7]
+        torch.cuda.get_device_properties.major = 11
+        torch.cuda.get_device_properties.minor = 7
+        torch.cuda.ipc_collect = lambda *args, **kwargs: None
+        torch.cuda.utilization = lambda *args, **kwargs: 0
+        if hasattr(torch.xpu, "getDeviceIdListForCard"):
+            torch.cuda.getDeviceIdListForCard = torch.xpu.getDeviceIdListForCard
+            torch.cuda.get_device_id_list_per_card = torch.xpu.getDeviceIdListForCard
+        else:
+            torch.cuda.getDeviceIdListForCard = torch.xpu.get_device_id_list_per_card
+            torch.cuda.get_device_id_list_per_card = (
+                torch.xpu.get_device_id_list_per_card
+            )
+
+        ipex_hijacks()
+        attention_init()
+        try:
+            from .diffusers import ipex_diffusers
+
+            ipex_diffusers()
+        except Exception:  # pylint: disable=broad-exception-caught
+            pass
+    except Exception as e:
+        return False, e
+    return True, None

rvc/jit/__init__.py

@@ -0,0 +1 @@
+from .jit import load_inputs, get_jit_model, export_jit_model, save_pickle

rvc/jit/jit.py

@@ -0,0 +1,78 @@
+import pickle
+from io import BytesIO
+from collections import OrderedDict
+import os
+
+import torch
+
+
+def load_pickle(path: str):
+    with open(path, "rb") as f:
+        return pickle.load(f)
+
+
+def save_pickle(ckpt: dict, save_path: str):
+    with open(save_path, "wb") as f:
+        pickle.dump(ckpt, f)
+
+
+def load_inputs(path: torch.serialization.FILE_LIKE, device: str, is_half=False):
+    parm = torch.load(path, map_location=torch.device("cpu"))
+    for key in parm.keys():
+        parm[key] = parm[key].to(device)
+        if is_half and parm[key].dtype == torch.float32:
+            parm[key] = parm[key].half()
+        elif not is_half and parm[key].dtype == torch.float16:
+            parm[key] = parm[key].float()
+    return parm
+
+
+def export_jit_model(
+    model: torch.nn.Module,
+    mode: str = "trace",
+    inputs: dict = None,
+    device=torch.device("cpu"),
+    is_half: bool = False,
+) -> dict:
+    model = model.half() if is_half else model.float()
+    model.eval()
+    if mode == "trace":
+        assert inputs is not None
+        model_jit = torch.jit.trace(model, example_kwarg_inputs=inputs)
+    elif mode == "script":
+        model_jit = torch.jit.script(model)
+    model_jit.to(device)
+    model_jit = model_jit.half() if is_half else model_jit.float()
+    buffer = BytesIO()
+    # model_jit=model_jit.cpu()
+    torch.jit.save(model_jit, buffer)
+    del model_jit
+    cpt = OrderedDict()
+    cpt["model"] = buffer.getvalue()
+    cpt["is_half"] = is_half
+    return cpt
+
+
+def get_jit_model(model_path: str, is_half: bool, device: str, exporter):
+    jit_model_path = model_path.rstrip(".pth")
+    jit_model_path += ".half.jit" if is_half else ".jit"
+    ckpt = None
+
+    if os.path.exists(jit_model_path):
+        ckpt = load_pickle(jit_model_path)
+        model_device = ckpt["device"]
+        if model_device != str(device):
+            del ckpt
+            ckpt = None
+
+    if ckpt is None:
+        ckpt = exporter(
+            model_path=model_path,
+            mode="script",
+            inputs_path=None,
+            save_path=jit_model_path,
+            device=device,
+            is_half=is_half,
+        )
+
+    return ckpt

rvc/layers/attentions.py

@@ -0,0 +1,292 @@
+import math
+from typing import Optional
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        out_channels: int,
+        n_heads: int,
+        p_dropout: float = 0.0,
+        window_size: Optional[int] = None,
+        heads_share: bool = True,
+        block_length: Optional[int] = None,
+        proximal_bias: bool = False,
+        proximal_init: bool = False,
+    ):
+        super(MultiHeadAttention, self).__init__()
+        assert channels % n_heads == 0
+
+        self.channels = channels
+        self.out_channels = out_channels
+        self.n_heads = n_heads
+        self.p_dropout = p_dropout
+        self.window_size = window_size
+        self.heads_share = heads_share
+        self.block_length = block_length
+        self.proximal_bias = proximal_bias
+        self.proximal_init = proximal_init
+        self.attn = None
+
+        self.k_channels = channels // n_heads
+        self.conv_q = nn.Conv1d(channels, channels, 1)
+        self.conv_k = nn.Conv1d(channels, channels, 1)
+        self.conv_v = nn.Conv1d(channels, channels, 1)
+        self.conv_o = nn.Conv1d(channels, out_channels, 1)
+        self.drop = nn.Dropout(p_dropout)
+
+        if window_size is not None:
+            n_heads_rel = 1 if heads_share else n_heads
+            rel_stddev = self.k_channels**-0.5
+            self.emb_rel_k = nn.Parameter(
+                torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels)
+                * rel_stddev
+            )
+            self.emb_rel_v = nn.Parameter(
+                torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels)
+                * rel_stddev
+            )
+
+        nn.init.xavier_uniform_(self.conv_q.weight)
+        nn.init.xavier_uniform_(self.conv_k.weight)
+        nn.init.xavier_uniform_(self.conv_v.weight)
+        if proximal_init:
+            with torch.no_grad():
+                self.conv_k.weight.copy_(self.conv_q.weight)
+                self.conv_k.bias.copy_(self.conv_q.bias)
+
+    def __call__(
+        self,
+        x: torch.Tensor,
+        c: torch.Tensor,
+        attn_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        return super().__call__(x, c, attn_mask=attn_mask)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        c: torch.Tensor,
+        attn_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        q = self.conv_q(x)
+        k = self.conv_k(c)
+        v = self.conv_v(c)
+
+        x, _ = self._attention(q, k, v, mask=attn_mask)
+
+        x = self.conv_o(x)
+        return x
+
+    def _attention(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+    ):
+        # reshape [b, d, t] -> [b, n_h, t, d_k]
+        b, d, t_s = key.size()
+        t_t = query.size(2)
+        query = query.view(b, self.n_heads, self.k_channels, t_t).transpose(2, 3)
+        key = key.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)
+        value = value.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)
+
+        scores = torch.matmul(query / math.sqrt(self.k_channels), key.transpose(-2, -1))
+        if self.window_size is not None:
+            assert (
+                t_s == t_t
+            ), "Relative attention is only available for self-attention."
+            key_relative_embeddings = self._get_relative_embeddings(self.emb_rel_k, t_s)
+            rel_logits = self._matmul_with_relative_keys(
+                query / math.sqrt(self.k_channels), key_relative_embeddings
+            )
+            scores_local = self._relative_position_to_absolute_position(rel_logits)
+            scores = scores + scores_local
+        if self.proximal_bias:
+            assert t_s == t_t, "Proximal bias is only available for self-attention."
+            scores = scores + self._attention_bias_proximal(t_s).to(
+                device=scores.device, dtype=scores.dtype
+            )
+        if mask is not None:
+            scores = scores.masked_fill(mask == 0, -1e4)
+            if self.block_length is not None:
+                assert (
+                    t_s == t_t
+                ), "Local attention is only available for self-attention."
+                block_mask = (
+                    torch.ones_like(scores)
+                    .triu(-self.block_length)
+                    .tril(self.block_length)
+                )
+                scores = scores.masked_fill(block_mask == 0, -1e4)
+        p_attn = F.softmax(scores, dim=-1)  # [b, n_h, t_t, t_s]
+        p_attn = self.drop(p_attn)
+        output = torch.matmul(p_attn, value)
+        if self.window_size is not None:
+            relative_weights = self._absolute_position_to_relative_position(p_attn)
+            value_relative_embeddings = self._get_relative_embeddings(
+                self.emb_rel_v, t_s
+            )
+            output = output + self._matmul_with_relative_values(
+                relative_weights, value_relative_embeddings
+            )
+        output = (
+            output.transpose(2, 3).contiguous().view(b, d, t_t)
+        )  # [b, n_h, t_t, d_k] -> [b, d, t_t]
+        return output, p_attn
+
+    def _matmul_with_relative_values(self, x, y):
+        """
+        x: [b, h, l, m]
+        y: [h or 1, m, d]
+        ret: [b, h, l, d]
+        """
+        ret = torch.matmul(x, y.unsqueeze(0))
+        return ret
+
+    def _matmul_with_relative_keys(self, x, y):
+        """
+        x: [b, h, l, d]
+        y: [h or 1, m, d]
+        ret: [b, h, l, m]
+        """
+        ret = torch.matmul(x, y.unsqueeze(0).transpose(-2, -1))
+        return ret
+
+    def _get_relative_embeddings(self, relative_embeddings, length: int):
+        # max_relative_position = 2 * self.window_size + 1
+        # Pad first before slice to avoid using cond ops.
+        pad_length: int = max(length - (self.window_size + 1), 0)
+        slice_start_position = max((self.window_size + 1) - length, 0)
+        slice_end_position = slice_start_position + 2 * length - 1
+        if pad_length > 0:
+            padded_relative_embeddings = F.pad(
+                relative_embeddings,
+                [0, 0, pad_length, pad_length, 0, 0],
+            )
+        else:
+            padded_relative_embeddings = relative_embeddings
+        used_relative_embeddings = padded_relative_embeddings[
+            :, slice_start_position:slice_end_position
+        ]
+        return used_relative_embeddings
+
+    def _relative_position_to_absolute_position(self, x):
+        """
+        x: [b, h, l, 2*l-1]
+        ret: [b, h, l, l]
+        """
+        batch, heads, length, _ = x.size()
+        # Concat columns of pad to shift from relative to absolute indexing.
+        x = F.pad(
+            x,
+            [0, 1, 0, 0, 0, 0, 0, 0],
+        )
+
+        # Concat extra elements so to add up to shape (len+1, 2*len-1).
+        x_flat = x.view([batch, heads, length * 2 * length])
+        x_flat = F.pad(x_flat, [0, length - 1, 0, 0, 0, 0])
+
+        # Reshape and slice out the padded elements.
+        x_final = x_flat.view([batch, heads, length + 1, 2 * length - 1])[
+            :, :, :length, length - 1 :
+        ]
+        return x_final
+
+    def _absolute_position_to_relative_position(self, x):
+        """
+        x: [b, h, l, l]
+        ret: [b, h, l, 2*l-1]
+        """
+        batch, heads, length, _ = x.size()
+        # padd along column
+        x = F.pad(x, [0, length - 1, 0, 0, 0, 0, 0, 0])
+        x_flat = x.view([batch, heads, (length**2) + (length * (length - 1))])
+        # add 0's in the beginning that will skew the elements after reshape
+        x_flat = F.pad(x_flat, [length, 0, 0, 0, 0, 0])
+        x_final = x_flat.view([batch, heads, length, 2 * length])[:, :, :, 1:]
+        return x_final
+
+    def _attention_bias_proximal(self, length: int):
+        """Bias for self-attention to encourage attention to close positions.
+        Args:
+          length: an integer scalar.
+        Returns:
+          a Tensor with shape [1, 1, length, length]
+        """
+        r = torch.arange(length, dtype=torch.float32)
+        diff = torch.unsqueeze(r, 0) - torch.unsqueeze(r, 1)
+        return torch.unsqueeze(torch.unsqueeze(-torch.log1p(torch.abs(diff)), 0), 0)
+
+
+class FFN(nn.Module):
+    """
+    Feed-Forward Network
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        filter_channels: int,
+        kernel_size: int,
+        p_dropout: float = 0.0,
+        activation: Optional[str] = None,
+        causal: bool = False,
+    ):
+        super(FFN, self).__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.filter_channels = filter_channels
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.activation = activation
+        self.causal = causal
+        self.is_activation = True if activation == "gelu" else False
+
+        self.conv_1 = nn.Conv1d(in_channels, filter_channels, kernel_size)
+        self.conv_2 = nn.Conv1d(filter_channels, out_channels, kernel_size)
+        self.drop = nn.Dropout(p_dropout)
+
+    def __call__(self, x: torch.Tensor, x_mask: torch.Tensor) -> torch.Tensor:
+        return super().__call__(x, x_mask)
+
+    def forward(self, x: torch.Tensor, x_mask: torch.Tensor) -> torch.Tensor:
+        x = self.conv_1(self._padding(x, x_mask))
+        if self.is_activation:
+            x = x * torch.sigmoid(1.702 * x)
+        else:
+            x = torch.relu(x)
+        x = self.drop(x)
+
+        x = self.conv_2(self._padding(x, x_mask))
+        return x * x_mask
+
+    def _padding(self, x: torch.Tensor, x_mask: torch.Tensor) -> torch.Tensor:
+        if self.causal:
+            return self._causal_padding(x * x_mask)
+        return self._same_padding(x * x_mask)
+
+    def _causal_padding(self, x):
+        if self.kernel_size == 1:
+            return x
+        pad_l: int = self.kernel_size - 1
+        pad_r: int = 0
+        # padding = [[0, 0], [0, 0], [pad_l, pad_r]]
+        x = F.pad(x, [pad_l, pad_r, 0, 0, 0, 0])
+        return x
+
+    def _same_padding(self, x):
+        if self.kernel_size == 1:
+            return x
+        pad_l: int = (self.kernel_size - 1) // 2
+        pad_r: int = self.kernel_size // 2
+        # padding = [[0, 0], [0, 0], [pad_l, pad_r]]
+        x = F.pad(x, [pad_l, pad_r, 0, 0, 0, 0])
+        return x

rvc/layers/discriminators.py

@@ -0,0 +1,172 @@
+from typing import List, Tuple
+
+import torch
+from torch import nn
+from torch.nn import Conv1d, Conv2d
+from torch.nn import functional as F
+from torch.nn.utils import spectral_norm, weight_norm
+
+from .residuals import LRELU_SLOPE
+from .utils import get_padding
+
+
+class MultiPeriodDiscriminator(torch.nn.Module):
+    """
+    version: 'v1' or 'v2'
+    """
+
+    def __init__(
+        self, version: str, use_spectral_norm: bool = False, has_xpu: bool = False
+    ):
+        super(MultiPeriodDiscriminator, self).__init__()
+        periods = (
+            (2, 3, 5, 7, 11, 17) if version == "v1" else (2, 3, 5, 7, 11, 17, 23, 37)
+        )
+
+        self.discriminators = nn.ModuleList(
+            [
+                DiscriminatorS(use_spectral_norm=use_spectral_norm),
+                *(
+                    DiscriminatorP(
+                        i, use_spectral_norm=use_spectral_norm, has_xpu=has_xpu
+                    )
+                    for i in periods
+                ),
+            ]
+        )
+
+    def __call__(self, y: torch.Tensor, y_hat: torch.Tensor) -> Tuple[
+        List[torch.Tensor],
+        List[torch.Tensor],
+        List[List[torch.Tensor]],
+        List[List[torch.Tensor]],
+    ]:
+        return super().__call__(y, y_hat)
+
+    def forward(self, y: torch.Tensor, y_hat: torch.Tensor) -> Tuple[
+        List[torch.Tensor],
+        List[torch.Tensor],
+        List[List[torch.Tensor]],
+        List[List[torch.Tensor]],
+    ]:
+        y_d_rs = []
+        y_d_gs = []
+        fmap_rs = []
+        fmap_gs = []
+
+        for d in self.discriminators:
+            y_d_r, fmap_r = d(y)
+            y_d_g, fmap_g = d(y_hat)
+            y_d_rs.append(y_d_r)
+            y_d_gs.append(y_d_g)
+            fmap_rs.append(fmap_r)
+            fmap_gs.append(fmap_g)
+
+        return y_d_rs, y_d_gs, fmap_rs, fmap_gs
+
+
+class DiscriminatorS(torch.nn.Module):
+    def __init__(self, use_spectral_norm: bool = False):
+        super(DiscriminatorS, self).__init__()
+        norm_f = spectral_norm if use_spectral_norm else weight_norm
+
+        self.convs = nn.ModuleList(
+            [
+                norm_f(Conv1d(1, 16, 15, 1, padding=7)),
+                norm_f(Conv1d(16, 64, 41, 4, groups=4, padding=20)),
+                norm_f(Conv1d(64, 256, 41, 4, groups=16, padding=20)),
+                norm_f(Conv1d(256, 1024, 41, 4, groups=64, padding=20)),
+                norm_f(Conv1d(1024, 1024, 41, 4, groups=256, padding=20)),
+                norm_f(Conv1d(1024, 1024, 5, 1, padding=2)),
+            ]
+        )
+        self.conv_post = norm_f(Conv1d(1024, 1, 3, 1, padding=1))
+
+    def __call__(self, x: torch.Tensor) -> Tuple[torch.Tensor, List[torch.Tensor]]:
+        return super().__call__(x)
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, List[torch.Tensor]]:
+        fmap = []
+
+        for l in self.convs:
+            x = l(x)
+            x = F.leaky_relu(x, LRELU_SLOPE)
+            fmap.append(x)
+
+        x = self.conv_post(x)
+        fmap.append(x)
+        x = torch.flatten(x, 1, -1)
+
+        return x, fmap
+
+
+class DiscriminatorP(torch.nn.Module):
+    def __init__(
+        self,
+        period: int,
+        kernel_size: int = 5,
+        stride: int = 3,
+        use_spectral_norm: bool = False,
+        has_xpu: bool = False,
+    ):
+        super(DiscriminatorP, self).__init__()
+        self.period = period
+        self.has_xpu = has_xpu
+        norm_f = spectral_norm if use_spectral_norm else weight_norm
+        sequence = (1, 32, 128, 512, 1024)
+        convs_padding = (get_padding(kernel_size, 1), 0)
+
+        self.convs = nn.ModuleList()
+        for i in range(len(sequence) - 1):
+            self.convs.append(
+                norm_f(
+                    Conv2d(
+                        sequence[i],
+                        sequence[i + 1],
+                        (kernel_size, 1),
+                        (stride, 1),
+                        padding=convs_padding,
+                    )
+                )
+            )
+        self.convs.append(
+            norm_f(
+                Conv2d(
+                    1024,
+                    1024,
+                    (kernel_size, 1),
+                    1,
+                    padding=convs_padding,
+                )
+            )
+        )
+        self.conv_post = norm_f(Conv2d(1024, 1, (3, 1), 1, padding=(1, 0)))
+
+    def __call__(self, x: torch.Tensor) -> Tuple[torch.Tensor, List[torch.Tensor]]:
+        return super().__call__(x)
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, List[torch.Tensor]]:
+        fmap = []
+
+        # 1d to 2d
+        b, c, t = x.shape
+        if t % self.period != 0:  # pad first
+            n_pad = self.period - (t % self.period)
+            if self.has_xpu and x.dtype == torch.bfloat16:
+                x = F.pad(x.to(dtype=torch.float16), (0, n_pad), "reflect").to(
+                    dtype=torch.bfloat16
+                )
+            else:
+                x = F.pad(x, (0, n_pad), "reflect")
+            t = t + n_pad
+        x = x.view(b, c, t // self.period, self.period)
+
+        for l in self.convs:
+            x = l(x)
+            x = F.leaky_relu(x, LRELU_SLOPE)
+            fmap.append(x)
+        x = self.conv_post(x)
+        fmap.append(x)
+        x = torch.flatten(x, 1, -1)
+
+        return x, fmap

rvc/layers/encoders.py

@@ -0,0 +1,221 @@
+import math
+from typing import Tuple, Optional
+
+import torch
+from torch import nn
+
+from .attentions import MultiHeadAttention, FFN
+from .norms import LayerNorm, WN
+from .utils import sequence_mask
+
+
+class Encoder(nn.Module):
+    def __init__(
+        self,
+        hidden_channels: int,
+        filter_channels: int,
+        n_heads: int,
+        n_layers: int,
+        kernel_size: int = 1,
+        p_dropout: float = 0.0,
+        window_size: int = 10,
+    ):
+        super(Encoder, self).__init__()
+
+        self.hidden_channels = hidden_channels
+        self.filter_channels = filter_channels
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.window_size = window_size
+
+        self.drop = nn.Dropout(p_dropout)
+        self.attn_layers = nn.ModuleList()
+        self.norm_layers_1 = nn.ModuleList()
+        self.ffn_layers = nn.ModuleList()
+        self.norm_layers_2 = nn.ModuleList()
+
+        for _ in range(self.n_layers):
+            self.attn_layers.append(
+                MultiHeadAttention(
+                    hidden_channels,
+                    hidden_channels,
+                    n_heads,
+                    p_dropout=p_dropout,
+                    window_size=window_size,
+                )
+            )
+            self.norm_layers_1.append(LayerNorm(hidden_channels))
+            self.ffn_layers.append(
+                FFN(
+                    hidden_channels,
+                    hidden_channels,
+                    filter_channels,
+                    kernel_size,
+                    p_dropout=p_dropout,
+                )
+            )
+            self.norm_layers_2.append(LayerNorm(hidden_channels))
+
+    def __call__(self, x: torch.Tensor, x_mask: torch.Tensor) -> torch.Tensor:
+        return super().__call__(x, x_mask)
+
+    def forward(self, x: torch.Tensor, x_mask: torch.Tensor) -> torch.Tensor:
+        attn_mask = x_mask.unsqueeze(2) * x_mask.unsqueeze(-1)
+        x = x * x_mask
+        for attn, norm1, ffn, norm2 in zip(
+            self.attn_layers,
+            self.norm_layers_1,
+            self.ffn_layers,
+            self.norm_layers_2,
+        ):
+            y = attn(x, x, attn_mask)
+            y = self.drop(y)
+            x = norm1(x + y)
+
+            y = ffn(x, x_mask)
+            y = self.drop(y)
+            x = norm2(x + y)
+        x = x * x_mask
+        return x
+
+
+class TextEncoder(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        hidden_channels: int,
+        filter_channels: int,
+        n_heads: int,
+        n_layers: int,
+        kernel_size: int,
+        p_dropout: float,
+        f0: bool = True,
+    ):
+        super(TextEncoder, self).__init__()
+
+        self.out_channels = out_channels
+        self.hidden_channels = hidden_channels
+        self.filter_channels = filter_channels
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.kernel_size = kernel_size
+        self.p_dropout = float(p_dropout)
+
+        self.emb_phone = nn.Linear(in_channels, hidden_channels)
+        self.lrelu = nn.LeakyReLU(0.1, inplace=True)
+        if f0 == True:
+            self.emb_pitch = nn.Embedding(256, hidden_channels)  # pitch 256
+        self.encoder = Encoder(
+            hidden_channels,
+            filter_channels,
+            n_heads,
+            n_layers,
+            kernel_size,
+            float(p_dropout),
+        )
+        self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
+
+    def __call__(
+        self,
+        phone: torch.Tensor,
+        pitch: torch.Tensor,
+        lengths: torch.Tensor,
+        skip_head: Optional[int] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        return super().__call__(
+            phone,
+            pitch,
+            lengths,
+            skip_head=skip_head,
+        )
+
+    def forward(
+        self,
+        phone: torch.Tensor,
+        pitch: torch.Tensor,
+        lengths: torch.Tensor,
+        skip_head: Optional[int] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        x = self.emb_phone(phone)
+        if pitch is not None:
+            x += self.emb_pitch(pitch)
+        x = x * math.sqrt(self.hidden_channels)  # [b, t, h]
+        x = self.lrelu(x)
+        x = torch.transpose(x, 1, -1)  # [b, h, t]
+        x_mask = torch.unsqueeze(
+            sequence_mask(lengths, x.size(2)),
+            1,
+        ).to(x.dtype)
+        x = self.encoder(x * x_mask, x_mask)
+        if skip_head is not None:
+            head = int(skip_head)
+            x = x[:, :, head:]
+            x_mask = x_mask[:, :, head:]
+        stats: torch.Tensor = self.proj(x) * x_mask
+        m, logs = torch.split(stats, self.out_channels, dim=1)
+        return m, logs, x_mask
+
+
+class PosteriorEncoder(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        hidden_channels: int,
+        kernel_size: int,
+        dilation_rate: int,
+        n_layers: int,
+        gin_channels=0,
+    ):
+        super(PosteriorEncoder, self).__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.dilation_rate = dilation_rate
+        self.n_layers = n_layers
+        self.gin_channels = gin_channels
+
+        self.pre = nn.Conv1d(in_channels, hidden_channels, 1)
+        self.enc = WN(
+            hidden_channels,
+            kernel_size,
+            dilation_rate,
+            n_layers,
+            gin_channels=gin_channels,
+        )
+        self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
+
+    def __call__(
+        self, x: torch.Tensor, x_lengths: torch.Tensor, g: Optional[torch.Tensor] = None
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        return super().__call__(x, x_lengths, g=g)
+
+    def forward(
+        self, x: torch.Tensor, x_lengths: torch.Tensor, g: Optional[torch.Tensor] = None
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        x_mask = torch.unsqueeze(
+            sequence_mask(x_lengths, x.size(2)),
+            1,
+        ).to(x.dtype)
+        x = self.pre(x) * x_mask
+        x = self.enc(x, x_mask, g=g)
+        stats = self.proj(x) * x_mask
+        m, logs = torch.split(stats, self.out_channels, dim=1)
+        z = (m + torch.randn_like(m) * torch.exp(logs)) * x_mask
+        return z, m, logs, x_mask
+
+    def remove_weight_norm(self):
+        self.enc.remove_weight_norm()
+
+    def __prepare_scriptable__(self):
+        for hook in self.enc._forward_pre_hooks.values():
+            if (
+                hook.__module__ == "torch.nn.utils.weight_norm"
+                and hook.__class__.__name__ == "WeightNorm"
+            ):
+                torch.nn.utils.remove_weight_norm(self.enc)
+        return self

rvc/layers/generators.py

@@ -0,0 +1,228 @@
+from typing import Optional, List, Tuple
+
+import torch
+from torch import nn
+from torch.nn import Conv1d, ConvTranspose1d
+from torch.nn import functional as F
+from torch.nn.utils import remove_weight_norm, weight_norm
+
+from .residuals import ResBlock1, ResBlock2, LRELU_SLOPE
+from .utils import call_weight_data_normal_if_Conv
+
+
+class Generator(torch.nn.Module):
+    def __init__(
+        self,
+        initial_channel: int,
+        resblock: str,
+        resblock_kernel_sizes: List[int],
+        resblock_dilation_sizes: List[List[int]],
+        upsample_rates: List[int],
+        upsample_initial_channel: int,
+        upsample_kernel_sizes: List[int],
+        gin_channels: int = 0,
+    ):
+        super(Generator, self).__init__()
+        self.num_kernels = len(resblock_kernel_sizes)
+        self.num_upsamples = len(upsample_rates)
+
+        self.conv_pre = Conv1d(
+            initial_channel, upsample_initial_channel, 7, 1, padding=3
+        )
+
+        self.ups = nn.ModuleList()
+        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
+            self.ups.append(
+                weight_norm(
+                    ConvTranspose1d(
+                        upsample_initial_channel // (2**i),
+                        upsample_initial_channel // (2 ** (i + 1)),
+                        k,
+                        u,
+                        padding=(k - u) // 2,
+                    )
+                )
+            )
+
+        self.resblocks = nn.ModuleList()
+        resblock_module = ResBlock1 if resblock == "1" else ResBlock2
+        for i in range(len(self.ups)):
+            ch = upsample_initial_channel // (2 ** (i + 1))
+            for k, d in zip(resblock_kernel_sizes, resblock_dilation_sizes):
+                self.resblocks.append(resblock_module(ch, k, d))
+
+        self.conv_post = Conv1d(ch, 1, 7, 1, padding=3, bias=False)
+        self.ups.apply(call_weight_data_normal_if_Conv)
+
+        if gin_channels != 0:
+            self.cond = nn.Conv1d(gin_channels, upsample_initial_channel, 1)
+
+    def __call__(
+        self,
+        x: torch.Tensor,
+        g: Optional[torch.Tensor] = None,
+        n_res: Optional[int] = None,
+    ) -> torch.Tensor:
+        return super().__call__(x, g=g, n_res=n_res)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        g: Optional[torch.Tensor] = None,
+        n_res: Optional[int] = None,
+    ):
+        if n_res is not None:
+            n = int(n_res)
+            if n != x.shape[-1]:
+                x = F.interpolate(x, size=n, mode="linear")
+
+        x = self.conv_pre(x)
+        if g is not None:
+            x = x + self.cond(g)
+
+        for i in range(self.num_upsamples):
+            x = F.leaky_relu(x, LRELU_SLOPE)
+            x = self.ups[i](x)
+            n = i * self.num_kernels
+            xs = self.resblocks[n](x)
+            for j in range(1, self.num_kernels):
+                xs += self.resblocks[n + j](x)
+            x = xs / self.num_kernels
+
+        x = F.leaky_relu(x)
+        x = self.conv_post(x)
+        x = torch.tanh(x)
+
+        return x
+
+    def __prepare_scriptable__(self):
+        for l in self.ups:
+            for hook in l._forward_pre_hooks.values():
+                # The hook we want to remove is an instance of WeightNorm class, so
+                # normally we would do `if isinstance(...)` but this class is not accessible
+                # because of shadowing, so we check the module name directly.
+                # https://github.com/pytorch/pytorch/blob/be0ca00c5ce260eb5bcec3237357f7a30cc08983/torch/nn/utils/__init__.py#L3
+                if (
+                    hook.__module__ == "torch.nn.utils.weight_norm"
+                    and hook.__class__.__name__ == "WeightNorm"
+                ):
+                    torch.nn.utils.remove_weight_norm(l)
+
+        for l in self.resblocks:
+            for hook in l._forward_pre_hooks.values():
+                if (
+                    hook.__module__ == "torch.nn.utils.weight_norm"
+                    and hook.__class__.__name__ == "WeightNorm"
+                ):
+                    torch.nn.utils.remove_weight_norm(l)
+        return self
+
+    def remove_weight_norm(self):
+        for l in self.ups:
+            remove_weight_norm(l)
+        for l in self.resblocks:
+            l.remove_weight_norm()
+
+
+class SineGenerator(torch.nn.Module):
+    """Definition of sine generator
+    SineGenerator(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-wavefrom (default 0.1)
+    noise_std: std of Gaussian noise (default 0.003)
+    voiced_thoreshold: 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(torch.pi) or cos(0)
+    """
+
+    def __init__(
+        self,
+        samp_rate: int,
+        harmonic_num: int = 0,
+        sine_amp: float = 0.1,
+        noise_std: float = 0.003,
+        voiced_threshold: int = 0,
+    ):
+        super(SineGenerator, self).__init__()
+        self.sine_amp = sine_amp
+        self.noise_std = noise_std
+        self.harmonic_num = harmonic_num
+        self.dim = harmonic_num + 1
+        self.sampling_rate = samp_rate
+        self.voiced_threshold = voiced_threshold
+
+    def __call__(
+        self, f0: torch.Tensor, upp: int
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        return super().__call__(f0, upp)
+
+    def forward(
+        self, f0: torch.Tensor, upp: int
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """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)
+        """
+        with torch.no_grad():
+            f0 = f0[:, None].transpose(1, 2)
+            f0_buf = torch.zeros(f0.shape[0], f0.shape[1], self.dim, device=f0.device)
+            # fundamental component
+            f0_buf[:, :, 0] = f0[:, :, 0]
+            for idx in range(self.harmonic_num):
+                f0_buf[:, :, idx + 1] = f0_buf[:, :, 0] * (
+                    idx + 2
+                )  # idx + 2: the (idx+1)-th overtone, (idx+2)-th harmonic
+            rad_values = (
+                f0_buf / self.sampling_rate
+            ) % 1  ###%1意味着n_har的乘积无法后处理优化
+            rand_ini = torch.rand(
+                f0_buf.shape[0], f0_buf.shape[2], device=f0_buf.device
+            )
+            rand_ini[:, 0] = 0
+            rad_values[:, 0, :] = rad_values[:, 0, :] + rand_ini
+            tmp_over_one = torch.cumsum(
+                rad_values, 1
+            )  # % 1  #####%1意味着后面的cumsum无法再优化
+            tmp_over_one *= upp
+            tmp_over_one: torch.Tensor = F.interpolate(
+                tmp_over_one.transpose(2, 1),
+                scale_factor=float(upp),
+                mode="linear",
+                align_corners=True,
+            ).transpose(2, 1)
+            rad_values: torch.Tensor = F.interpolate(
+                rad_values.transpose(2, 1), scale_factor=float(upp), mode="nearest"
+            ).transpose(
+                2, 1
+            )  #######
+            tmp_over_one %= 1
+            tmp_over_one_idx = (tmp_over_one[:, 1:, :] - tmp_over_one[:, :-1, :]) < 0
+            cumsum_shift = torch.zeros_like(rad_values)
+            cumsum_shift[:, 1:, :] = tmp_over_one_idx * -1.0
+            sine_waves = torch.sin(
+                torch.cumsum(rad_values + cumsum_shift, dim=1) * 2 * torch.pi
+            )
+            sine_waves = sine_waves * self.sine_amp
+            uv = self._f02uv(f0)
+            uv: torch.Tensor = F.interpolate(
+                uv.transpose(2, 1), scale_factor=float(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, uv, noise
+
+    def _f02uv(self, f0):
+        # generate uv signal
+        uv = torch.ones_like(f0)
+        uv = uv * (f0 > self.voiced_threshold)
+        if uv.device.type == "privateuseone":  # for DirectML
+            uv = uv.float()
+        return uv

rvc/layers/norms.py

@@ -0,0 +1,148 @@
+from typing import Optional
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from .utils import activate_add_tanh_sigmoid_multiply
+
+
+class LayerNorm(nn.Module):
+    def __init__(self, channels: int, eps: float = 1e-5):
+        super(LayerNorm, self).__init__()
+        self.channels = channels
+        self.eps = eps
+
+        self.gamma = nn.Parameter(torch.ones(channels))
+        self.beta = nn.Parameter(torch.zeros(channels))
+
+    def forward(self, x: torch.Tensor):
+        x = x.transpose(1, -1)
+        x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)
+        return x.transpose(1, -1)
+
+
+class WN(torch.nn.Module):
+    def __init__(
+        self,
+        hidden_channels: int,
+        kernel_size: int,
+        dilation_rate: int,
+        n_layers: int,
+        gin_channels: int = 0,
+        p_dropout: int = 0,
+    ):
+        super(WN, self).__init__()
+        assert kernel_size % 2 == 1
+        self.hidden_channels = hidden_channels
+        self.kernel_size = (kernel_size,)
+        self.dilation_rate = dilation_rate
+        self.n_layers = n_layers
+        self.gin_channels = gin_channels
+        self.p_dropout = float(p_dropout)
+
+        self.in_layers = torch.nn.ModuleList()
+        self.res_skip_layers = torch.nn.ModuleList()
+        self.drop = nn.Dropout(float(p_dropout))
+
+        if gin_channels != 0:
+            cond_layer = torch.nn.Conv1d(
+                gin_channels, 2 * hidden_channels * n_layers, 1
+            )
+            self.cond_layer = torch.nn.utils.weight_norm(cond_layer, name="weight")
+
+        for i in range(n_layers):
+            dilation = dilation_rate**i
+            padding = int((kernel_size * dilation - dilation) / 2)
+            in_layer = torch.nn.Conv1d(
+                hidden_channels,
+                2 * hidden_channels,
+                kernel_size,
+                dilation=dilation,
+                padding=padding,
+            )
+            in_layer = torch.nn.utils.weight_norm(in_layer, name="weight")
+            self.in_layers.append(in_layer)
+
+            # last one is not necessary
+            if i < n_layers - 1:
+                res_skip_channels = 2 * hidden_channels
+            else:
+                res_skip_channels = hidden_channels
+
+            res_skip_layer = torch.nn.Conv1d(hidden_channels, res_skip_channels, 1)
+            res_skip_layer = torch.nn.utils.weight_norm(res_skip_layer, name="weight")
+            self.res_skip_layers.append(res_skip_layer)
+
+    def __call__(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        g: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        return super().__call__(x, x_mask, g=g)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        g: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        output = torch.zeros_like(x)
+
+        if g is not None:
+            g = self.cond_layer(g)
+
+        for i, (in_layer, res_skip_layer) in enumerate(
+            zip(self.in_layers, self.res_skip_layers)
+        ):
+            x_in: torch.Tensor = in_layer(x)
+            if g is not None:
+                cond_offset = i * 2 * self.hidden_channels
+                g_l = g[:, cond_offset : cond_offset + 2 * self.hidden_channels, :]
+            else:
+                g_l = torch.zeros_like(x_in)
+
+            acts = activate_add_tanh_sigmoid_multiply(x_in, g_l, self.hidden_channels)
+            acts: torch.Tensor = self.drop(acts)
+
+            res_skip_acts: torch.Tensor = res_skip_layer(acts)
+            if i < self.n_layers - 1:
+                res_acts = res_skip_acts[:, : self.hidden_channels, :]
+                x = (x + res_acts) * x_mask
+                output = output + res_skip_acts[:, self.hidden_channels :, :]
+            else:
+                output = output + res_skip_acts
+        return output * x_mask
+
+    def remove_weight_norm(self):
+        if self.gin_channels != 0:
+            torch.nn.utils.remove_weight_norm(self.cond_layer)
+        for l in self.in_layers:
+            torch.nn.utils.remove_weight_norm(l)
+        for l in self.res_skip_layers:
+            torch.nn.utils.remove_weight_norm(l)
+
+    def __prepare_scriptable__(self):
+        if self.gin_channels != 0:
+            for hook in self.cond_layer._forward_pre_hooks.values():
+                if (
+                    hook.__module__ == "torch.nn.utils.weight_norm"
+                    and hook.__class__.__name__ == "WeightNorm"
+                ):
+                    torch.nn.utils.remove_weight_norm(self.cond_layer)
+        for l in self.in_layers:
+            for hook in l._forward_pre_hooks.values():
+                if (
+                    hook.__module__ == "torch.nn.utils.weight_norm"
+                    and hook.__class__.__name__ == "WeightNorm"
+                ):
+                    torch.nn.utils.remove_weight_norm(l)
+        for l in self.res_skip_layers:
+            for hook in l._forward_pre_hooks.values():
+                if (
+                    hook.__module__ == "torch.nn.utils.weight_norm"
+                    and hook.__class__.__name__ == "WeightNorm"
+                ):
+                    torch.nn.utils.remove_weight_norm(l)
+        return self

rvc/layers/nsf.py

@@ -0,0 +1,216 @@
+from typing import Optional, List
+import math
+
+import torch
+from torch import nn
+from torch.nn import Conv1d, ConvTranspose1d
+from torch.nn import functional as F
+from torch.nn.utils import remove_weight_norm, weight_norm
+
+from .generators import SineGenerator
+from .residuals import ResBlock1, ResBlock2, LRELU_SLOPE
+from .utils import call_weight_data_normal_if_Conv
+
+
+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: int,
+        harmonic_num: int = 0,
+        sine_amp: float = 0.1,
+        add_noise_std: float = 0.003,
+        voiced_threshod: int = 0,
+    ):
+        super(SourceModuleHnNSF, self).__init__()
+
+        self.sine_amp = sine_amp
+        self.noise_std = add_noise_std
+        # to produce sine waveforms
+        self.l_sin_gen = SineGenerator(
+            sampling_rate, harmonic_num, sine_amp, add_noise_std, voiced_threshod
+        )
+        # 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 __call__(self, x: torch.Tensor, upp: int = 1) -> torch.Tensor:
+        return super().__call__(x, upp=upp)
+
+    def forward(self, x: torch.Tensor, upp: int = 1) -> torch.Tensor:
+        sine_wavs, _, _ = self.l_sin_gen(x, upp)
+        sine_wavs = sine_wavs.to(dtype=self.l_linear.weight.dtype)
+        sine_merge: torch.Tensor = self.l_tanh(self.l_linear(sine_wavs))
+        return sine_merge  # , None, None  # noise, uv
+
+
+class NSFGenerator(torch.nn.Module):
+    def __init__(
+        self,
+        initial_channel: int,
+        resblock: str,
+        resblock_kernel_sizes: List[int],
+        resblock_dilation_sizes: List[List[int]],
+        upsample_rates: List[int],
+        upsample_initial_channel: int,
+        upsample_kernel_sizes: List[int],
+        gin_channels: int,
+        sr: int,
+    ):
+        super(NSFGenerator, self).__init__()
+        self.num_kernels = len(resblock_kernel_sizes)
+        self.num_upsamples = len(upsample_rates)
+
+        self.f0_upsamp = torch.nn.Upsample(scale_factor=math.prod(upsample_rates))
+        self.m_source = SourceModuleHnNSF(sampling_rate=sr, harmonic_num=0)
+        self.noise_convs = nn.ModuleList()
+        self.conv_pre = Conv1d(
+            initial_channel, upsample_initial_channel, 7, 1, padding=3
+        )
+        resblock = ResBlock1 if resblock == "1" else ResBlock2
+
+        self.ups = nn.ModuleList()
+        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
+            c_cur = upsample_initial_channel // (2 ** (i + 1))
+            self.ups.append(
+                weight_norm(
+                    ConvTranspose1d(
+                        upsample_initial_channel // (2**i),
+                        upsample_initial_channel // (2 ** (i + 1)),
+                        k,
+                        u,
+                        padding=(k - u) // 2,
+                    )
+                )
+            )
+            if i + 1 < len(upsample_rates):
+                stride_f0 = math.prod(upsample_rates[i + 1 :])
+                self.noise_convs.append(
+                    Conv1d(
+                        1,
+                        c_cur,
+                        kernel_size=stride_f0 * 2,
+                        stride=stride_f0,
+                        padding=stride_f0 // 2,
+                    )
+                )
+            else:
+                self.noise_convs.append(Conv1d(1, c_cur, kernel_size=1))
+
+        self.resblocks = nn.ModuleList()
+        for i in range(len(self.ups)):
+            ch: int = upsample_initial_channel // (2 ** (i + 1))
+            for j, (k, d) in enumerate(
+                zip(resblock_kernel_sizes, resblock_dilation_sizes)
+            ):
+                self.resblocks.append(resblock(ch, k, d))
+
+        self.conv_post = Conv1d(ch, 1, 7, 1, padding=3, bias=False)
+        self.ups.apply(call_weight_data_normal_if_Conv)
+
+        if gin_channels != 0:
+            self.cond = nn.Conv1d(gin_channels, upsample_initial_channel, 1)
+
+        self.upp = math.prod(upsample_rates)
+
+        self.lrelu_slope = LRELU_SLOPE
+
+    def __call__(
+        self,
+        x: torch.Tensor,
+        f0: torch.Tensor,
+        g: Optional[torch.Tensor] = None,
+        n_res: Optional[int] = None,
+    ) -> torch.Tensor:
+        return super().__call__(x, f0, g=g, n_res=n_res)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        f0: torch.Tensor,
+        g: Optional[torch.Tensor] = None,
+        n_res: Optional[int] = None,
+    ) -> torch.Tensor:
+        har_source = self.m_source(f0, self.upp)
+        har_source = har_source.transpose(1, 2)
+
+        if n_res is not None:
+            n_res = int(n_res)
+            if n_res * self.upp != har_source.shape[-1]:
+                har_source = F.interpolate(
+                    har_source, size=n_res * self.upp, mode="linear"
+                )
+            if n_res != x.shape[-1]:
+                x = F.interpolate(x, size=n_res, mode="linear")
+
+        x = self.conv_pre(x)
+        if g is not None:
+            x = x + self.cond(g)
+        # torch.jit.script() does not support direct indexing of torch modules
+        # That's why I wrote this
+        for i, (ups, noise_convs) in enumerate(zip(self.ups, self.noise_convs)):
+            if i < self.num_upsamples:
+                x = F.leaky_relu(x, self.lrelu_slope)
+                x = ups(x)
+                x_source = noise_convs(har_source)
+                x = x + x_source
+                xs: Optional[torch.Tensor] = None
+                l = [i * self.num_kernels + j for j in range(self.num_kernels)]
+                for j, resblock in enumerate(self.resblocks):
+                    if j in l:
+                        if xs is None:
+                            xs = resblock(x)
+                        else:
+                            xs += resblock(x)
+                # This assertion cannot be ignored! \
+                # If ignored, it will cause torch.jit.script() compilation errors
+                assert isinstance(xs, torch.Tensor)
+                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):
+        for l in self.ups:
+            remove_weight_norm(l)
+        for l in self.resblocks:
+            l.remove_weight_norm()
+
+    def __prepare_scriptable__(self):
+        for l in self.ups:
+            for hook in l._forward_pre_hooks.values():
+                # The hook we want to remove is an instance of WeightNorm class, so
+                # normally we would do `if isinstance(...)` but this class is not accessible
+                # because of shadowing, so we check the module name directly.
+                # https://github.com/pytorch/pytorch/blob/be0ca00c5ce260eb5bcec3237357f7a30cc08983/torch/nn/utils/__init__.py#L3
+                if (
+                    hook.__module__ == "torch.nn.utils.weight_norm"
+                    and hook.__class__.__name__ == "WeightNorm"
+                ):
+                    torch.nn.utils.remove_weight_norm(l)
+        for l in self.resblocks:
+            for hook in self.resblocks._forward_pre_hooks.values():
+                if (
+                    hook.__module__ == "torch.nn.utils.weight_norm"
+                    and hook.__class__.__name__ == "WeightNorm"
+                ):
+                    torch.nn.utils.remove_weight_norm(l)
+        return self

rvc/layers/residuals.py

@@ -0,0 +1,353 @@
+from typing import Optional, List, Tuple
+
+import torch
+from torch import nn
+from torch.nn import Conv1d
+from torch.nn import functional as F
+from torch.nn.utils import remove_weight_norm, weight_norm
+
+from .norms import WN
+from .utils import (
+    get_padding,
+    call_weight_data_normal_if_Conv,
+)
+
+LRELU_SLOPE = 0.1
+
+
+class ResBlock1(torch.nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        kernel_size: int = 3,
+        dilation: List[int] = (1, 3, 5),
+    ):
+        super(ResBlock1, self).__init__()
+
+        self.convs1 = nn.ModuleList()
+        for d in dilation:
+            self.convs1.append(
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=d,
+                        padding=get_padding(kernel_size, d),
+                    )
+                ),
+            )
+        self.convs1.apply(call_weight_data_normal_if_Conv)
+
+        self.convs2 = nn.ModuleList()
+        for _ in dilation:
+            self.convs2.append(
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=get_padding(kernel_size, 1),
+                    )
+                ),
+            )
+        self.convs2.apply(call_weight_data_normal_if_Conv)
+        self.lrelu_slope = LRELU_SLOPE
+
+    def __call__(
+        self,
+        x: torch.Tensor,
+        x_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        return super().__call__(x, x_mask=x_mask)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        for c1, c2 in zip(self.convs1, self.convs2):
+            xt = F.leaky_relu(x, self.lrelu_slope)
+            if x_mask is not None:
+                xt = xt * x_mask
+            xt = c1(xt)
+            xt = F.leaky_relu(xt, self.lrelu_slope)
+            if x_mask is not None:
+                xt = xt * x_mask
+            xt = c2(xt)
+            x = xt + x
+        if x_mask is not None:
+            x = x * x_mask
+        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)
+
+    def __prepare_scriptable__(self):
+        for l in self.convs1:
+            for hook in l._forward_pre_hooks.values():
+                if (
+                    hook.__module__ == "torch.nn.utils.weight_norm"
+                    and hook.__class__.__name__ == "WeightNorm"
+                ):
+                    torch.nn.utils.remove_weight_norm(l)
+        for l in self.convs2:
+            for hook in l._forward_pre_hooks.values():
+                if (
+                    hook.__module__ == "torch.nn.utils.weight_norm"
+                    and hook.__class__.__name__ == "WeightNorm"
+                ):
+                    torch.nn.utils.remove_weight_norm(l)
+        return self
+
+
+class ResBlock2(torch.nn.Module):
+    """
+    Actually this module is not used currently
+    because all configs specified "resblock": "1"
+    """
+
+    def __init__(
+        self,
+        channels: int,
+        kernel_size=3,
+        dilation: List[int] = (1, 3),
+    ):
+        super(ResBlock2, self).__init__()
+        self.convs = nn.ModuleList()
+        for d in dilation:
+            self.convs.append(
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=d,
+                        padding=get_padding(kernel_size, d),
+                    )
+                ),
+            )
+        self.convs.apply(call_weight_data_normal_if_Conv)
+        self.lrelu_slope = LRELU_SLOPE
+
+    def __call__(
+        self,
+        x: torch.Tensor,
+        x_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        return super().__call__(x, x_mask=x_mask)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        for c in self.convs:
+            xt = F.leaky_relu(x, self.lrelu_slope)
+            if x_mask is not None:
+                xt = xt * x_mask
+            xt = c(xt)
+            x = xt + x
+        if x_mask is not None:
+            x = x * x_mask
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs:
+            remove_weight_norm(l)
+
+    def __prepare_scriptable__(self):
+        for l in self.convs:
+            for hook in l._forward_pre_hooks.values():
+                if (
+                    hook.__module__ == "torch.nn.utils.weight_norm"
+                    and hook.__class__.__name__ == "WeightNorm"
+                ):
+                    torch.nn.utils.remove_weight_norm(l)
+        return self
+
+
+class ResidualCouplingLayer(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        hidden_channels: int,
+        kernel_size: int,
+        dilation_rate: int,
+        n_layers: int,
+        p_dropout: int = 0,
+        gin_channels: int = 0,
+        mean_only: bool = False,
+    ):
+        assert channels % 2 == 0, "channels should be divisible by 2"
+        super(ResidualCouplingLayer, self).__init__()
+        self.channels = channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.dilation_rate = dilation_rate
+        self.n_layers = n_layers
+        self.half_channels = channels // 2
+        self.mean_only = mean_only
+
+        self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)
+        self.enc = WN(
+            hidden_channels,
+            kernel_size,
+            dilation_rate,
+            n_layers,
+            p_dropout=float(p_dropout),
+            gin_channels=gin_channels,
+        )
+        self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)
+        self.post.weight.data.zero_()
+        self.post.bias.data.zero_()
+
+    def __call__(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        g: Optional[torch.Tensor] = None,
+        reverse: bool = False,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        return super().__call__(x, x_mask, g=g, reverse=reverse)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        g: Optional[torch.Tensor] = None,
+        reverse: bool = False,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        x0, x1 = torch.split(x, [self.half_channels] * 2, 1)
+        h = self.pre(x0) * x_mask
+        h = self.enc(h, x_mask, g=g)
+        stats = self.post(h) * x_mask
+        if not self.mean_only:
+            m, logs = torch.split(stats, [self.half_channels] * 2, 1)
+        else:
+            m = stats
+            logs = torch.zeros_like(m)
+
+        if not reverse:
+            x1 = m + x1 * torch.exp(logs) * x_mask
+            x = torch.cat([x0, x1], 1)
+            logdet = torch.sum(logs, [1, 2])
+            return x, logdet
+
+        x1 = (x1 - m) * torch.exp(-logs) * x_mask
+        x = torch.cat([x0, x1], 1)
+        return x, torch.zeros([1])
+
+    def remove_weight_norm(self):
+        self.enc.remove_weight_norm()
+
+    def __prepare_scriptable__(self):
+        for hook in self.enc._forward_pre_hooks.values():
+            if (
+                hook.__module__ == "torch.nn.utils.weight_norm"
+                and hook.__class__.__name__ == "WeightNorm"
+            ):
+                torch.nn.utils.remove_weight_norm(self.enc)
+        return self
+
+
+class ResidualCouplingBlock(nn.Module):
+    class Flip(nn.Module):
+        """
+        torch.jit.script() Compiled functions
+        can't take variable number of arguments or
+        use keyword-only arguments with defaults
+        """
+
+        def forward(
+            self,
+            x: torch.Tensor,
+            x_mask: torch.Tensor,
+            g: Optional[torch.Tensor] = None,
+            reverse: bool = False,
+        ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+            x = torch.flip(x, [1])
+            if not reverse:
+                logdet = torch.zeros(x.size(0)).to(dtype=x.dtype, device=x.device)
+                return x, logdet
+            else:
+                return x, torch.zeros([1], device=x.device)
+
+    def __init__(
+        self,
+        channels: int,
+        hidden_channels: int,
+        kernel_size: int,
+        dilation_rate: int,
+        n_layers: int,
+        n_flows: int = 4,
+        gin_channels: int = 0,
+    ):
+        super(ResidualCouplingBlock, self).__init__()
+        self.channels = channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.dilation_rate = dilation_rate
+        self.n_layers = n_layers
+        self.n_flows = n_flows
+        self.gin_channels = gin_channels
+
+        self.flows = nn.ModuleList()
+        for _ in range(n_flows):
+            self.flows.append(
+                ResidualCouplingLayer(
+                    channels,
+                    hidden_channels,
+                    kernel_size,
+                    dilation_rate,
+                    n_layers,
+                    gin_channels=gin_channels,
+                    mean_only=True,
+                )
+            )
+            self.flows.append(self.Flip())
+
+    def __call__(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        g: Optional[torch.Tensor] = None,
+        reverse: bool = False,
+    ) -> torch.Tensor:
+        return super().__call__(x, x_mask, g=g, reverse=reverse)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        g: Optional[torch.Tensor] = None,
+        reverse: bool = False,
+    ) -> torch.Tensor:
+        if not reverse:
+            for flow in self.flows:
+                x, _ = flow(x, x_mask, g=g, reverse=reverse)
+        else:
+            for flow in reversed(self.flows):
+                x, _ = flow.forward(x, x_mask, g=g, reverse=reverse)
+        return x
+
+    def remove_weight_norm(self):
+        for i in range(self.n_flows):
+            self.flows[i * 2].remove_weight_norm()
+
+    def __prepare_scriptable__(self):
+        for i in range(self.n_flows):
+            for hook in self.flows[i * 2]._forward_pre_hooks.values():
+                if (
+                    hook.__module__ == "torch.nn.utils.weight_norm"
+                    and hook.__class__.__name__ == "WeightNorm"
+                ):
+                    torch.nn.utils.remove_weight_norm(self.flows[i * 2])
+        return self

rvc/layers/synthesizers.py

@@ -0,0 +1,396 @@
+from typing import Optional, List, Union
+
+import torch
+from torch import nn
+
+
+from .encoders import TextEncoder, PosteriorEncoder
+from .generators import Generator
+from .nsf import NSFGenerator
+from .residuals import ResidualCouplingBlock
+from .utils import (
+    slice_on_last_dim,
+    rand_slice_segments_on_last_dim,
+)
+
+
+class SynthesizerTrnMsNSFsid(nn.Module):
+    def __init__(
+        self,
+        spec_channels: int,
+        segment_size: int,
+        inter_channels: int,
+        hidden_channels: int,
+        filter_channels: int,
+        n_heads: int,
+        n_layers: int,
+        kernel_size: int,
+        p_dropout: int,
+        resblock: str,
+        resblock_kernel_sizes: List[int],
+        resblock_dilation_sizes: List[List[int]],
+        upsample_rates: List[int],
+        upsample_initial_channel: int,
+        upsample_kernel_sizes: List[int],
+        spk_embed_dim: int,
+        gin_channels: int,
+        sr: Optional[Union[str, int]],
+        encoder_dim: int,
+        use_f0: bool,
+    ):
+        super().__init__()
+        if isinstance(sr, str):
+            sr = {
+                "32k": 32000,
+                "40k": 40000,
+                "48k": 48000,
+            }[sr]
+        self.spec_channels = spec_channels
+        self.inter_channels = inter_channels
+        self.hidden_channels = hidden_channels
+        self.filter_channels = filter_channels
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.kernel_size = kernel_size
+        self.p_dropout = float(p_dropout)
+        self.resblock = resblock
+        self.resblock_kernel_sizes = resblock_kernel_sizes
+        self.resblock_dilation_sizes = resblock_dilation_sizes
+        self.upsample_rates = upsample_rates
+        self.upsample_initial_channel = upsample_initial_channel
+        self.upsample_kernel_sizes = upsample_kernel_sizes
+        self.segment_size = segment_size
+        self.gin_channels = gin_channels
+        self.spk_embed_dim = spk_embed_dim
+
+        self.enc_p = TextEncoder(
+            encoder_dim,
+            inter_channels,
+            hidden_channels,
+            filter_channels,
+            n_heads,
+            n_layers,
+            kernel_size,
+            float(p_dropout),
+            f0=use_f0,
+        )
+        if use_f0:
+            self.dec = NSFGenerator(
+                inter_channels,
+                resblock,
+                resblock_kernel_sizes,
+                resblock_dilation_sizes,
+                upsample_rates,
+                upsample_initial_channel,
+                upsample_kernel_sizes,
+                gin_channels=gin_channels,
+                sr=sr,
+            )
+        else:
+            self.dec = Generator(
+                inter_channels,
+                resblock,
+                resblock_kernel_sizes,
+                resblock_dilation_sizes,
+                upsample_rates,
+                upsample_initial_channel,
+                upsample_kernel_sizes,
+                gin_channels=gin_channels,
+            )
+        self.enc_q = PosteriorEncoder(
+            spec_channels,
+            inter_channels,
+            hidden_channels,
+            5,
+            1,
+            16,
+            gin_channels=gin_channels,
+        )
+        self.flow = ResidualCouplingBlock(
+            inter_channels, hidden_channels, 5, 1, 3, gin_channels=gin_channels
+        )
+        self.emb_g = nn.Embedding(self.spk_embed_dim, gin_channels)
+
+    def remove_weight_norm(self):
+        self.dec.remove_weight_norm()
+        self.flow.remove_weight_norm()
+        if hasattr(self, "enc_q"):
+            self.enc_q.remove_weight_norm()
+
+    def __prepare_scriptable__(self):
+        for hook in self.dec._forward_pre_hooks.values():
+            # The hook we want to remove is an instance of WeightNorm class, so
+            # normally we would do `if isinstance(...)` but this class is not accessible
+            # because of shadowing, so we check the module name directly.
+            # https://github.com/pytorch/pytorch/blob/be0ca00c5ce260eb5bcec3237357f7a30cc08983/torch/nn/utils/__init__.py#L3
+            if (
+                hook.__module__ == "torch.nn.utils.weight_norm"
+                and hook.__class__.__name__ == "WeightNorm"
+            ):
+                torch.nn.utils.remove_weight_norm(self.dec)
+        for hook in self.flow._forward_pre_hooks.values():
+            if (
+                hook.__module__ == "torch.nn.utils.weight_norm"
+                and hook.__class__.__name__ == "WeightNorm"
+            ):
+                torch.nn.utils.remove_weight_norm(self.flow)
+        if hasattr(self, "enc_q"):
+            for hook in self.enc_q._forward_pre_hooks.values():
+                if (
+                    hook.__module__ == "torch.nn.utils.weight_norm"
+                    and hook.__class__.__name__ == "WeightNorm"
+                ):
+                    torch.nn.utils.remove_weight_norm(self.enc_q)
+        return self
+
+    @torch.jit.ignore
+    def forward(
+        self,
+        phone: torch.Tensor,
+        phone_lengths: torch.Tensor,
+        y: torch.Tensor,
+        y_lengths: torch.Tensor,
+        ds: Optional[torch.Tensor] = None,
+        pitch: Optional[torch.Tensor] = None,
+        pitchf: Optional[torch.Tensor] = None,
+    ):  # 这里ds是id，[bs,1]
+        # print(1,pitch.shape)#[bs,t]
+        g = self.emb_g(ds).unsqueeze(-1)  # [b, 256, 1]##1是t，广播的
+        m_p, logs_p, x_mask = self.enc_p(phone, pitch, phone_lengths)
+        z, m_q, logs_q, y_mask = self.enc_q(y, y_lengths, g=g)
+        z_p = self.flow(z, y_mask, g=g)
+        z_slice, ids_slice = rand_slice_segments_on_last_dim(
+            z, y_lengths, self.segment_size
+        )
+        if pitchf is not None:
+            pitchf = slice_on_last_dim(pitchf, ids_slice, self.segment_size)
+            o = self.dec(z_slice, pitchf, g=g)
+        else:
+            o = self.dec(z_slice, g=g)
+        return o, ids_slice, x_mask, y_mask, (z, z_p, m_p, logs_p, m_q, logs_q)
+
+    @torch.jit.export
+    def infer(
+        self,
+        phone: torch.Tensor,
+        phone_lengths: torch.Tensor,
+        sid: torch.Tensor,
+        pitch: Optional[torch.Tensor] = None,
+        pitchf: Optional[torch.Tensor] = None,  # nsff0
+        skip_head: Optional[int] = None,
+        return_length: Optional[int] = None,
+        return_length2: Optional[int] = None,
+    ):
+        g = self.emb_g(sid).unsqueeze(-1)
+        if skip_head is not None and return_length is not None:
+            head = int(skip_head)
+            length = int(return_length)
+            flow_head = head - 24
+            if flow_head < 0:
+                flow_head = 0
+            dec_head = head - flow_head
+            m_p, logs_p, x_mask = self.enc_p(phone, pitch, phone_lengths, flow_head)
+            z_p = (m_p + torch.exp(logs_p) * torch.randn_like(m_p) * 0.66666) * x_mask
+            z = self.flow(z_p, x_mask, g=g, reverse=True)
+            z = z[:, :, dec_head : dec_head + length]
+            x_mask = x_mask[:, :, dec_head : dec_head + length]
+            if pitchf is not None:
+                pitchf = pitchf[:, head : head + length]
+        else:
+            m_p, logs_p, x_mask = self.enc_p(phone, pitch, phone_lengths)
+            z_p = (m_p + torch.exp(logs_p) * torch.randn_like(m_p) * 0.66666) * x_mask
+            z = self.flow(z_p, x_mask, g=g, reverse=True)
+        del z_p, m_p, logs_p
+        if pitchf is not None:
+            o = self.dec(
+                z * x_mask,
+                pitchf,
+                g=g,
+                n_res=return_length2,
+            )
+        else:
+            o = self.dec(z * x_mask, g=g, n_res=return_length2)
+        del x_mask, z
+        return o  # , x_mask, (z, z_p, m_p, logs_p)
+
+
+class SynthesizerTrnMs256NSFsid(SynthesizerTrnMsNSFsid):
+    def __init__(
+        self,
+        spec_channels: int,
+        segment_size: int,
+        inter_channels: int,
+        hidden_channels: int,
+        filter_channels: int,
+        n_heads: int,
+        n_layers: int,
+        kernel_size: int,
+        p_dropout: int,
+        resblock: str,
+        resblock_kernel_sizes: List[int],
+        resblock_dilation_sizes: List[List[int]],
+        upsample_rates: List[int],
+        upsample_initial_channel: int,
+        upsample_kernel_sizes: List[int],
+        spk_embed_dim: int,
+        gin_channels: int,
+        sr: Union[str, int],
+    ):
+        super().__init__(
+            spec_channels,
+            segment_size,
+            inter_channels,
+            hidden_channels,
+            filter_channels,
+            n_heads,
+            n_layers,
+            kernel_size,
+            p_dropout,
+            resblock,
+            resblock_kernel_sizes,
+            resblock_dilation_sizes,
+            upsample_rates,
+            upsample_initial_channel,
+            upsample_kernel_sizes,
+            spk_embed_dim,
+            gin_channels,
+            sr,
+            256,
+            True,
+        )
+
+
+class SynthesizerTrnMs768NSFsid(SynthesizerTrnMsNSFsid):
+    def __init__(
+        self,
+        spec_channels: int,
+        segment_size: int,
+        inter_channels: int,
+        hidden_channels: int,
+        filter_channels: int,
+        n_heads: int,
+        n_layers: int,
+        kernel_size: int,
+        p_dropout: int,
+        resblock: str,
+        resblock_kernel_sizes: List[int],
+        resblock_dilation_sizes: List[List[int]],
+        upsample_rates: List[int],
+        upsample_initial_channel: int,
+        upsample_kernel_sizes: List[int],
+        spk_embed_dim: int,
+        gin_channels: int,
+        sr: Union[str, int],
+    ):
+        super().__init__(
+            spec_channels,
+            segment_size,
+            inter_channels,
+            hidden_channels,
+            filter_channels,
+            n_heads,
+            n_layers,
+            kernel_size,
+            p_dropout,
+            resblock,
+            resblock_kernel_sizes,
+            resblock_dilation_sizes,
+            upsample_rates,
+            upsample_initial_channel,
+            upsample_kernel_sizes,
+            spk_embed_dim,
+            gin_channels,
+            sr,
+            768,
+            True,
+        )
+
+
+class SynthesizerTrnMs256NSFsid_nono(SynthesizerTrnMsNSFsid):
+    def __init__(
+        self,
+        spec_channels: int,
+        segment_size: int,
+        inter_channels: int,
+        hidden_channels: int,
+        filter_channels: int,
+        n_heads: int,
+        n_layers: int,
+        kernel_size: int,
+        p_dropout: int,
+        resblock: str,
+        resblock_kernel_sizes: List[int],
+        resblock_dilation_sizes: List[List[int]],
+        upsample_rates: List[int],
+        upsample_initial_channel: int,
+        upsample_kernel_sizes: List[int],
+        spk_embed_dim: int,
+        gin_channels: int,
+        sr=None,
+    ):
+        super().__init__(
+            spec_channels,
+            segment_size,
+            inter_channels,
+            hidden_channels,
+            filter_channels,
+            n_heads,
+            n_layers,
+            kernel_size,
+            p_dropout,
+            resblock,
+            resblock_kernel_sizes,
+            resblock_dilation_sizes,
+            upsample_rates,
+            upsample_initial_channel,
+            upsample_kernel_sizes,
+            spk_embed_dim,
+            gin_channels,
+            256,
+            False,
+        )
+
+
+class SynthesizerTrnMs768NSFsid_nono(SynthesizerTrnMsNSFsid):
+    def __init__(
+        self,
+        spec_channels: int,
+        segment_size: int,
+        inter_channels: int,
+        hidden_channels: int,
+        filter_channels: int,
+        n_heads: int,
+        n_layers: int,
+        kernel_size: int,
+        p_dropout: int,
+        resblock: str,
+        resblock_kernel_sizes: List[int],
+        resblock_dilation_sizes: List[List[int]],
+        upsample_rates: List[int],
+        upsample_initial_channel: int,
+        upsample_kernel_sizes: List[int],
+        spk_embed_dim: int,
+        gin_channels: int,
+        sr=None,
+    ):
+        super().__init__(
+            spec_channels,
+            segment_size,
+            inter_channels,
+            hidden_channels,
+            filter_channels,
+            n_heads,
+            n_layers,
+            kernel_size,
+            p_dropout,
+            resblock,
+            resblock_kernel_sizes,
+            resblock_dilation_sizes,
+            upsample_rates,
+            upsample_initial_channel,
+            upsample_kernel_sizes,
+            spk_embed_dim,
+            gin_channels,
+            768,
+            False,
+        )

rvc/layers/transforms.py

@@ -0,0 +1,209 @@
+from typing import Optional
+
+import numpy as np
+import torch
+from torch.nn import functional as F
+
+DEFAULT_MIN_BIN_WIDTH = 1e-3
+DEFAULT_MIN_BIN_HEIGHT = 1e-3
+DEFAULT_MIN_DERIVATIVE = 1e-3
+
+
+def piecewise_rational_quadratic_transform(
+    inputs: torch.Tensor,
+    unnormalized_widths: torch.Tensor,
+    unnormalized_heights: torch.Tensor,
+    unnormalized_derivatives: torch.Tensor,
+    inverse: bool = False,
+    tails: Optional[str] = None,
+    tail_bound: float = 1.0,
+    min_bin_width=DEFAULT_MIN_BIN_WIDTH,
+    min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
+    min_derivative=DEFAULT_MIN_DERIVATIVE,
+):
+    if tails is None:
+        spline_fn = rational_quadratic_spline
+        spline_kwargs = {}
+    else:
+        spline_fn = unconstrained_rational_quadratic_spline
+        spline_kwargs = {"tails": tails, "tail_bound": tail_bound}
+
+    outputs, logabsdet = spline_fn(
+        inputs=inputs,
+        unnormalized_widths=unnormalized_widths,
+        unnormalized_heights=unnormalized_heights,
+        unnormalized_derivatives=unnormalized_derivatives,
+        inverse=inverse,
+        min_bin_width=min_bin_width,
+        min_bin_height=min_bin_height,
+        min_derivative=min_derivative,
+        **spline_kwargs
+    )
+    return outputs, logabsdet
+
+
+def searchsorted(bin_locations, inputs, eps=1e-6):
+    bin_locations[..., -1] += eps
+    return torch.sum(inputs[..., None] >= bin_locations, dim=-1) - 1
+
+
+def unconstrained_rational_quadratic_spline(
+    inputs: torch.Tensor,
+    unnormalized_widths: torch.Tensor,
+    unnormalized_heights: torch.Tensor,
+    unnormalized_derivatives: torch.Tensor,
+    inverse: bool = False,
+    tails: str = "linear",
+    tail_bound: float = 1.0,
+    min_bin_width=DEFAULT_MIN_BIN_WIDTH,
+    min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
+    min_derivative=DEFAULT_MIN_DERIVATIVE,
+):
+    inside_interval_mask = (inputs >= -tail_bound) & (inputs <= tail_bound)
+    outside_interval_mask = ~inside_interval_mask
+
+    outputs = torch.zeros_like(inputs)
+    logabsdet = torch.zeros_like(inputs)
+
+    if tails == "linear":
+        unnormalized_derivatives = F.pad(unnormalized_derivatives, pad=(1, 1))
+        constant = np.log(np.exp(1 - min_derivative) - 1)
+        unnormalized_derivatives[..., 0] = constant
+        unnormalized_derivatives[..., -1] = constant
+
+        outputs[outside_interval_mask] = inputs[outside_interval_mask]
+        logabsdet[outside_interval_mask] = 0
+    else:
+        raise RuntimeError("{} tails are not implemented.".format(tails))
+
+    (
+        outputs[inside_interval_mask],
+        logabsdet[inside_interval_mask],
+    ) = rational_quadratic_spline(
+        inputs=inputs[inside_interval_mask],
+        unnormalized_widths=unnormalized_widths[inside_interval_mask, :],
+        unnormalized_heights=unnormalized_heights[inside_interval_mask, :],
+        unnormalized_derivatives=unnormalized_derivatives[inside_interval_mask, :],
+        inverse=inverse,
+        left=-tail_bound,
+        right=tail_bound,
+        bottom=-tail_bound,
+        top=tail_bound,
+        min_bin_width=min_bin_width,
+        min_bin_height=min_bin_height,
+        min_derivative=min_derivative,
+    )
+
+    return outputs, logabsdet
+
+
+def rational_quadratic_spline(
+    inputs: torch.Tensor,
+    unnormalized_widths: torch.Tensor,
+    unnormalized_heights: torch.Tensor,
+    unnormalized_derivatives: torch.Tensor,
+    inverse: bool = False,
+    left=0.0,
+    right=1.0,
+    bottom=0.0,
+    top=1.0,
+    min_bin_width=DEFAULT_MIN_BIN_WIDTH,
+    min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
+    min_derivative=DEFAULT_MIN_DERIVATIVE,
+):
+    if torch.min(inputs) < left or torch.max(inputs) > right:
+        raise ValueError("Input to a transform is not within its domain")
+
+    num_bins = unnormalized_widths.shape[-1]
+
+    if min_bin_width * num_bins > 1.0:
+        raise ValueError("Minimal bin width too large for the number of bins")
+    if min_bin_height * num_bins > 1.0:
+        raise ValueError("Minimal bin height too large for the number of bins")
+
+    widths = F.softmax(unnormalized_widths, dim=-1)
+    widths = min_bin_width + (1 - min_bin_width * num_bins) * widths
+    cumwidths = torch.cumsum(widths, dim=-1)
+    cumwidths = F.pad(cumwidths, pad=(1, 0), mode="constant", value=0.0)
+    cumwidths = (right - left) * cumwidths + left
+    cumwidths[..., 0] = left
+    cumwidths[..., -1] = right
+    widths = cumwidths[..., 1:] - cumwidths[..., :-1]
+
+    derivatives = min_derivative + F.softplus(unnormalized_derivatives)
+
+    heights = F.softmax(unnormalized_heights, dim=-1)
+    heights = min_bin_height + (1 - min_bin_height * num_bins) * heights
+    cumheights = torch.cumsum(heights, dim=-1)
+    cumheights = F.pad(cumheights, pad=(1, 0), mode="constant", value=0.0)
+    cumheights = (top - bottom) * cumheights + bottom
+    cumheights[..., 0] = bottom
+    cumheights[..., -1] = top
+    heights = cumheights[..., 1:] - cumheights[..., :-1]
+
+    if inverse:
+        bin_idx = searchsorted(cumheights, inputs)[..., None]
+    else:
+        bin_idx = searchsorted(cumwidths, inputs)[..., None]
+
+    input_cumwidths = cumwidths.gather(-1, bin_idx)[..., 0]
+    input_bin_widths = widths.gather(-1, bin_idx)[..., 0]
+
+    input_cumheights = cumheights.gather(-1, bin_idx)[..., 0]
+    delta = heights / widths
+    input_delta = delta.gather(-1, bin_idx)[..., 0]
+
+    input_derivatives = derivatives.gather(-1, bin_idx)[..., 0]
+    input_derivatives_plus_one = derivatives[..., 1:].gather(-1, bin_idx)[..., 0]
+
+    input_heights = heights.gather(-1, bin_idx)[..., 0]
+
+    if inverse:
+        a = (inputs - input_cumheights) * (
+            input_derivatives + input_derivatives_plus_one - 2 * input_delta
+        ) + input_heights * (input_delta - input_derivatives)
+        b = input_heights * input_derivatives - (inputs - input_cumheights) * (
+            input_derivatives + input_derivatives_plus_one - 2 * input_delta
+        )
+        c = -input_delta * (inputs - input_cumheights)
+
+        discriminant = b.pow(2) - 4 * a * c
+        assert (discriminant >= 0).all()
+
+        root = (2 * c) / (-b - torch.sqrt(discriminant))
+        outputs = root * input_bin_widths + input_cumwidths
+
+        theta_one_minus_theta = root * (1 - root)
+        denominator = input_delta + (
+            (input_derivatives + input_derivatives_plus_one - 2 * input_delta)
+            * theta_one_minus_theta
+        )
+        derivative_numerator = input_delta.pow(2) * (
+            input_derivatives_plus_one * root.pow(2)
+            + 2 * input_delta * theta_one_minus_theta
+            + input_derivatives * (1 - root).pow(2)
+        )
+        logabsdet = torch.log(derivative_numerator) - 2 * torch.log(denominator)
+
+        return outputs, -logabsdet
+    else:
+        theta = (inputs - input_cumwidths) / input_bin_widths
+        theta_one_minus_theta = theta * (1 - theta)
+
+        numerator = input_heights * (
+            input_delta * theta.pow(2) + input_derivatives * theta_one_minus_theta
+        )
+        denominator = input_delta + (
+            (input_derivatives + input_derivatives_plus_one - 2 * input_delta)
+            * theta_one_minus_theta
+        )
+        outputs = input_cumheights + numerator / denominator
+
+        derivative_numerator = input_delta.pow(2) * (
+            input_derivatives_plus_one * theta.pow(2)
+            + 2 * input_delta * theta_one_minus_theta
+            + input_derivatives * (1 - theta).pow(2)
+        )
+        logabsdet = torch.log(derivative_numerator) - 2 * torch.log(denominator)
+
+        return outputs, logabsdet

rvc/layers/utils.py

@@ -0,0 +1,82 @@
+from typing import List, Optional, Tuple, Iterator
+
+import torch
+
+
+def call_weight_data_normal_if_Conv(m: torch.nn.Module):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        mean = 0.0
+        std = 0.01
+        m.weight.data.normal_(mean, std)
+
+
+def get_padding(kernel_size: int, dilation=1) -> int:
+    return int((kernel_size * dilation - dilation) / 2)
+
+
+def slice_on_last_dim(
+    x: torch.Tensor,
+    start_indices: List[int],
+    segment_size=4,
+) -> torch.Tensor:
+    new_shape = [*x.shape]
+    new_shape[-1] = segment_size
+    ret = torch.empty(new_shape, device=x.device)
+    for i in range(x.size(0)):
+        idx_str = start_indices[i]
+        idx_end = idx_str + segment_size
+        ret[i, ..., :] = x[i, ..., idx_str:idx_end]
+    return ret
+
+
+def rand_slice_segments_on_last_dim(
+    x: torch.Tensor,
+    x_lengths: int = None,
+    segment_size=4,
+) -> Tuple[torch.Tensor, List[int]]:
+    b, _, t = x.size()
+    if x_lengths is None:
+        x_lengths = t
+    ids_str_max = x_lengths - segment_size + 1
+    ids_str = (torch.rand([b]).to(device=x.device) * ids_str_max).to(dtype=torch.long)
+    ret = slice_on_last_dim(x, ids_str, segment_size)
+    return ret, ids_str
+
+
+@torch.jit.script
+def activate_add_tanh_sigmoid_multiply(
+    input_a: torch.Tensor, input_b: torch.Tensor, n_channels: int
+) -> torch.Tensor:
+    in_act = input_a + input_b
+    t_act = torch.tanh(in_act[:, :n_channels, :])
+    s_act = torch.sigmoid(in_act[:, n_channels:, :])
+    acts = t_act * s_act
+    return acts
+
+
+def sequence_mask(
+    length: torch.Tensor,
+    max_length: Optional[int] = None,
+) -> torch.BoolTensor:
+    if max_length is None:
+        max_length = int(length.max())
+    x = torch.arange(max_length, dtype=length.dtype, device=length.device)
+    return x.unsqueeze(0) < length.unsqueeze(1)
+
+
+def total_grad_norm(
+    parameters: Iterator[torch.nn.Parameter],
+    norm_type: float = 2.0,
+) -> float:
+    norm_type = float(norm_type)
+    total_norm = 0.0
+
+    for p in parameters:
+        if p.grad is None:
+            continue
+        param_norm = p.grad.data.norm(norm_type)
+        total_norm += float(param_norm.item()) ** norm_type
+    total_norm = total_norm ** (1.0 / norm_type)
+
+    return total_norm

rvc/onnx/__init__.py

@@ -0,0 +1,2 @@
+from .infer import RVC
+from .exporter import export_onnx

rvc/onnx/exporter.py

@@ -0,0 +1,52 @@
+import torch
+
+from .synthesizer import SynthesizerTrnMsNSFsid
+
+
+def export_onnx(from_cpkt_pth: str, to_onnx_pth: str) -> str:
+    cpt = torch.load(from_cpkt_pth, map_location="cpu")
+    cpt["config"][-3] = cpt["weight"]["emb_g.weight"].shape[0]
+    vec_channels = 256 if cpt.get("version", "v1") == "v1" else 768
+
+    test_phone = torch.rand(1, 200, vec_channels)  # hidden unit
+    test_phone_lengths = torch.tensor([200]).long()  # hidden unit 长度（貌似没啥用）
+    test_pitch = torch.randint(size=(1, 200), low=5, high=255)  # 基频（单位赫兹）
+    test_pitchf = torch.rand(1, 200)  # nsf基频
+    test_ds = torch.LongTensor([0])  # 说话人ID
+    test_rnd = torch.rand(1, 192, 200)  # 噪声（加入随机因子）
+
+    device = "cpu"  # 导出时设备（不影响使用模型）
+
+    net_g = SynthesizerTrnMsNSFsid(
+        *cpt["config"], encoder_dim=vec_channels
+    )  # fp32导出（C++要支持fp16必须手动将内存重新排列所以暂时不用fp16）
+    net_g.load_state_dict(cpt["weight"], strict=False)
+    input_names = ["phone", "phone_lengths", "pitch", "pitchf", "ds", "rnd"]
+    output_names = [
+        "audio",
+    ]
+    # net_g.construct_spkmixmap() #多角色混合轨道导出
+    torch.onnx.export(
+        net_g,
+        (
+            test_phone.to(device),
+            test_phone_lengths.to(device),
+            test_pitch.to(device),
+            test_pitchf.to(device),
+            test_ds.to(device),
+            test_rnd.to(device),
+        ),
+        to_onnx_pth,
+        dynamic_axes={
+            "phone": [1],
+            "pitch": [1],
+            "pitchf": [1],
+            "rnd": [2],
+        },
+        do_constant_folding=False,
+        opset_version=17,
+        verbose=False,
+        input_names=input_names,
+        output_names=output_names,
+    )
+    return "Finished"

rvc/onnx/infer.py

@@ -0,0 +1,146 @@
+import typing
+import os
+
+import librosa
+import numpy as np
+import onnxruntime
+
+from rvc.f0 import (
+    PM,
+    Harvest,
+    Dio,
+    F0Predictor,
+)
+
+
+class Model:
+    def __init__(
+        self,
+        path: typing.Union[str, bytes, os.PathLike],
+        device: typing.Literal["cpu", "cuda", "dml"] = "cpu",
+    ):
+        if device == "cpu":
+            providers = ["CPUExecutionProvider"]
+        elif device == "cuda":
+            providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+        elif device == "dml":
+            providers = ["DmlExecutionProvider"]
+        else:
+            raise RuntimeError("Unsportted Device")
+        self.model = onnxruntime.InferenceSession(path, providers=providers)
+
+
+class ContentVec(Model):
+    def __init__(
+        self,
+        vec_path: typing.Union[str, bytes, os.PathLike],
+        device: typing.Literal["cpu", "cuda", "dml"] = "cpu",
+    ):
+        super().__init__(vec_path, device)
+
+    def __call__(self, wav: np.ndarray[typing.Any, np.dtype]):
+        return self.forward(wav)
+
+    def forward(self, wav: np.ndarray[typing.Any, np.dtype]):
+        if wav.ndim == 2:  # double channels
+            wav = wav.mean(-1)
+        assert wav.ndim == 1, wav.ndim
+        wav = np.expand_dims(np.expand_dims(wav, 0), 0)
+        onnx_input = {self.model.get_inputs()[0].name: wav}
+        logits = self.model.run(None, onnx_input)[0]
+        return logits.transpose(0, 2, 1)
+
+
+predictors: typing.Dict[str, F0Predictor] = {
+    "pm": PM,
+    "harvest": Harvest,
+    "dio": Dio,
+}
+
+
+def get_f0_predictor(
+    f0_method: str, hop_length: int, sampling_rate: int
+) -> F0Predictor:
+    return predictors[f0_method](hop_length=hop_length, sampling_rate=sampling_rate)
+
+
+class RVC(Model):
+    def __init__(
+        self,
+        model_path: typing.Union[str, bytes, os.PathLike],
+        hop_len=512,
+        vec_path: typing.Union[str, bytes, os.PathLike] = "vec-768-layer-12.onnx",
+        device: typing.Literal["cpu", "cuda", "dml"] = "cpu",
+    ):
+        super().__init__(model_path, device)
+        self.vec_model = ContentVec(vec_path, device)
+        self.hop_len = hop_len
+
+    def infer(
+        self,
+        wav: np.ndarray[typing.Any, np.dtype],
+        wav_sr: int,
+        model_sr: int = 40000,
+        sid: int = 0,
+        f0_method="dio",
+        f0_up_key=0,
+    ) -> np.ndarray[typing.Any, np.dtype[np.int16]]:
+        f0_min = 50
+        f0_max = 1100
+        f0_mel_min = 1127 * np.log(1 + f0_min / 700)
+        f0_mel_max = 1127 * np.log(1 + f0_max / 700)
+        f0_predictor = get_f0_predictor(
+            f0_method,
+            self.hop_len,
+            model_sr,
+        )
+        org_length = len(wav)
+        if org_length / wav_sr > 50.0:
+            raise RuntimeError("wav max length exceeded")
+
+        hubert = self.vec_model(librosa.resample(wav, orig_sr=wav_sr, target_sr=16000))
+        hubert = np.repeat(hubert, 2, axis=2).transpose(0, 2, 1).astype(np.float32)
+        hubert_length = hubert.shape[1]
+
+        pitchf = f0_predictor.compute_f0(wav, hubert_length)
+        pitchf = pitchf * 2 ** (f0_up_key / 12)
+        pitch = pitchf.copy()
+        f0_mel = 1127 * np.log(1 + pitch / 700)
+        f0_mel[f0_mel > 0] = (f0_mel[f0_mel > 0] - f0_mel_min) * 254 / (
+            f0_mel_max - f0_mel_min
+        ) + 1
+        f0_mel[f0_mel <= 1] = 1
+        f0_mel[f0_mel > 255] = 255
+        pitch = np.rint(f0_mel).astype(np.int64)
+
+        pitchf = pitchf.reshape(1, len(pitchf)).astype(np.float32)
+        pitch = pitch.reshape(1, len(pitch))
+        ds = np.array([sid]).astype(np.int64)
+
+        rnd = np.random.randn(1, 192, hubert_length).astype(np.float32)
+        hubert_length = np.array([hubert_length]).astype(np.int64)
+
+        out_wav = self.forward(hubert, hubert_length, pitch, pitchf, ds, rnd).squeeze()
+
+        out_wav = np.pad(out_wav, (0, 2 * self.hop_len), "constant")
+
+        return out_wav[0:org_length]
+
+    def forward(
+        self,
+        hubert: np.ndarray[typing.Any, np.dtype[np.float32]],
+        hubert_length: int,
+        pitch: np.ndarray[typing.Any, np.dtype[np.int64]],
+        pitchf: np.ndarray[typing.Any, np.dtype[np.float32]],
+        ds: np.ndarray[typing.Any, np.dtype[np.int64]],
+        rnd: np.ndarray[typing.Any, np.dtype[np.float32]],
+    ) -> np.ndarray[typing.Any, np.dtype[np.int16]]:
+        onnx_input = {
+            self.model.get_inputs()[0].name: hubert,
+            self.model.get_inputs()[1].name: hubert_length,
+            self.model.get_inputs()[2].name: pitch,
+            self.model.get_inputs()[3].name: pitchf,
+            self.model.get_inputs()[4].name: ds,
+            self.model.get_inputs()[5].name: rnd,
+        }
+        return (self.model.run(None, onnx_input)[0] * 32767).astype(np.int16)

rvc/onnx/synthesizer.py

@@ -0,0 +1,80 @@
+from typing import List, Optional, Union
+
+import torch
+
+from rvc.layers.synthesizers import SynthesizerTrnMsNSFsid as SynthesizerBase
+
+
+class SynthesizerTrnMsNSFsid(SynthesizerBase):
+    def __init__(
+        self,
+        spec_channels: int,
+        segment_size: int,
+        inter_channels: int,
+        hidden_channels: int,
+        filter_channels: int,
+        n_heads: int,
+        n_layers: int,
+        kernel_size: int,
+        p_dropout: int,
+        resblock: str,
+        resblock_kernel_sizes: List[int],
+        resblock_dilation_sizes: List[List[int]],
+        upsample_rates: List[int],
+        upsample_initial_channel: int,
+        upsample_kernel_sizes: List[int],
+        spk_embed_dim: int,
+        gin_channels: int,
+        sr: Optional[Union[str, int]],
+        encoder_dim: int,
+    ):
+        super().__init__(
+            spec_channels,
+            segment_size,
+            inter_channels,
+            hidden_channels,
+            filter_channels,
+            n_heads,
+            n_layers,
+            kernel_size,
+            p_dropout,
+            resblock,
+            resblock_kernel_sizes,
+            resblock_dilation_sizes,
+            upsample_rates,
+            upsample_initial_channel,
+            upsample_kernel_sizes,
+            spk_embed_dim,
+            gin_channels,
+            sr,
+            encoder_dim,
+            True,
+        )
+        self.speaker_map = None
+
+    def remove_weight_norm(self):
+        self.dec.remove_weight_norm()
+        self.flow.remove_weight_norm()
+        self.enc_q.remove_weight_norm()
+
+    def construct_spkmixmap(self):
+        self.speaker_map = torch.zeros((self.n_speaker, 1, 1, self.gin_channels))
+        for i in range(self.n_speaker):
+            self.speaker_map[i] = self.emb_g(torch.LongTensor([[i]]))
+        self.speaker_map = self.speaker_map.unsqueeze(0)
+
+    def forward(self, phone, phone_lengths, pitch, nsff0, g, rnd, max_len=None):
+        if self.speaker_map is not None:  # [N, S]  *  [S, B, 1, H]
+            g = g.reshape((g.shape[0], g.shape[1], 1, 1, 1))  # [N, S, B, 1, 1]
+            g = g * self.speaker_map  # [N, S, B, 1, H]
+            g = torch.sum(g, dim=1)  # [N, 1, B, 1, H]
+            g = g.transpose(0, -1).transpose(0, -2).squeeze(0)  # [B, H, N]
+        else:
+            g = g.unsqueeze(0)
+            g = self.emb_g(g).transpose(1, 2)
+
+        m_p, logs_p, x_mask = self.enc_p(phone, pitch, phone_lengths)
+        z_p = (m_p + torch.exp(logs_p) * rnd) * x_mask
+        z = self.flow(z_p, x_mask, g=g, reverse=True)
+        o = self.dec((z * x_mask)[:, :, :max_len], nsff0, g=g)
+        return o

rvc/synthesizer.py

@@ -0,0 +1,65 @@
+from collections import OrderedDict
+
+import torch
+
+from .layers.synthesizers import SynthesizerTrnMsNSFsid
+from .jit import load_inputs, export_jit_model, save_pickle
+
+
+def get_synthesizer(cpt: OrderedDict, device=torch.device("cpu")):
+    cpt["config"][-3] = cpt["weight"]["emb_g.weight"].shape[0]
+    if_f0 = cpt.get("f0", 1)
+    version = cpt.get("version", "v1")
+    if version == "v1":
+        encoder_dim = 256
+    elif version == "v2":
+        encoder_dim = 768
+    net_g = SynthesizerTrnMsNSFsid(
+        *cpt["config"],
+        encoder_dim=encoder_dim,
+        use_f0=if_f0 == 1,
+    )
+    del net_g.enc_q
+    net_g.load_state_dict(cpt["weight"], strict=False)
+    net_g = net_g.float()
+    net_g.eval().to(device)
+    net_g.remove_weight_norm()
+    return net_g, cpt
+
+
+def load_synthesizer(
+    pth_path: torch.serialization.FILE_LIKE, device=torch.device("cpu")
+):
+    return get_synthesizer(
+        torch.load(pth_path, map_location=torch.device("cpu")),
+        device,
+    )
+
+
+def synthesizer_jit_export(
+    model_path: str,
+    mode: str = "script",
+    inputs_path: str = None,
+    save_path: str = None,
+    device=torch.device("cpu"),
+    is_half=False,
+):
+    if not save_path:
+        save_path = model_path.rstrip(".pth")
+        save_path += ".half.jit" if is_half else ".jit"
+    if "cuda" in str(device) and ":" not in str(device):
+        device = torch.device("cuda:0")
+    from rvc.synthesizer import load_synthesizer
+
+    model, cpt = load_synthesizer(model_path, device)
+    assert isinstance(cpt, dict)
+    model.forward = model.infer
+    inputs = None
+    if mode == "trace":
+        inputs = load_inputs(inputs_path, device, is_half)
+    ckpt = export_jit_model(model, mode, inputs, device, is_half)
+    cpt.pop("weight")
+    cpt["model"] = ckpt["model"]
+    cpt["device"] = device
+    save_pickle(cpt, save_path)
+    return cpt