RvcPyInfer 0.1.0__py3-none-any.whl

RvcPyInfer/onnx/ModelSimplePool.py

@@ -0,0 +1,26 @@
+from collections import OrderedDict
+from collections.abc import Callable
+
+
+class ModelSimplePool[TKey, TModel]:
+    def __init__(self, factory: Callable[[TKey], TModel], permanent_size: int = 1) -> None:
+        self._factory = factory
+        self._pools = OrderedDict[TKey, TModel]()
+        self._permanent_size = permanent_size
+
+    def remove(self, key: TKey) -> None:
+        del self._pools[key]
+
+    def clear(self) -> None:
+        self._pools.clear()
+
+    def get(self, key: TKey) -> TModel:
+        model = self._pools.get(key)
+        if model is not None:
+            self._pools.move_to_end(key)
+        else:
+            model = self._factory(key)
+            self._pools[key] = model
+            if len(self._pools) > self._permanent_size:
+                self._pools.popitem(last=False)
+        return model

RvcPyInfer/onnx/RvcGen.py

@@ -0,0 +1,51 @@
+import numpy as np
+from numpy.typing import NDArray
+
+from ..error.InferEnvError import InferEnvError
+from ..InferProviders import InferProviders
+from ..path_utils import path as pathf
+from ..type_alist import Audio, PathLike
+from .model_loader import load_model
+
+
+class RvcGen:
+    def __init__(self, path: PathLike, model_out_sr: int, providers: InferProviders) -> None:
+        self.sr = model_out_sr
+        self.session, self.compiled_model = load_model(pathf(path), providers)
+        if self.session is not None:
+            self.channels = next(filter(lambda arg: arg.name == "rnd", self.session.get_inputs())).shape[1]
+        elif self.compiled_model is not None:
+            self.channels = next(filter(lambda arg: arg.any_name == "rnd", self.compiled_model.inputs)).partial_shape[1].get_length()
+
+    def infer(self, phone: NDArray[np.float32], mel_pitch: NDArray[np.int64], f0_pitch: NDArray[np.float32], sid: int = 0, seed: int | None = 1234) -> Audio:
+        assert phone.shape[0] == mel_pitch.shape[0] == f0_pitch.shape[0], "帧数应统一"
+        phone_len = phone.shape[0]
+        ds = np.array([sid], dtype=np.int64)
+        rnd = np.random.default_rng(seed).standard_normal((1, self.channels, phone_len), dtype=np.float32)
+        phone = phone[None, ...]
+        mel_pitch = mel_pitch.reshape(1, -1)
+        f0_pitch = f0_pitch.reshape(1, -1)
+
+        model_input = {
+            "rnd": rnd,
+            "ds": ds,
+            "phone": phone,
+            "phone_lengths": np.array([phone_len]),
+            "pitch": mel_pitch,
+            "pitchf": f0_pitch
+        }
+        if self.session is not None:
+            output = self.session.run(
+                output_names=["audio"],
+                input_feed=model_input
+            )[0]
+        elif self.compiled_model is not None:
+            infer_request = self.compiled_model.create_infer_request()
+            result_dict = infer_request.infer(inputs=model_input)
+            output = result_dict["audio"]
+            del infer_request, result_dict
+        else:
+            raise InferEnvError("没有任意一个支持的推理引擎，请至少安装一个可选模块")
+        # (1, sampling)
+        assert isinstance(output, np.ndarray), "输出应为 NDArray"
+        return output.squeeze(axis=1).squeeze(axis=0).astype(np.float32), self.sr

RvcPyInfer/onnx/export/Exporter.py

