RvcPyInfer 0.1.0__py3-none-any.whl

RvcPyInfer/InferProviders.py

@@ -0,0 +1,108 @@
+from enum import Flag, auto
+
+from .error.InferEnvError import InferEnvError
+
+
+class InferProviders(Flag):
+    NONE = auto() # 仅用于内部方便处理
+
+    # ort
+    ORT_CPU = auto()
+    FORCE_CUDA = auto()
+    FORCE_DML = auto()
+    CUDA = ORT_CPU | FORCE_CUDA
+    DML = ORT_CPU | FORCE_DML
+
+    # openvino
+    OPENVINO_AUTO = auto()
+    OPENVINO_CPU = auto()
+    OPENVINO_GPU = auto()
+    OPENVINO_NPU = auto()
+    OPENVINO_OPTIONAL_MULTI  = auto()
+    OPENVINO_OPTIONAL_HETERO = auto()
+    OPENVINO_ALL = OPENVINO_CPU | OPENVINO_GPU | OPENVINO_NPU | OPENVINO_OPTIONAL_MULTI
+
+    def is_ort(self) -> bool:
+        return bool((InferProviders.ORT_CPU | InferProviders.CUDA | InferProviders.DML) & self)
+    def is_ov(self) -> bool:
+        return bool((InferProviders.OPENVINO_CPU | InferProviders.OPENVINO_GPU | InferProviders.OPENVINO_NPU | InferProviders.OPENVINO_AUTO) & self)
+
+    def get_onnx_provider(self) -> list[str]:
+        res = []
+        if self.FORCE_CUDA in self:
+            res.append("CUDAExecutionProvider")
+        if self.FORCE_DML in self:
+            res.append("DmlExecutionProvider")
+        if self.ORT_CPU in self:
+            res.append("CPUExecutionProvider")
+        return res
+    
+    def get_openvino_device(self) -> str:
+        devices: list[str] = []
+        if self.OPENVINO_GPU in self:
+            # 现场找全
+            from openvino import Core  # pyright: ignore[reportMissingImports]
+            devices.extend(filter(lambda name: name.startswith("GPU"), Core().available_devices))
+        if self.OPENVINO_NPU in self:
+            devices.append("NPU")
+        if self.OPENVINO_CPU in self:
+            devices.append("CPU")
+        devices_str = ",".join(devices)
+        if self.OPENVINO_AUTO in self:
+            if len(devices) != 0:
+                return f"AUTO:{devices_str}"
+            else:
+                return "AUTO"
+        elif self.OPENVINO_OPTIONAL_MULTI in self:
+            return f"MULTI:{devices_str}"
+        elif self.OPENVINO_OPTIONAL_HETERO in self:
+            return f"HETERO:{devices_str}"
+        else:
+            return devices[0]
+
+    @staticmethod
+    def default() -> "InferProviders":
+        available = InferProviders.currently_available()
+
+        if InferProviders.FORCE_CUDA in available:
+            return InferProviders.CUDA
+        elif InferProviders.FORCE_DML in available:
+            return InferProviders.DML
+        elif InferProviders.OPENVINO_CPU in available:
+            # 有时候自动的异构并行反而比纯 CPU 还慢，这里保守点
+            return InferProviders.OPENVINO_CPU
+        elif InferProviders.ORT_CPU in available:
+            return InferProviders.ORT_CPU
+        else:
+            raise InferEnvError("没有任意一个支持的推理引擎，请至少安装一个可选模块")
+        
+    @staticmethod
+    def currently_available() -> "InferProviders":
+        from .infer_env import HAS_ONE, HAS_OPENVINO, HAS_ORT
+        res = InferProviders.NONE
+
+        if HAS_ORT:
+            import onnxruntime as ort  # pyright: ignore[reportMissingImports]
+            available_providers = ort.get_available_providers()
+            if "CPUExecutionProvider" in available_providers:
+                res = res | InferProviders.ORT_CPU
+            if "CUDAExecutionProvider" in available_providers:
+                res = res | InferProviders.FORCE_CUDA
+            if "DmlExecutionProvider" in available_providers:
+                res = res | InferProviders.FORCE_DML
+        if HAS_OPENVINO:
+            from openvino import Core  # pyright: ignore[reportMissingImports]
+            devices = Core().available_devices
+            res = res | InferProviders.OPENVINO_AUTO | InferProviders.OPENVINO_OPTIONAL_MULTI | InferProviders.OPENVINO_OPTIONAL_HETERO
+            if "CPU" in devices:
+                res = res | InferProviders.OPENVINO_CPU
+            if any(i.startswith("GPU") for i in devices):
+                res = res | InferProviders.OPENVINO_GPU
+            if "NPU" in devices:
+                res = res | InferProviders.OPENVINO_NPU
+        if not HAS_ONE: 
+            raise InferEnvError("没有任意一个支持的推理引擎，请至少安装一个可选模块")
+        
+        assert res != InferProviders.NONE, "我算是知道了，你们各个都身怀绝技"
+        res = res & ~InferProviders.NONE
+        return res

