optimize(train): combine extract f0 together

rvc/f0/__init__.py

@@ -1,10 +1 @@
-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"]
+from .gen import Generator

rvc/f0/gen.py

@@ -0,0 +1,127 @@
+from math import log
+from pathlib import Path
+from typing import Optional, Union, Literal, Tuple
+
+from numba import jit
+import numpy as np
+
+
+@jit(nopython=True)
+def post_process(
+    sr: int,
+    window: int,
+    f0: np.ndarray,
+    f0_up_key: int,
+    manual_x_pad: int,
+    f0_mel_min: float,
+    f0_mel_max: float,
+    manual_f0: Optional[Union[np.ndarray, list]]=None,
+) -> Tuple[np.ndarray, np.ndarray]:
+    f0 = np.multiply(f0, pow(2, f0_up_key / 12))
+    # with open("test.txt","w")as f:f.write("\n".join([str(i)for i in f0.tolist()]))
+    tf0 = sr // window  # 每秒f0点数
+    if manual_f0 is not None:
+        delta_t = np.round(
+            (manual_f0[:, 0].max() - manual_f0[:, 0].min()) * tf0 + 1
+        ).astype("int16")
+        replace_f0 = np.interp(
+            list(range(delta_t)), manual_f0[:, 0] * 100, manual_f0[:, 1]
+        )
+        shape = f0[manual_x_pad * tf0 : manual_x_pad * tf0 + len(replace_f0)].shape[0]
+        f0[manual_x_pad * tf0 : manual_x_pad * tf0 + len(replace_f0)] = replace_f0[:shape]
+    # with open("test_opt.txt","w")as f:f.write("\n".join([str(i)for i in f0.tolist()]))
+    f0_mel = 1127 * np.log(1 + f0 / 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
+    f0_coarse = np.rint(f0_mel).astype(np.int32)
+    return f0_coarse, f0  # 1-0
+
+
+class Generator(object):
+    def __init__(
+        self,
+        rmvpe_root: Path,
+        is_half: bool,
+        x_pad: int,
+        device = "cpu",
+        window = 160,
+        sr = 16000
+    ):
+        self.rmvpe_root = rmvpe_root
+        self.is_half = is_half
+        self.x_pad = x_pad
+        self.device = device
+        self.window = window
+        self.sr = sr
+
+    def calculate(
+        self,
+        x: np.ndarray,
+        p_len: int,
+        f0_up_key: int,
+        f0_method: Literal['pm', 'dio', 'harvest', 'crepe', 'rmvpe', 'fcpe'],
+        filter_radius: Optional[Union[int, float]],
+        manual_f0: Optional[Union[np.ndarray, list]]=None,
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        f0_min = 50
+        f0_max = 1100
+        if f0_method == "pm":
+            if not hasattr(self, "pm"):
+                from .pm import PM
+                self.pm = PM(self.window, f0_min, f0_max, self.sr)
+            f0 = self.pm.compute_f0(x, p_len=p_len)
+        elif f0_method == "dio":
+            if not hasattr(self, "dio"):
+                from .dio import Dio
+                self.dio = Dio(self.window, f0_min, f0_max, self.sr)
+            f0 = self.dio.compute_f0(x, p_len=p_len)
+        elif f0_method == "harvest":
+            if not hasattr(self, "harvest"):
+                from .harvest import Harvest
+                self.harvest = Harvest(self.window, f0_min, f0_max, self.sr)
+            f0 = self.harvest.compute_f0(x, p_len=p_len, filter_radius=filter_radius)
+        elif f0_method == "crepe":
+            if not hasattr(self, "crepe"):
+                from .crepe import CRePE
+                self.crepe = CRePE(
+                    self.window,
+                    f0_min,
+                    f0_max,
+                    self.sr,
+                    self.device,
+                )
+            f0 = self.crepe.compute_f0(x, p_len=p_len)
+        elif f0_method == "rmvpe":
+            if not hasattr(self, "rmvpe"):
+                from .rmvpe import RMVPE
+                self.rmvpe = RMVPE(
+                    str(self.rmvpe_root/"rmvpe.pt"),
+                    is_half=self.is_half,
+                    device=self.device,
+                    # use_jit=self.config.use_jit,
+                )
+            f0 = self.rmvpe.compute_f0(x, p_len=p_len, filter_radius=0.03)
+            if "privateuseone" in str(self.device):  # clean ortruntime memory
+                del self.rmvpe.model
+                del self.rmvpe
+        elif f0_method == "fcpe":
+            if not hasattr(self, "fcpe"):
+                from .fcpe import FCPE
+                self.fcpe = FCPE(
+                    self.window,
+                    f0_min,
+                    f0_max,
+                    self.sr,
+                    self.device,
+                )
+            f0 = self.fcpe.compute_f0(x, p_len=p_len)
+        else:
+            raise ValueError(f"f0 method {f0_method} has not yet been supported")
+
+        return post_process(
+            self.sr, self.window, f0, f0_up_key, self.x_pad,
+            1127 * log(1 + f0_min / 700),
+            1127 * log(1 + f0_max / 700),
+            manual_f0,
+        )

rvc/f0/pm.py

@@ -1,4 +1,4 @@
-from typing import Any, Optional
+from typing import Optional
 
 import numpy as np
 import parselmouth

rvc/onnx/infer.py

@@ -5,12 +5,7 @@ import librosa
 import numpy as np
 import onnxruntime
 
-from rvc.f0 import (
-    PM,
-    Harvest,
-    Dio,
-    F0Predictor,
-)
+from rvc.f0 import Generator
 
 
 class Model:
@@ -51,49 +46,28 @@ class ContentVec(Model):
         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,
+        model_sr=40000,
         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
+        self.f0_gen = Generator(None, False, 0, window=hop_len, sr=model_sr)
 
     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")
@@ -102,16 +76,8 @@ class RVC(Model):
         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)
+        pitch, pitchf = self.f0_gen.calculate(wav, hubert_length, f0_up_key, f0_method, None)
+        pitch = pitch.astype(np.int64)
 
         pitchf = pitchf.reshape(1, len(pitchf)).astype(np.float32)
         pitch = pitch.reshape(1, len(pitch))