@@ -0,0 +1,130 @@
+import hashlib
+import json
+import os
+import random
+import re
+import subprocess
+import uuid
+from datetime import datetime
+from importlib.resources import files
+
+from ...path_utils import path
+from ...type_alist import PathLike
+
+const_tips = """# ----------------------------------------------------------------
+# 这是自动生成的临时脚本，若您看见工具未能正常删除的话，您可以手动删除它
+# 您当然可以尝试运行了，甚至保存下来，脚本引用的是绝对路径
+# ----------------------------------------------------------------
+
+"""
+
+class Exporter:
+    def __init__(self, root: PathLike, model: PathLike, export_target: PathLike, runtime: PathLike | None = None) -> None:
+        self.root = path(root).resolve()
+        self.model = path(model).resolve()
+        self.target = path(export_target).resolve()
+        if runtime is None:
+            self.runtime = self.root / "runtime" / "python.exe"
+        else:
+            self.runtime = path(runtime).resolve()
+        self.parse_prefix = self.build_parse_prefix()
+        template = files("RvcPyInfer.resource.template").joinpath("export_onnx.py").read_text(encoding="utf-8")
+        self.template_digest = hashlib.sha256(template.encode("utf-8"))
+        self.script = self.format_template(template)
+
+    @staticmethod
+    def build_parse_prefix() -> str:
+        alphabet = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z']
+        tmp: list[str] = []
+        for _ in range(8):
+            tmp.append(alphabet[random.randint(0, len(alphabet) - 1)])
+        return "".join(tmp)
+
+    def format_template(self, template: str) -> str:
+        model = str(self.model).replace("\\", "/")
+        target = str(self.target).replace("\\", "/")
+        root = str(self.root).replace("\\", "/")
+        template = template.replace("/./..{code.InsertionPoint.0}.././", f'{model}')
+        template = template.replace("/./..{code.InsertionPoint.1}.././", target)
+        template = template.replace("/./..{code.InsertionPoint.2}.././", self.parse_prefix)
+        template = template.replace("/./..{code.InsertionPoint.3}.././", root)
+
+        args_digest = hashlib.sha256(f"model: {model}, target: {target}, root: {root}".encode())
+
+        gen_info = [
+            "\n"
+            "# ----------------------------------------------------------------\n"
+            "# -------- 脚本元信息 --------\n"
+            f"# 该脚本生成于: {datetime.now()}\n",
+            f"# 模板摘要: {self.template_digest.hexdigest()}\n"
+            f"# 非随机参数摘要: {args_digest.hexdigest()}\n"
+            f"# 模型: {model}\n"
+            f"# 目标: {target}\n"
+            f"# Rvc 项目根目录: {root}\n"
+            f"# 解析前缀: {self.parse_prefix}\n"
+            f"# 系统执行所用的 runtime: {self.runtime}\n"
+            "# ----------------------------------------------------------------\n"
+            "\n"
+        ]
+
+        gen_info = ''.join(gen_info)
+
+        return const_tips + gen_info + template
+    
+    def parse_stdout(self, stdout: str) -> dict[str, str | int]:
+        lines = stdout.split('\n')
+        info = {}
+        for line in lines:
+            line = line.strip()
+            if line.startswith(self.parse_prefix):
+                match = re.match(r'^\/\.\/\.\.(.*)\.\.\/\.\/$', line[len(self.parse_prefix):])
+                if match is None:
+                    continue
+                else:
+                    json_info = match.group(1)
+                    info |= json.loads(json_info)
+        return info
+    
+    def export(self) -> None:
+        tmp_py_path = (self.root / str(uuid.uuid4())).with_suffix(".py")
+        with open(tmp_py_path, mode="+w", encoding="utf-8") as f:
+            f.write(self.script)
+        try:
+            result = subprocess.run(
+                args=[str(self.runtime), str(tmp_py_path)],
+                capture_output=True,
+                text=True,
+                timeout=120,
+                cwd=self.root
+            )
+        finally:
+            try:
+                if tmp_py_path.exists():
+                    os.remove(tmp_py_path)
+            except Exception:
+                ...
+        info = self.parse_stdout(result.stdout)
+        sr = info.get("sr")
+        if sr is not None:
+            print(f"模型采样率: {sr}")
+        
+        if result.returncode == 0:
+            print("导出成功！")
+            return
+        else:
+            msg = info.get("msg")
+            if msg is None:
+                print("导出失败")
+            else:
+                if isinstance(msg, str):
+                    msg = msg.replace("\\\"", "\"").replace("\\n", "\n")
+                print(f"导出失败: {msg}")
+            print("-------- stdout --------")
+            print(result.stdout)
+            print("------------------------")
+            print()
+            print("-------- stderr --------")
+            print(result.stderr)
+            print("------------------------")
+            print()
+            print(f"return code: {result.returncode}")

RvcPyInfer/onnx/export/cli.py

