pymss-core 0.1.0__tar.gz

pymss_core-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 KitsuneX07
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

pymss_core-0.1.0/PKG-INFO

@@ -0,0 +1,113 @@
+Metadata-Version: 2.4
+Name: pymss-core
+Version: 0.1.0
+Summary: Core model, configuration, and checkpoint package for music source separation.
+Author-email: KitsuneX07 <ghast1085654218@163.com>
+Maintainer-email: baicai1145 <3423714059@qq.com>
+License-Expression: MIT
+Project-URL: Bug Tracker, https://github.com/pymss-project/pymss-core/issues
+Project-URL: Source Code, https://github.com/pymss-project/pymss-core
+Project-URL: Documentation, https://github.com/pymss-project/pymss-core/blob/main/README.md
+Keywords: music source separation,audio separation,music processing,machine learning,audio
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Multimedia :: Sound/Audio
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy>=1.26
+Requires-Dist: pyyaml>=6.0.1
+Requires-Dist: torch>=2.7.1
+Provides-Extra: mlx
+Requires-Dist: mlx>=0.31.0; (sys_platform == "darwin" and platform_machine == "arm64") and extra == "mlx"
+Dynamic: license-file
+
+# pymss-core
+
+[中文文档](README_CN.md)
+
+Core model, configuration, and checkpoint package for music source separation.
+
+`pymss-core` is the shared low-level package for higher-level projects such as
+`pymss` inference and `pymsst` training. It contains model definitions,
+configuration loading, and checkpoint compatibility helpers. It intentionally
+does not include inference DSP pipelines, chunked demixing, audio file I/O,
+model downloads, catalog management, CLI, HTTP server, WebUI, datasets, losses,
+or training loops.
+
+## Install
+
+```bash
+pip install pymss-core
+```
+
+For local development:
+
+```bash
+uv sync --dev
+```
+
+Optional MLX backend on Apple Silicon:
+
+```bash
+pip install "pymss-core[mlx]"
+```
+
+## Public API
+
+```python
+from pymss_core import (
+    get_model_from_config,
+    load_config,
+    load_model_weights,
+)
+
+model, config = get_model_from_config("bs_roformer", "config.yaml")
+load_model_weights(model, "model.ckpt", model_type="bs_roformer", strict=True)
+
+model.eval()
+```
+
+## Package Boundary
+
+Included:
+
+- YAML config loading with `AttrDict`
+- PyTorch model definitions under `pymss_core.modules`
+- Optional MLX backend implementations for supported model forward paths
+- Model factory: `get_model_from_config(model_type, config_path)`
+- Checkpoint helpers for common MSS checkpoint containers
+- Small model-internal DSP math needed to construct model structures
+- VR network structures and VR model parameter JSON files
+
+Excluded:
+
+- Audio file decoding/encoding
+- Resampling, preprocessing, and full inference DSP pipelines
+- Tensor-level chunked demixing runtime
+- Model catalog, aliases, downloads, and cache management
+- CLI, server, WebUI, and endpoint schemas
+- Dataset, augmentation, loss, metrics, and trainer code
+- Any default dependency on MLX, Librosa, tqdm, Lightning, FastAPI, Uvicorn,
+  PyAV, WandB, or training extras
+
+## Repository Roles
+
+```text
+pymss-core
+  shared model/config/checkpoint layer
+
+pymss
+  user-facing inference package built on pymss-core, with audio I/O and demix
+
+pymsst
+  training package built on pymss-core, with training data/loss/runtime code
+```

pymss_core-0.1.0/README.md

@@ -0,0 +1,81 @@
+# pymss-core
+
+[中文文档](README_CN.md)
+
+Core model, configuration, and checkpoint package for music source separation.
+
+`pymss-core` is the shared low-level package for higher-level projects such as
+`pymss` inference and `pymsst` training. It contains model definitions,
+configuration loading, and checkpoint compatibility helpers. It intentionally
+does not include inference DSP pipelines, chunked demixing, audio file I/O,
+model downloads, catalog management, CLI, HTTP server, WebUI, datasets, losses,
+or training loops.
+
+## Install
+
+```bash
+pip install pymss-core
+```
+
+For local development:
+
+```bash
+uv sync --dev
+```
+
+Optional MLX backend on Apple Silicon:
+
+```bash
+pip install "pymss-core[mlx]"
+```
+
+## Public API
+
+```python
+from pymss_core import (
+    get_model_from_config,
+    load_config,
+    load_model_weights,
+)
+
+model, config = get_model_from_config("bs_roformer", "config.yaml")
+load_model_weights(model, "model.ckpt", model_type="bs_roformer", strict=True)
+
+model.eval()
+```
+
+## Package Boundary
+
+Included:
+
+- YAML config loading with `AttrDict`
+- PyTorch model definitions under `pymss_core.modules`
+- Optional MLX backend implementations for supported model forward paths
+- Model factory: `get_model_from_config(model_type, config_path)`
+- Checkpoint helpers for common MSS checkpoint containers
+- Small model-internal DSP math needed to construct model structures
+- VR network structures and VR model parameter JSON files
+
+Excluded:
+
+- Audio file decoding/encoding
+- Resampling, preprocessing, and full inference DSP pipelines
+- Tensor-level chunked demixing runtime
+- Model catalog, aliases, downloads, and cache management
+- CLI, server, WebUI, and endpoint schemas
+- Dataset, augmentation, loss, metrics, and trainer code
+- Any default dependency on MLX, Librosa, tqdm, Lightning, FastAPI, Uvicorn,
+  PyAV, WandB, or training extras
+
+## Repository Roles
+
+```text
+pymss-core
+  shared model/config/checkpoint layer
+
+pymss
+  user-facing inference package built on pymss-core, with audio I/O and demix
+
+pymsst
+  training package built on pymss-core, with training data/loss/runtime code
+```