RvcPyInfer/InferTask.py

@@ -0,0 +1,210 @@
+from collections.abc import Callable
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import numpy as np
+import soundfile
+from numpy.typing import NDArray
+
+from .audio.audio_utils import crossfade, reSR, rms_frame_match, split_by_max_len_with_overlap, split_by_silence
+from .f0_utils import apply_rise_tone, build_f0extract_func, f0_to_mel, normalized_mel
+from .type_alist import Audio, AudioLike, F0ExtractAlgorithm, F0ExtractAlgorithmList, FileLike
+
+if TYPE_CHECKING:
+    from .RvcContext import RvcContext
+
+class InferTask:
+    audios: list[Audio]
+    def __init__(self, 
+            # -- 必填 --
+            context: "RvcContext",
+            vec_model: Path,
+            gen_model: tuple[Path, int],
+            *audios: AudioLike,
+
+            # 可选扩展
+            sid: int = 0,
+            seed: int | None = 1234,
+
+            # -- 特征索引 --
+            index_path: Path | None = None,
+            index_rate: float = 0.33,
+            index_k: int = 8, # 检索的最近特征数量
+
+            # -- f0 提取 --
+            f0extract_algorithm: F0ExtractAlgorithm,
+            f0_up_semitone: float = 0,
+            f0_min: float = 50,
+            f0_max: float = 1100,
+
+            # -- 强制切片与交叉淡化 --
+            slice_max_len: int = 30,
+            slice_overlap_len: int = 5,
+
+            # -- 静音切片 --
+            silence_frame_len: int = 20,
+            silence_hop_len: int = 10,
+            silence_thresh_db: float = -40,
+            silence_min_silence_duration_ms: float = 800.0,
+            silence_max_transition_ms: float = 100.0,
+
+            # -- RMS 包络匹配 --
+            rms_match_frame_len: int = 20,
+            rms_match_hop_len: int = 10,
+            rms_match_mix: float = 1.0, # 一般不用改
+            ) -> None:
+        self.context = context
+        self.vec = vec_model.resolve()
+        self.gen = (gen_model[0].resolve(), gen_model[1])
+        self.audio_likes = list(audios)
+        self.sid = sid
+        self.seed = seed
+        self.index_path = index_path
+        self.index_rate = index_rate
+        self.index_k = index_k
+        self.f0extract_algorithm = f0extract_algorithm
+        self.f0_up_semitone = f0_up_semitone
+        self.f0_min = f0_min
+        self.f0_max = f0_max
+        self.slice_max_len = slice_max_len
+        self.slice_overlap_len = slice_overlap_len
+        self.silence_frame_len = silence_frame_len
+        self.silence_hop_len = silence_hop_len
+        self.silence_thresh_db = silence_thresh_db
+        self.silence_min_silence_duration_ms = silence_min_silence_duration_ms
+        self.silence_max_transition_ms = silence_max_transition_ms
+        self.rms_match_frame_len = rms_match_frame_len
+        self.rms_match_hop_len = rms_match_hop_len
+        self.rms_match_mix = rms_match_mix
+
+    def read(self) -> None:
+        self.audios = []
+        for audio_like in self.audio_likes:
+            if isinstance(audio_like, tuple) and len(audio_like) == 2 and isinstance(audio_like[0], np.ndarray) and isinstance(audio_like[1], int):
+                self.audios.append(audio_like)
+            else:
+                assert not isinstance(audio_like, tuple) # 这个类型检查真的很奇怪有没有
+                data, sr = soundfile.read(
+                    audio_like, dtype="float32"
+                )
+                data = data.astype(np.float32)
+                if len(data.shape) != 1:
+                    data = data.mean(axis=1, dtype=np.float32)
+                self.audios.append((data, sr))
+        del self.audio_likes
+
+    def core_chunk_infer(self, audio: Audio, f0extract_func: Callable[[Audio, int], NDArray[np.float32]]) -> Audio:
+        model = self.context._vec_pool.get(self.vec)
+        feats = model.infer(audio)
+
+        f0 = f0extract_func(audio, feats.shape[0]) # 这里是没有批处理维度的
+        f0 = apply_rise_tone(f0, self.f0_up_semitone)
+
+        mel = normalized_mel(f0_to_mel(f0), f0_to_mel(self.f0_min), f0_to_mel(self.f0_max))
+        mel = np.rint(mel).astype(np.int64)
+
+        if self.index_rate > 1e-6 and self.index_path is not None and self.context._index_pool is not None:
+            index = self.context._index_pool.get(self.index_path)
+            feats = index.apply_index(feats, self.index_rate, self.index_k)
+
+        model = self.context._gen_pool.get(self.gen)
+        res = model.infer(
+            phone=feats,
+            f0_pitch=f0, mel_pitch=mel,
+            sid=self.sid, seed=self.seed
+        )
+
+        res_d, res_sr = res
+        source_d, source_sr = audio
+
+        source_len = len(source_d)
+        res_len = source_len * res_sr // source_sr
+        
+        pad_len = res_len - len(res_d)
+        if pad_len == 0:
+            return res
+        elif pad_len < 0:
+            return (res_d[:res_len], res_sr)
+        else:
+            return (
+                np.pad(res_d, [0, pad_len], mode="constant"),
+                res_sr
+            )
+    
+    def chunk_infer(self, audio: Audio, f0extract_func: Callable[[Audio, int], NDArray[np.float32]]) -> Audio:
+        splited = split_by_max_len_with_overlap( # 就算不足一个切片长也不会报错哒
+            audio,
+            max_len=self.slice_max_len,
+            overlap_len=self.slice_overlap_len
+            )
+        del audio # 节约点内存
+        def handle(c):
+            return self.core_chunk_infer(c, f0extract_func)
+        infer = [handle(i) for i in splited]
+        res = crossfade(
+            infer,
+            overlap_len=self.slice_overlap_len
+            )
+        return res
+
+    def gen_infer(self, callback: Callable[[int, Audio], None]) -> None:
+        assert self.f0extract_algorithm in F0ExtractAlgorithmList, f"{self.f0extract_algorithm} 不是受支持的 f0 提取算法"
+        func = build_f0extract_func(self.f0extract_algorithm, self.f0_min, self.f0_max) # pyright: ignore[reportArgumentType]
+        for i, audio in enumerate(self.audios):
+            splited = split_by_silence(
+                audio,
+                frame_len=self.silence_frame_len,
+                hop_len=self.silence_hop_len,
+                silence_thresh_db=self.silence_thresh_db,
+                min_silence_duration_ms=self.silence_min_silence_duration_ms,
+                max_transition_ms=self.silence_max_transition_ms
+            )
+            chunks: list[Audio] = []
+            for chunk, is_sil in splited:
+                if is_sil:
+                    chunks.append(reSR(chunk, target_sr=self.gen[1])) # 内部判断相等会直接返回
+                else:
+                    chunks.append(self.chunk_infer(chunk, func))
+            datas = [data for data, sr in chunks]
+            res = (np.concatenate(datas), self.gen[1])
+            res = rms_frame_match(
+                source=res,
+                target=reSR(audio, target_sr=res[1]),
+                frame_len=self.rms_match_frame_len,
+                hop_len=self.rms_match_hop_len,
+                mix=self.rms_match_mix
+            )
+            callback(i, res)
+                
+    def run(self) -> list[Audio]:
+        self.read()
+
+        res = []
+        self.gen_infer(
+            lambda _, a: res.append(a)
+        )
+        return res
+    
+    def run_and_save(self, *output: FileLike, 
+                     subtype: str | None = None,
+                     format: str | None = None):
+        self.read()
+
+        if len(self.audios) != len(output):
+            raise ValueError(f"音频数量 {len(self.audios)} 与输出数量 {len(output)} 不匹配")
+        def save(i: int, a: Audio) -> None:
+            d, s = a
+            soundfile.write(output[i], 
+                            d, s, 
+                            subtype=subtype,
+                            format=format)
+        self.gen_infer(
+            save
+        )
+
+    def run_and_callback(self, 
+                        callback: Callable[[int, Audio], None]) -> None:
+        self.read()
+        self.gen_infer(
+            callback=callback
+        )