@@ -0,0 +1,104 @@
+import argparse
+import sys
+from pathlib import Path
+
+from .direct_read_sr import direct_read_sr
+from .Exporter import Exporter
+
+
+def export_command(args) -> None:
+    """处理 export 子命令的具体逻辑"""
+    root_path = Path(args.root).resolve()
+    model_path = Path(args.model).resolve()
+    target_path = Path(args.target).resolve()
+    runtime_path = Path(args.runtime).resolve() if args.runtime else None
+
+    # 前置检查
+    if not root_path.is_dir():
+        print(f"[错误] 根目录不存在或不是文件夹: {root_path}")
+        sys.exit(1)
+        
+    if not model_path.is_file():
+        print(f"[错误] 找不到模型文件: {model_path}")
+        sys.exit(1)
+
+    # 确保输出目录存在
+    if not target_path.parent.exists():
+        print(f"[提示] 创建输出目录: {target_path.parent}")
+        target_path.parent.mkdir(parents=True, exist_ok=True)
+
+    # 执行导出
+    try:
+        print(">>> 开始初始化导出器...")
+        exporter = Exporter(
+            root=root_path,
+            model=model_path,
+            export_target=target_path,
+            runtime=runtime_path
+        )
+        
+        print(">>> 正在执行导出流程，请耐心等待...")
+        exporter.export()
+        
+    except FileNotFoundError as e:
+        print(f"\n[错误] 找不到必要的文件: {e}")
+        print("请检查模板文件 (export_onnx.py) 是否存在于 resource/template 目录下，或运行时路径是否正确。")
+        sys.exit(1)
+    except Exception as e:
+        print(f"\n[错误] 导出过程中发生未预期的异常: {e}")
+        import traceback
+        traceback.print_exc()
+        sys.exit(1)
+
+def show_sr_command(args) -> None:
+    """处理 show-sr 子命令的具体逻辑"""
+    model_path = Path(args.model).resolve()
+
+    if not model_path.is_file():
+        print(f"[错误] 找不到模型文件: {model_path}")
+        sys.exit(1)
+
+    sr = direct_read_sr(model_path)
+    print(f"模型采样率是: {sr}")
+
+def main() -> None:
+    # 主解析器
+    parser = argparse.ArgumentParser(
+        description="RvcPyInfer 命令行工具",
+        formatter_class=argparse.RawTextHelpFormatter
+    )
+    
+    # 创建子命令解析器
+    subparsers = parser.add_subparsers(dest="command", help="可用的子命令", required=True)
+
+    # 注册 export 子命令
+    export_parser = subparsers.add_parser(
+        "export", 
+        help="导出 RVC 模型为 ONNX 格式",
+        description="将 RVC .pth 模型导出为 ONNX 格式"
+    )
+    # export 子命令的参数
+    export_parser.add_argument("-r", "--root", type=str, required=True, help="RVC 项目根目录路径")
+    export_parser.add_argument("-m", "--model", type=str, required=True, help="输入的 RVC 模型路径")
+    export_parser.add_argument("-t", "--target", type=str, required=True, help="导出的 ONNX 模型目标路径")
+    export_parser.add_argument("--runtime", type=str, default=None, help="自定义 Python 运行时路径 (默认: <root>/runtime/python.exe)")
+    # 绑定执行函数
+    export_parser.set_defaults(func=export_command)
+
+    # 注册 show-sr 子命令
+    show_sr_parser = subparsers.add_parser(
+        "show-sr",
+        help="直接查看 RVC 模型的输出采样率",
+        description="查看 RVC 模型的输出采样率"
+    )
+    show_sr_parser.add_argument("-m", "--model", type=str, required=True, help="输入的 RVC 模型路径")
+    show_sr_parser.set_defaults(func=show_sr_command)
+
+    # 解析参数
+    args = parser.parse_args()
+
+    # 调用对应的处理函数
+    args.func(args)
+
+if __name__ == "__main__":
+    main()

RvcPyInfer/onnx/export/direct_read_sr.py

