pymss-core 0.1.0__py3-none-any.whl

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/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/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/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)

pymss_core/modules/apollo_mlx.py

@@ -0,0 +1,242 @@
+import torch
+
+from .bs_roformer.mlx_attention import _mlx_dtype, _torch_to_mlx_array, mlx_to_torch_mps
+from .look2hear.apollo import BSNet, ConvActNorm1d, ICB, RMSNorm
+
+
+def torch_to_mlx_input(tensor, dtype):
+    import mlx.core as mx
+
+    return mx.array(tensor.detach().to(dtype=dtype).cpu().numpy())
+
+
+def _mlx_param(module, name, tensor, dtype):
+    cache = getattr(module, "_pymss_mlx_full_param_cache", None)
+    if cache is None:
+        cache = {}
+        module._pymss_mlx_full_param_cache = cache
+    key = (name, tensor.data_ptr(), tensor._version, tuple(tensor.shape), dtype)
+    cached = cache.get(name)
+    if cached is not None and cached[0] == key:
+        return cached[1]
+    value = _torch_to_mlx_array(tensor, dtype)
+    cache[name] = (key, value)
+    return value
+
+
+def _reflect_pad_last(x, pad):
+    import mlx.core as mx
+
+    if pad <= 0:
+        return x
+    if x.shape[-1] <= pad:
+        raise ValueError("reflect padding requires input length greater than padding")
+    return mx.concatenate((x[..., 1 : pad + 1][..., ::-1], x, x[..., -pad - 1 : -1][..., ::-1]), axis=-1)
+
+
+def _stft(module, raw_audio, dtype):
+    import mlx.core as mx
+
+    batch_channels, length = raw_audio.shape
+    n_fft = module.win
+    hop = module.stride
+    x = _reflect_pad_last(raw_audio.astype(dtype), n_fft // 2)
+    frames = 1 + (x.shape[-1] - n_fft) // hop
+    framed = mx.as_strided(x, shape=(batch_channels, frames, n_fft), strides=(x.shape[-1], hop, 1))
+    window = _torch_to_mlx_array(module.window, torch.float32).astype(dtype)
+    spec = mx.fft.rfft(framed * window, n=n_fft, axis=-1)
+    return mx.moveaxis(spec, -1, -2), {"length": length, "n_fft": n_fft, "hop": hop, "window": window, "dtype": dtype}
+
+
+def _istft(spec, context):
+    import mlx.core as mx
+
+    n_fft = context["n_fft"]
+    hop = context["hop"]
+    frames = mx.fft.irfft(mx.moveaxis(spec, -2, -1), n=n_fft, axis=-1).astype(context["dtype"]) * context["window"]
+    frame_count = frames.shape[1]
+    full_length = n_fft + hop * (frame_count - 1)
+    positions = mx.arange(n_fft)[None, :] + hop * mx.arange(frame_count)[:, None]
+    audio = mx.zeros((frames.shape[0], full_length), dtype=context["dtype"]).at[:, positions].add(frames)
+    denom_frames = mx.broadcast_to(mx.square(context["window"])[None, :], (frame_count, n_fft))
+    denom = mx.zeros((full_length,), dtype=context["dtype"]).at[positions].add(denom_frames)
+    audio = audio / mx.maximum(denom[None, :], mx.array(1e-11, dtype=context["dtype"]))
+    pad = n_fft // 2
+    return audio[..., pad : pad + context["length"]]
+
+
+def _conv1d_ncl(conv, x, dtype):
+    import mlx.core as mx
+
+    weight = _mlx_param(conv, "weight", conv.weight, dtype).transpose(0, 2, 1)
+    y = mx.conv1d(
+        x.transpose(0, 2, 1),
+        weight,
+        stride=conv.stride[0],
+        padding=conv.padding[0],
+        dilation=conv.dilation[0],
+        groups=conv.groups,
+    )
+    if conv.bias is not None:
+        y = y + _mlx_param(conv, "bias", conv.bias, dtype)
+    return y.transpose(0, 2, 1)
+
+
+def _rms_norm(module, x, dtype):
+    import mlx.core as mx
+
+    batch, channels, frames = x.shape
+    groups = int(module.groups)
+    y = x.astype(mx.float32).reshape(batch, groups, channels // groups, frames)
+    y = y * mx.rsqrt(mx.mean(mx.square(y), axis=2, keepdims=True) + module.eps)
+    y = y.reshape(batch, channels, frames).astype(x.dtype)
+    return y * _mlx_param(module, "weight", module.weight, dtype).reshape(1, -1, 1)
+
+
+def _silu(x):
+    import mlx.core as mx
+
+    return x * mx.sigmoid(x)
+
+
+def _glu_channel(x):
+    import mlx.core as mx
+
+    a, b = mx.split(x, 2, axis=1)
+    return a * mx.sigmoid(b)
+
+
+def _module_forward(module, x, dtype):
+    if isinstance(module, torch.nn.Sequential):
+        for child in module:
+            x = _module_forward(child, x, dtype)
+        return x
+    if isinstance(module, torch.nn.Conv1d):
+        return _conv1d_ncl(module, x, dtype)
+    if isinstance(module, RMSNorm):
+        return _rms_norm(module, x, dtype)
+    if isinstance(module, torch.nn.SiLU):
+        return _silu(x)
+    if isinstance(module, torch.nn.GLU):
+        return _glu_channel(x)
+    if isinstance(module, ConvActNorm1d):
+        return _conv_act_norm(module, x, dtype)
+    if isinstance(module, ICB):
+        return _module_forward(module.blocks, x, dtype)
+    if isinstance(module, BSNet):
+        return _bsnet(module, x, dtype)
+    raise TypeError(f"unsupported Apollo layer for MLX full backend: {type(module).__name__}")
+
+
+def _conv_act_norm(module, x, dtype):
+    y = _conv1d_ncl(module.conv[0], x, dtype)
+    y = _rms_norm(module.conv[1], y, dtype)
+    y = _conv1d_ncl(module.conv[2], y, dtype)
+    y = _silu(y)
+    y = _conv1d_ncl(module.conv[4], y, dtype)
+    if module.causal:
+        y = y[..., : -module.kernel + 1]
+    return x + y
+
+
+def _apply_rope(module, x, dtype):
+    import mlx.core as mx
+
+    seq_len = x.shape[-2]
+    cos = _torch_to_mlx_array(module.cos_freq[:seq_len], dtype).reshape(1, 1, seq_len, -1)
+    sin = _torch_to_mlx_array(module.sin_freq[:seq_len], dtype).reshape(1, 1, seq_len, -1)
+    even, odd = x[..., 0::2], x[..., 1::2]
+    cos_e = cos[..., 0::2]
+    sin_e = sin[..., 0::2]
+    out = mx.zeros_like(x)
+    out = out.at[..., 0::2].add(even * cos_e - odd * sin_e)
+    out = out.at[..., 1::2].add(odd * cos_e + even * sin_e)
+    return out
+
+
+def _roformer(module, x, dtype):
+    import mlx.core as mx
+
+    batch, _, frames = x.shape
+    x_norm = _rms_norm(module.input_norm, x, dtype)
+    qkv = _conv1d_ncl(module.weight, x_norm, dtype)
+    qkv = qkv.reshape(batch, module.num_head, module.hidden_size * 3, frames).transpose(0, 1, 3, 2)
+    q, k, v = mx.split(qkv, 3, axis=-1)
+    q = _apply_rope(module, q, dtype)
+    k = _apply_rope(module, k, dtype)
+    attn = mx.fast.scaled_dot_product_attention(q, k, v, scale=module.hidden_size**-0.5, mask=None)
+    out = attn.transpose(0, 1, 3, 2).reshape(batch, -1, frames)
+    out = _conv1d_ncl(module.output, out, dtype) + x
+
+    hidden = _rms_norm(module.MLP[0], out, dtype)
+    hidden = _conv1d_ncl(module.MLP[1], hidden, dtype)
+    hidden = _silu(hidden)
+    gate, z = mx.split(hidden, 2, axis=1)
+    return out + _conv1d_ncl(module.MLP_output, _silu(gate) * z, dtype)
+
+
+def _bsnet(module, x, dtype):
+    batch, bands, channels, frames = x.shape
+    band = x.transpose(0, 3, 2, 1).reshape(batch * frames, channels, bands)
+    band = _roformer(module.band_net, band, dtype)
+    band = band.reshape(batch, frames, channels, bands).transpose(0, 3, 2, 1)
+    seq = _module_forward(module.seq_net, band.reshape(batch * bands, channels, frames), dtype)
+    return seq.reshape(batch, bands, channels, frames)
+
+
+def _feature_extractor(module, raw_audio, dtype):
+    import mlx.core as mx
+
+    mx_dtype = _mlx_dtype(dtype)
+    batch, channels, samples = raw_audio.shape
+    spec, _ = _stft(module, raw_audio.reshape(batch * channels, samples), mx_dtype)
+    features = []
+    powers = []
+    band_index = 0
+    for width, bn in zip(module.band_width, module.BN):
+        sub = spec[:, band_index : band_index + width]
+        power = mx.sqrt(mx.sum(mx.square(sub.real) + mx.square(sub.imag), axis=1, keepdims=True) + module.eps)
+        norm = sub / power
+        inp = mx.concatenate((norm.real, norm.imag, mx.log(power)), axis=1)
+        features.append(_module_forward(bn, inp.astype(mx_dtype), dtype))
+        powers.append(power)
+        band_index += width
+    return mx.stack(features, axis=1), spec
+
+
+def _estimate_spec(module, feature, batch_channels, dtype):
+    import mlx.core as mx
+
+    specs = []
+    for band_feature, output, width in zip(mx.split(feature, feature.shape[1], axis=1), module.output, module.band_width):
+        band_feature = band_feature[:, 0]
+        ri = _module_forward(output, band_feature, dtype).reshape(batch_channels, 2, width, -1)
+        specs.append(ri[:, 0] + (1j * ri[:, 1]))
+    return mx.concatenate(specs, axis=1)
+
+
+def mlx_forward_apollo_mx(module, raw_audio, dtype=torch.float16):
+    if dtype not in (torch.float16, torch.float32):
+        raise TypeError("MLX full Apollo supports torch.float16 or torch.float32 compute dtype")
+    mx_dtype = _mlx_dtype(dtype)
+    raw_audio = raw_audio.astype(mx_dtype)
+    batch, channels, samples = raw_audio.shape
+    feature, _ = _feature_extractor(module, raw_audio, dtype)
+    for block in module.net:
+        feature = _bsnet(block, feature, dtype)
+    est_spec = _estimate_spec(module, feature, batch * channels, dtype)
+    return _istft(
+        est_spec,
+        {
+            "length": samples,
+            "n_fft": module.win,
+            "hop": module.stride,
+            "window": _torch_to_mlx_array(module.window, torch.float32).astype(mx_dtype),
+            "dtype": mx_dtype,
+        },
+    ).reshape(batch, channels, -1)
+
+
+def mlx_forward_apollo(module, raw_audio, dtype=torch.float16):
+    x_mx = torch_to_mlx_input(raw_audio, dtype=dtype)
+    return mlx_to_torch_mps(mlx_forward_apollo_mx(module, x_mx, dtype), raw_audio)