RvcPyInfer/RvcContext.py

@@ -0,0 +1,125 @@
+import warnings
+from pathlib import Path
+
+from .index.RvcFeatIndex import RvcFeatIndex
+from .InferProviders import InferProviders
+from .InferTask import InferTask
+from .onnx.ContentVec import ContentVec
+from .onnx.ModelSimplePool import ModelSimplePool
+from .onnx.RvcGen import RvcGen
+from .path_utils import path
+from .type_alist import AudioLike, F0ExtractAlgorithm, PathLike
+from .warn.InferEnvWarn import InferEnvWarn
+
+
+# 请注意，RvcContext 并不是线程安全的
+class RvcContext:
+    def __init__(self,
+                 providers: InferProviders = InferProviders.default(),
+                 vec_pool_permanent_size: int = 1,
+                 gen_pool_permanent_size: int = 2,
+                 index_pool_permanent_size: int = 2) -> None:
+        self._providers = providers
+        self._vec_pool = ModelSimplePool[Path, ContentVec](
+            lambda p: ContentVec(p.resolve(), self._providers),
+            vec_pool_permanent_size
+        )        
+        self._gen_pool = ModelSimplePool[tuple[Path, int], RvcGen](
+            lambda args: RvcGen(args[0].resolve(), args[1], self._providers),
+            gen_pool_permanent_size,
+        )
+
+        from .infer_env import HAS_FAISS
+        if HAS_FAISS:
+            self._index_pool = ModelSimplePool[Path, RvcFeatIndex](
+                lambda p: RvcFeatIndex(p.resolve()),
+                index_pool_permanent_size
+            )
+        else:
+            self._index_pool = None
+
+    def clear_pool(self) -> None:
+        self._vec_pool.clear()
+        self._gen_pool.clear()
+
+    def build_task(self, 
+            # -- 必填 --
+            vec_model: PathLike, 
+            gen_model: PathLike, 
+            gen_model_sr: int,
+            *audios: AudioLike,
+
+            # 可选扩展
+            sid: int = 0,
+            seed: int | None = 1234,
+
+            # -- 特征索引 --
+            index_path: PathLike | None = None,
+            index_rate: float = 0.33,
+            index_k: int = 8, # 检索的最近特征数量
+
+            # -- f0 提取 --
+            f0extract_algorithm: F0ExtractAlgorithm = "dio",
+            f0_up_semitone: float = 0,
+            f0_min: float = 50,
+            f0_max: float = 1100,
+
+            # -- 强制切片与交叉淡化 --
+            slice_max_len: int = 30,
+            slice_overlap_len: int = 5,
+
+            # -- 静音切片 --
+            silence_frame_len: int = 20,
+            silence_hop_len: int = 10,
+            silence_thresh_db: float = -40,
+            silence_min_silence_duration_ms: float = 800.0,
+            silence_max_transition_ms: float = 100.0,
+
+            # -- RMS 包络匹配 --
+            rms_match_frame_len: int = 20,
+            rms_match_hop_len: int = 10,
+            rms_match_mix: float = 1.0, # 一般不用改
+            ) -> "InferTask":
+        vec = path(vec_model)
+        gen = path(gen_model)
+
+        if not vec.exists():
+            raise FileNotFoundError(f"找不到模型: {vec}")
+        if not gen.exists():
+            raise FileNotFoundError(f"找不到模型: {gen}")
+        
+        if index_path is not None:
+            from .infer_env import HAS_FAISS
+            if not HAS_FAISS:
+                warnings.warn("您正在使用特征索引，但并未安装 faiss。如果您想使用特征索引请安装 [index] 可选依赖", InferEnvWarn)
+                index_path = None
+            else:
+                index_path = path(index_path)
+                if not index_path.exists():
+                    raise FileNotFoundError(f"找不到特征索引: {index_path}")
+
+        return InferTask(
+            self,
+            vec,
+            (gen, gen_model_sr),
+            *audios,
+            sid=sid,
+            seed=seed,
+            index_path=index_path,
+            index_rate=index_rate,
+            index_k=index_k,
+            f0extract_algorithm=f0extract_algorithm,
+            f0_up_semitone=f0_up_semitone,
+            f0_min=f0_min,
+            f0_max=f0_max,
+            slice_max_len=slice_max_len,
+            slice_overlap_len=slice_overlap_len,
+            silence_frame_len=silence_frame_len,
+            silence_hop_len=silence_hop_len,
+            silence_thresh_db=silence_thresh_db,
+            silence_min_silence_duration_ms=silence_min_silence_duration_ms,
+            silence_max_transition_ms=silence_max_transition_ms,
+            rms_match_frame_len=rms_match_frame_len,
+            rms_match_hop_len=rms_match_hop_len,
+            rms_match_mix=rms_match_mix,
+        )