@@ -0,0 +1,109 @@
+import collections
+import io
+import pickle
+import zipfile
+from typing import Any
+
+from ...type_alist import PathLike
+
+
+class DroppedObject:
+    """
+    哨兵值：静默吸收所有参数和状态，代表被丢弃的不安全对象。
+    采用单例模式，所有被丢弃的对象都将指向这同一个实例。
+    """
+    _instance = None
+
+    def __new__(cls, *args: Any, **kwargs: Any) -> "DroppedObject":
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        # 吸收所有 pickle 传入的构造参数
+        pass
+    
+    def __setstate__(self, state: Any) -> None:
+        # 吸收所有 pickle 传入的属性状态
+        pass
+
+    def __repr__(self) -> str:
+        return "<DroppedObject>"
+
+    def __bool__(self) -> bool:
+        # 方便在 if 判断中识别：if val is DROPPED: ...
+        return False
+    
+class SafeStandardUnpickler(pickle.Unpickler):
+    """遇到非白名单类型时，静默替换为 DroppedObject 哨兵"""
+    
+    SAFE_CLASSES = {
+        'builtins': {
+            v.__name__: v for v in [
+                type(None), type(...), type(NotImplemented),
+                int, float, complex, bool, str, bytes, bytearray, memoryview,
+                list, tuple, dict, set, frozenset, range, slice, object
+            ]
+        },
+        'collections': {
+            'OrderedDict': collections.OrderedDict,
+            'defaultdict': collections.defaultdict,
+            'deque': collections.deque,
+            'Counter': collections.Counter,
+        },
+    }
+
+    def find_class(self, module: str, name: str):
+        if module in self.SAFE_CLASSES and name in self.SAFE_CLASSES[module]:
+            return self.SAFE_CLASSES[module][name]
+        
+        # 静默丢弃：返回哨兵类，不抛出异常
+        return DroppedObject
+    
+    def persistent_load(self, pid: Any) -> DroppedObject:
+        """
+        PyTorch 用 persistent_id 机制存储 tensor 的 storage。
+        我们不需要 tensor 数据，直接返回哨兵值。
+        """
+        return DroppedObject()
+    
+def safe_load_pth(filepath: PathLike) -> Any:
+    """
+    安全加载 PyTorch .pth/.pt 文件。
+    自动处理 ZIP 格式（PyTorch 1.6+）和旧版原始 pickle 格式。
+    """
+    with open(filepath, 'rb') as f:
+        magic = f.read(2)
+        f.seek(0)
+
+        if magic == b'PK':
+            # ===== ZIP 格式（PyTorch 1.6+）=====
+            # .pth 文件本质是一个 ZIP，里面包含:
+            #   archive/data.pkl  ← 结构信息（我们要的）
+            #   archive/data/0    ← tensor 原始数据（不需要）
+            with zipfile.ZipFile(filepath, 'r') as zf:
+                # 找到 .pkl 文件
+                pkl_names = [n for n in zf.namelist()
+                             if n.endswith('.pkl')]
+                if not pkl_names:
+                    raise ValueError(
+                        f"ZIP 格式的 .pth 文件中未找到 .pkl 文件，"
+                        f"包含的文件: {zf.namelist()}"
+                    )
+                pkl_data = zf.read(pkl_names[0])
+                return SafeStandardUnpickler(io.BytesIO(pkl_data)).load()
+        else:
+            # ===== 旧版原始 pickle 格式 =====
+            return SafeStandardUnpickler(f).load()
+
+def direct_read_sr(model: PathLike) -> int:
+    loaded_data = safe_load_pth(model)
+    config = loaded_data.get("config")
+    if config is None:
+        raise ValueError("找不到 config")
+    if len(config) < 18: 
+        raise ValueError("config 长度不足 18")
+    sr = config[17]
+    if not isinstance(sr, int):
+        raise ValueError("采样率不是整数")
+    return sr

RvcPyInfer/onnx/model_loader.py

@@ -0,0 +1,41 @@
+import warnings
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from ..error.InferEnvError import InferEnvError
+from ..InferProviders import InferProviders
+from ..warn.InferModelWarn import InferModelWarn
+
+if TYPE_CHECKING:
+    from ..infer_env import HAS_OPENVINO, HAS_ORT
+    if HAS_ORT:
+        import onnxruntime as ort  # pyright: ignore[reportMissingImports]
+    if HAS_OPENVINO:
+        import openvino  # pyright: ignore[reportMissingImports]
+
+def load_model(path: Path, providers: InferProviders) -> "tuple[ort.InferenceSession | None, openvino.CompiledModel | None]": # pyright: ignore[reportAttributeAccessIssue]
+    session = None
+    compiled_model = None
+    if providers.is_ort():
+            import onnxruntime as ort  # pyright: ignore[reportMissingImports]
+            session = ort.InferenceSession(
+                path.with_suffix(".onnx"), providers=providers.get_onnx_provider()
+            )
+    elif providers.is_ov():
+        from openvino import convert_model, save_model  # pyright: ignore[reportMissingImports]
+
+        from ..ov.OVCoreSingleton import core
+        if (path.parent / path.stem).with_suffix(".xml").exists():
+            model = core.read_model(str((path.parent / path.stem).with_suffix(".xml").resolve()))
+            compiled_model = core.compile_model(model=model, device_name=providers.get_openvino_device())
+        elif (path.parent / path.stem).with_suffix(".onnx").exists():
+            warnings.warn("您正在尝试使用 OpenVINO 运行 onnx，模型将自动转换并存档", InferModelWarn)
+            model = convert_model(str((path.parent / path.stem).with_suffix(".onnx").resolve()))
+            save_model(model, str((path.parent / path.stem).with_suffix(".xml").resolve()), compress_to_fp16=False)
+            compiled_model = core.compile_model(model=model, device_name=providers.get_openvino_device())
+        else:
+            raise FileNotFoundError(f"找不到模型: {path.parent / path.stem}")
+    else:
+        raise InferEnvError("没有任意一个支持的推理引擎或提供者 Flag 错误，请至少安装一个可选模块")
+    
+    return session, compiled_model

RvcPyInfer/ov/OVCoreSingleton.py

@@ -0,0 +1,3 @@
+import openvino as ov  # pyright: ignore[reportMissingImports]
+
+core = ov.Core()

RvcPyInfer/path_utils.py

@@ -0,0 +1,10 @@
+from pathlib import Path
+
+from .type_alist import PathLike
+
+
+def path(path: PathLike) -> Path:
+    if isinstance(path, str):
+        return Path(path)
+    else:
+        return path