pymss_core-0.1.0/pymss_core/__init__.py

@@ -0,0 +1,25 @@
+"""Core model, configuration, and checkpoint API for music source separation.
+
+`pymss_core` contains the shared pieces used by higher-level packages:
+configuration loading, model construction, model definitions, and
+checkpoint/state-dict helpers. It intentionally does not provide file audio
+I/O, inference DSP pipelines, chunked demixing, model catalog downloads, CLI,
+HTTP server, or WebUI functionality.
+"""
+
+from .checkpoint import load_checkpoint, load_model_weights, load_state_dict, unwrap_state_dict
+from .config import AttrDict, ConfigLoader, load_config, to_attrdict, to_plain
+from .utils import get_model_from_config
+
+__all__ = (
+    "AttrDict",
+    "ConfigLoader",
+    "get_model_from_config",
+    "load_checkpoint",
+    "load_config",
+    "load_model_weights",
+    "load_state_dict",
+    "to_attrdict",
+    "to_plain",
+    "unwrap_state_dict",
+)

pymss_core-0.1.0/pymss_core/checkpoint.py

@@ -0,0 +1,127 @@
+"""Checkpoint helpers shared by inference and training frontends."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from types import ModuleType
+from typing import Any
+
+import torch
+
+
+STATE_DICT_KEYS = ("state", "state_dict", "model_state_dict")
+
+
+def unwrap_state_dict(checkpoint: Any) -> Any:
+    """Return the model state dict from common MSS checkpoint containers."""
+    if isinstance(checkpoint, dict):
+        for key in STATE_DICT_KEYS:
+            if key in checkpoint:
+                return checkpoint[key]
+    return checkpoint
+
+
+def _install_demucs_pickle_stubs() -> dict[str, ModuleType | None]:
+    import sys
+    import types
+
+    module_names = ("demucs", "demucs.demucs", "demucs.hdemucs", "demucs.htdemucs")
+    previous = {name: sys.modules.get(name) for name in module_names}
+    package = sys.modules.setdefault("demucs", types.ModuleType("demucs"))
+    package.__path__ = []
+    for module_name, class_names in {
+        "demucs": ("Demucs",),
+        "hdemucs": ("HDemucs", "HTDemucs"),
+        "htdemucs": ("HTDemucs",),
+    }.items():
+        full_name = f"demucs.{module_name}"
+        module = sys.modules.setdefault(full_name, types.ModuleType(full_name))
+        setattr(package, module_name, module)
+        for class_name in class_names:
+            if not hasattr(module, class_name):
+                setattr(module, class_name, type(class_name, (), {"__module__": full_name}))
+    return previous
+
+
+def _restore_modules(previous: dict[str, ModuleType | None]) -> None:
+    import sys
+
+    for name, module in previous.items():
+        if module is None:
+            sys.modules.pop(name, None)
+        else:
+            sys.modules[name] = module
+
+
+def _torch_load(path: str | Path, *, map_location="cpu", weights_only: bool | None = None, mmap: bool = True) -> Any:
+    kwargs: dict[str, Any] = {"map_location": map_location}
+    if weights_only is not None:
+        kwargs["weights_only"] = weights_only
+    if mmap:
+        kwargs["mmap"] = True
+    try:
+        return torch.load(path, **kwargs)
+    except TypeError:
+        kwargs.pop("mmap", None)
+        try:
+            return torch.load(path, **kwargs)
+        except TypeError:
+            kwargs.pop("weights_only", None)
+            return torch.load(path, **kwargs)
+
+
+def load_checkpoint(
+    path: str | Path,
+    *,
+    model_type: str | None = None,
+    map_location: str | torch.device = "cpu",
+    weights_only: bool | None = None,
+    mmap: bool = True,
+) -> Any:
+    """Load a checkpoint package with compatibility for common MSS formats."""
+    model_type = (model_type or "").lower()
+    if model_type in {"htdemucs", "demucs", "legacy_demucs", "legacy_tasnet"}:
+        previous = _install_demucs_pickle_stubs()
+        try:
+            return _torch_load(path, map_location=map_location, weights_only=False, mmap=mmap)
+        finally:
+            _restore_modules(previous)
+    if model_type == "apollo":
+        weights_only = False if weights_only is None else weights_only
+    return _torch_load(path, map_location=map_location, weights_only=weights_only, mmap=mmap)
+
+
+def load_state_dict(
+    path: str | Path,
+    *,
+    model_type: str | None = None,
+    map_location: str | torch.device = "cpu",
+    weights_only: bool | None = None,
+    mmap: bool = True,
+) -> Any:
+    """Load and unwrap the model state dict from a checkpoint file."""
+    return unwrap_state_dict(
+        load_checkpoint(
+            path,
+            model_type=model_type,
+            map_location=map_location,
+            weights_only=weights_only,
+            mmap=mmap,
+        )
+    )
+
+
+def load_model_weights(
+    model: torch.nn.Module,
+    checkpoint_or_path: Any,
+    *,
+    model_type: str | None = None,
+    strict: bool = True,
+    map_location: str | torch.device = "cpu",
+) -> Any:
+    """Load weights from a checkpoint package or file into a model."""
+    if isinstance(checkpoint_or_path, (str, Path)):
+        state_dict = load_state_dict(checkpoint_or_path, model_type=model_type, map_location=map_location)
+    else:
+        state_dict = unwrap_state_dict(checkpoint_or_path)
+    return model.load_state_dict(state_dict, strict=strict)

pymss_core-0.1.0/pymss_core/config.py

@@ -0,0 +1,135 @@
+import re
+
+import yaml
+
+
+class ConfigLoader(yaml.FullLoader):
+    """YAML loader used by pymss-core model configuration files."""
+
+    pass
+
+
+ConfigLoader.add_implicit_resolver(
+    "tag:yaml.org,2002:float",
+    re.compile(
+        r"""^[-+]?(
+            ([0-9][0-9_]*)?\.[0-9_]+([eE][-+]?[0-9]+)?
+            |[0-9][0-9_]*[eE][-+]?[0-9]+
+            |\.(inf|Inf|INF)
+            |\.(nan|NaN|NAN)
+        )$""",
+        re.X,
+    ),
+    list("-+0123456789."),
+)
+
+
+class AttrDict(dict):
+    """Dictionary that recursively exposes keys as attributes.
+
+    Args:
+        data (Mapping | None, optional): Data value. Defaults to None.
+        **kwargs: Additional keyword arguments.
+
+    Example:
+        >>> cfg = AttrDict({"audio": {"chunk_size": 485100}})
+        >>> cfg.audio.chunk_size
+        485100"""
+
+    def __init__(self, data=None, **kwargs):
+        """Initialize the instance.
+
+        Args:
+            data (Mapping | None, optional): Data value. Defaults to None.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            None: This method completes for its side effects."""
+        super().__init__()
+        for key, value in dict(data or {}, **kwargs).items():
+            self[key] = value
+
+    def __getattr__(self, key):
+        """Return a missing attribute from the underlying mapping.
+
+        Args:
+            key (str): Key value.
+
+        Returns:
+            Any: Computed result."""
+        try:
+            return self[key]
+        except KeyError as exc:
+            raise AttributeError(key) from exc
+
+    def __setattr__(self, key, value):
+        """Store an attribute assignment in the underlying mapping.
+
+        Args:
+            key (str): Key value.
+            value (Any): Value value.
+
+        Returns:
+            None: This method completes for its side effects."""
+        self[key] = value
+
+    def __setitem__(self, key, value):
+        """Store an item after recursively converting nested dictionaries.
+
+        Args:
+            key (str): Key value.
+            value (Any): Value value.
+
+        Returns:
+            None: This method completes for its side effects."""
+        super().__setitem__(key, to_attrdict(value))
+
+
+def to_attrdict(value):
+    """Recursively convert dictionaries to AttrDict objects.
+
+    Args:
+        value (Any): Value value.
+
+    Returns:
+        Any: Converted value with nested dictionaries wrapped as AttrDict."""
+    if isinstance(value, dict):
+        return value if isinstance(value, AttrDict) else AttrDict(value)
+    if isinstance(value, list):
+        return [to_attrdict(item) for item in value]
+    if isinstance(value, tuple):
+        return tuple(to_attrdict(item) for item in value)
+    return value
+
+
+def to_plain(value):
+    """Recursively convert AttrDict objects back to plain Python containers.
+
+    Args:
+        value (Any): Value value.
+
+    Returns:
+        Any: Converted value using plain dictionaries, lists, and tuples."""
+    if isinstance(value, AttrDict):
+        return {key: to_plain(item) for key, item in value.items()}
+    if isinstance(value, list):
+        return [to_plain(item) for item in value]
+    if isinstance(value, tuple):
+        return tuple(to_plain(item) for item in value)
+    return value
+
+
+def load_config(path):
+    """Load a YAML model configuration file.
+
+    Args:
+        path (str | os.PathLike): File system path.
+
+    Returns:
+        AttrDict: Parsed configuration with attribute access.
+
+    Example:
+        >>> config = load_config("config.yaml")
+        >>> config.inference.batch_size"""
+    with open(path, encoding="utf-8") as f:
+        return to_attrdict(yaml.load(f, Loader=ConfigLoader))

pymss_core-0.1.0/pymss_core/modules/_dsp.py

@@ -0,0 +1,84 @@
+"""Small DSP helpers needed by model definitions."""
+
+from __future__ import annotations
+
+import numpy as np
+
+
+def hz_to_midi(hz):
+    """Convert frequencies in Hz to MIDI note numbers."""
+    hz = np.asarray(hz)
+    return 69.0 + 12.0 * np.log2(hz / 440.0)
+
+
+def midi_to_hz(midi):
+    """Convert MIDI note numbers to frequencies in Hz."""
+    midi = np.asarray(midi)
+    return 440.0 * np.power(2.0, (midi - 69.0) / 12.0)
+
+
+def _hz_to_mel(frequencies, *, htk=False):
+    frequencies = np.asarray(frequencies, dtype=np.float64)
+    if htk:
+        return 2595.0 * np.log10(1.0 + frequencies / 700.0)
+
+    f_sp = 200.0 / 3
+    mels = frequencies / f_sp
+    min_log_hz = 1000.0
+    min_log_mel = min_log_hz / f_sp
+    logstep = np.log(6.4) / 27.0
+    log_t = frequencies >= min_log_hz
+    mels = np.array(mels, copy=True)
+    mels[log_t] = min_log_mel + np.log(frequencies[log_t] / min_log_hz) / logstep
+    return mels
+
+
+def _mel_to_hz(mels, *, htk=False):
+    mels = np.asarray(mels, dtype=np.float64)
+    if htk:
+        return 700.0 * (np.power(10.0, mels / 2595.0) - 1.0)
+
+    f_sp = 200.0 / 3
+    freqs = f_sp * mels
+    min_log_hz = 1000.0
+    min_log_mel = min_log_hz / f_sp
+    logstep = np.log(6.4) / 27.0
+    log_t = mels >= min_log_mel
+    freqs = np.array(freqs, copy=True)
+    freqs[log_t] = min_log_hz * np.exp(logstep * (mels[log_t] - min_log_mel))
+    return freqs
+
+
+def mel_frequencies(n_mels, *, fmin=0.0, fmax=11025.0, htk=False):
+    """Return center frequencies on the mel scale, including endpoints."""
+    min_mel = _hz_to_mel(fmin, htk=htk)
+    max_mel = _hz_to_mel(fmax, htk=htk)
+    return _mel_to_hz(np.linspace(min_mel, max_mel, int(n_mels)), htk=htk)
+
+
+def fft_frequencies(*, sr, n_fft):
+    """Return FFT bin center frequencies."""
+    return np.linspace(0.0, float(sr) / 2.0, int(1 + n_fft // 2), endpoint=True)
+
+
+def mel_filterbank(sr, n_fft, n_mels=128, fmin=0.0, fmax=None, htk=False, norm="slaney", dtype=np.float32):
+    """Create a triangular mel filterbank for model initialization."""
+    if fmax is None:
+        fmax = float(sr) / 2.0
+
+    mel_f = mel_frequencies(int(n_mels) + 2, fmin=fmin, fmax=fmax, htk=htk)
+    fft_f = fft_frequencies(sr=sr, n_fft=n_fft)
+
+    fdiff = np.diff(mel_f)
+    ramps = np.subtract.outer(mel_f, fft_f)
+    lower = -ramps[:-2] / fdiff[:-1, np.newaxis]
+    upper = ramps[2:] / fdiff[1:, np.newaxis]
+    weights = np.maximum(0.0, np.minimum(lower, upper))
+
+    if norm == "slaney":
+        enorm = 2.0 / (mel_f[2 : int(n_mels) + 2] - mel_f[: int(n_mels)])
+        weights *= enorm[:, np.newaxis]
+    elif norm is not None:
+        raise ValueError(f"Unsupported mel filterbank norm: {norm!r}")
+
+    return weights.astype(dtype, copy=False)