RvcPyInfer/__init__.py

@@ -0,0 +1,42 @@
+from . import f0_utils
+from ._version import __version__  # noqa: F401
+from .audio import audio_utils
+from .cli import main as ModelInferCLI
+from .error.InferEnvError import InferEnvError
+from .error.NotSupportedAlgorithmError import NotSupportedAlgorithmError
+from .InferProviders import InferProviders
+from .InferTask import InferTask
+from .onnx.ContentVec import ContentVec
+from .onnx.export import Exporter as OnnxExporter
+from .onnx.export.cli import main as ModelExportToolCLI
+from .onnx.export.direct_read_sr import direct_read_sr
+from .onnx.ModelSimplePool import ModelSimplePool
+from .onnx.RvcGen import RvcGen
+from .RvcContext import RvcContext
+from .type_alist import Audio, AudioLike, F0ExtractAlgorithm, F0ExtractAlgorithmList, FileLike, PathLike
+from .warn.InferEnvWarn import InferEnvWarn
+from .warn.InferModelWarn import InferModelWarn
+from .warn.InferWarn import InferWarn
+
+# 用户肯定不会希望我的 __version__ 去污染 ta 的版本号
+
+__all__ = [
+    "Audio", "AudioLike", "PathLike", "FileLike", "F0ExtractAlgorithm", "F0ExtractAlgorithmList",
+    "InferEnvError", 
+    "NotSupportedAlgorithmError",
+    "InferWarn",
+    "InferEnvWarn",
+    "InferModelWarn",
+    "f0_utils",
+    "audio_utils",
+    "InferProviders",
+    "RvcContext",
+    "InferTask",
+    "OnnxExporter",
+    "direct_read_sr",
+    "ModelExportToolCLI",
+    "ContentVec",
+    "RvcGen",
+    "ModelSimplePool",
+    "ModelInferCLI"
+]

RvcPyInfer/__main__.py

@@ -0,0 +1,4 @@
+from .cli import main as infer_main
+
+if __name__ == "__main__":
+    infer_main()

RvcPyInfer/_version.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"