gamemlx 0.1.0__py3-none-any.whl

gamemlx/__init__.py

@@ -0,0 +1,73 @@
+"""gamemlx: pure-MLX inference package for GAME (vocal MIDI extraction)."""
+from .config import GamemlxConfig, load_config
+from .model import GAMEModel, load_model as _load_model_local
+
+__all__ = ["GAMEModel", "load_model", "load_config", "GamemlxConfig", "extract"]
+
+
+def load_model(name_or_path: str = "medium", precision: str = "fp32"):
+    """Load a model by alias (small/medium/large), HF repo id, or local path.
+
+    Aliases download from HuggingFace (ModelScope fallback) into
+    ~/.cache/gamemlx on first use. Returns (GAMEModel, lang_map).
+    """
+    from .config import load_lang_map
+    from .download import download_model
+    model_dir = download_model(name_or_path, precision="both")
+    model = _load_model_local(model_dir, precision=precision)
+    lang_map = load_lang_map(model_dir)
+    return model, lang_map
+
+
+def extract(model, path, language=None, output_dir=None, tempo=120.0,
+            seg_threshold=0.2, seg_radius=0.02, est_threshold=0.2,
+            t0=0.0, nsteps=8, rng_seed=None):
+    """High-level API: extract MIDI from a file or directory.
+
+    `model` is a GAMEModel (see load_model). Mirrors gamemlx.cli extract.
+    """
+    import pathlib
+
+    import mlx.core as mx
+    import numpy as np
+
+    from .cli import _MidiWriter, _parse_filemap
+    from .config import load_lang_map
+    from .download import CACHE_DIR
+    from .slicer import Slicer
+    from .utils import get_language_id, load_audio
+
+    path = pathlib.Path(path)
+    filemap = _parse_filemap(path, {".wav", ".flac", ".mp3", ".aac", ".ogg"})
+    if output_dir is None:
+        output_dir = path if path.is_dir() else path.parent
+    output_dir = pathlib.Path(output_dir)
+
+    sr = model.config.features.audio_sample_rate
+    slicer = Slicer(sr=sr, threshold=-40., min_length=1000, min_interval=200,
+                    max_sil_kept=100)
+    radius_frames = round(seg_radius / model.timestep)
+    ts = mx.array([t0 + i * (1 - t0) / nsteps for i in range(nsteps)],
+                  dtype=mx.float32)
+
+    writer = _MidiWriter(output_dir, tempo)
+    for key, filepath in filemap.items():
+        lang_map = None  # language resolved by caller knowledge; default id 0
+        language_id = get_language_id(language, lang_map if language else None)
+        wav = load_audio(filepath, sr=sr)
+        for chunk in slicer.slice(wav):
+            cw = np.asarray(chunk["waveform"], dtype=np.float32)
+            duration = cw.shape[0] / sr
+            rng = np.random.default_rng(rng_seed) if rng_seed is not None else None
+            dur, pres, sc = model(
+                mx.array(cw[None, :]), mx.array([[duration]], dtype=mx.float32),
+                boundary_threshold=seg_threshold, boundary_radius=radius_frames,
+                score_threshold=est_threshold, language=mx.array([language_id]),
+                t=ts, rng_matrices=rng,
+            )
+            mx.eval(dur, pres, sc)
+            writer.add_chunk(key, chunk["offset"], duration,
+                             np.asarray(dur[0]), np.asarray(sc[0]),
+                             np.asarray(pres[0]))
+    outs = [writer.flush(key) for key in list(filemap.keys())]
+    return outs

gamemlx/cli.py

@@ -0,0 +1,199 @@
+"""gamemlx CLI: extract MIDI from audio, mirroring GAME/infer.py extract."""
+import glob
+import json
+import pathlib
+
+import click
+
+from .config import load_lang_map
+from .download import download_model
+from .model import load_model
+from .slicer import Slicer
+from .utils import get_language_id, load_audio
+
+
+def _validate_exts(ctx, param, value):
+    return {"." + ext.strip().lower() for ext in value.split(",")}
+
+
+def _validate_output_formats(ctx, param, value):
+    formats = {fmt.strip().lower() for fmt in value.split(",")}
+    supported = {"mid", "txt", "csv"}
+    if not formats.issubset(supported):
+        raise click.BadParameter(f"Unsupported formats: {formats - supported}.")
+    return formats
+
+
+def _parse_filemap(path: pathlib.Path, exts: set) -> dict[str, pathlib.Path]:
+    if path.is_file():
+        return {path.name: path}
+    files = [f for f in path.rglob("*") if f.is_file() and f.suffix.lower() in exts]
+    filemap = {f.relative_to(path).as_posix(): f for f in files}
+    if not filemap:
+        raise FileNotFoundError(f"No audio files found in directory: {path}")
+    return filemap
+
+
+class _MidiWriter:
+    """Equivalent of GAME SaveCombinedMidiFileCallback (aggregation + sorting)."""
+
+    def __init__(self, output_dir: pathlib.Path, tempo: float):
+        self.output_dir = output_dir
+        self.tempo = tempo
+        self.notes: dict[str, list] = {}
+
+    def add_chunk(self, key, offset, length, durations, scores, presence):
+        import mlx.core as mx
+        import numpy as np
+        d = np.asarray(durations)
+        s = np.asarray(scores)
+        p = np.asarray(presence).astype(bool)
+        # exact torch: pad(d,(1,0)).cumsum.clamp(max=length).add(offset)
+        note_onset = np.concatenate([[0.0], np.cumsum(d)])[:len(d) + 1]
+        note_onset = np.minimum(note_onset, length) + offset
+        note_offset = np.minimum(np.cumsum(d), length) + offset
+        for onset, off, score, valid in zip(note_onset, note_offset, s, p):
+            if off - onset <= 0 or not valid:
+                continue
+            self.notes.setdefault(key, []).append((onset, off, float(score)))
+
+    def flush(self, key):
+        import mido
+        notes = self.notes.pop(key, [])
+        notes.sort(key=lambda x: (x[0], x[1], x[2]))
+        # clamp overlapping notes like torch callback
+        last_time = 0.0
+        cleaned = []
+        for onset, offset, pitch in notes:
+            onset = max(onset, last_time)
+            offset = max(offset, onset)
+            if offset <= onset:
+                continue
+            last_time = offset
+            cleaned.append((onset, offset, pitch))
+        track = mido.MidiTrack()
+        track.append(mido.MetaMessage("set_tempo", tempo=mido.bpm2tempo(self.tempo), time=0))
+        last_time = 0
+        for onset, offset, pitch in cleaned:
+            onset_ticks = round(onset * self.tempo * 8)
+            offset_ticks = round(offset * self.tempo * 8)
+            midi_pitch = round(pitch)
+            if offset_ticks <= onset_ticks:
+                continue
+            track.append(mido.Message("note_on", note=midi_pitch, time=onset_ticks - last_time))
+            track.append(mido.Message("note_off", note=midi_pitch, time=offset_ticks - onset_ticks))
+            last_time = offset_ticks
+        filepath = (self.output_dir / key).with_suffix(".mid")
+        filepath.parent.mkdir(parents=True, exist_ok=True)
+        with mido.MidiFile(charset="utf8") as midi_file:
+            midi_file.tracks.append(track)
+            midi_file.save(filepath)
+        return filepath
+
+
+@click.group()
+def main():
+    pass
+
+
+@main.command(name="extract", help="Extract MIDI from single or multiple audio files.")
+@click.argument("path", type=click.Path(exists=True, dir_okay=True, file_okay=True,
+                                        readable=True, path_type=pathlib.Path))
+@click.option("-m", "--model", type=str, required=True,
+              help="Model alias (small|medium|large, auto-download) or path to a local MLX model directory.")
+@click.option("-l", "--language", type=str, default=None, show_default=False,
+              help="Language code for better segmentation if supported.")
+@click.option("--output-formats", type=str, default="mid", show_default=True,
+              callback=_validate_output_formats,
+              help="Output formats: mid, txt, csv.")
+@click.option("--tempo", type=click.FloatRange(min=0, min_open=True), default=120,
+              show_default=True, help="Tempo (BPM) for MIDI files.")
+@click.option("--output-dir", type=click.Path(dir_okay=True, file_okay=False,
+                                              writable=True, path_type=pathlib.Path),
+              default=None, help="Directory for results (default: alongside input).")
+@click.option("--precision", type=click.Choice(["fp32", "fp16"]), default="fp32",
+              show_default=True, help="Compute precision.")
+@click.option("--seg-threshold", type=float, default=0.2, show_default=True,
+              help="Boundary decoding threshold.")
+@click.option("--seg-radius", type=float, default=0.02, show_default=True,
+              help="Boundary decoding radius in seconds (divided by timestep, rounded).")
+@click.option("--est-threshold", type=float, default=0.2, show_default=True,
+              help="Presence threshold.")
+@click.option("--t0", type=float, default=0.0, show_default=True, help="D3PM start T.")
+@click.option("--nsteps", type=int, default=8, show_default=True, help="D3PM steps.")
+@click.option("--input-formats", type=str, default="wav,flac,mp3,aac,ogg",
+              show_default=True, callback=_validate_exts)
+@click.option("--rng-seed", type=int, default=None,
+              help="numpy seed for d3pm randomness (default: nondeterministic).")
+@click.pass_context
+def extract(ctx, path, model, language, output_formats, tempo, output_dir, precision,
+            seg_threshold, seg_radius, est_threshold, t0, nsteps, input_formats, rng_seed):
+    import time
+
+    import mlx.core as mx
+    import numpy as np
+
+    filemap = _parse_filemap(path, input_formats)
+    if output_dir is None:
+        output_dir = path if path.is_dir() else path.parent
+
+    t_start = time.time()
+    # -m accepts alias (small/medium/large -> auto-download), HF repo id, or local dir
+    model_arg = str(model)
+    if pathlib.Path(model_arg).is_dir():
+        model_dir = pathlib.Path(model_arg)
+    else:
+        model_dir = download_model(model_arg, precision="both")
+    m = load_model(model_dir, precision=precision)
+    sr = m.config.features.audio_sample_rate
+    timestep = m.timestep
+    lang_map = load_lang_map(model_dir)
+    language_id = get_language_id(language, lang_map)
+
+    slicer = Slicer(sr=sr, threshold=-40., min_length=1000, min_interval=200,
+                    max_sil_kept=100)
+    radius_frames = round(seg_radius / timestep)
+
+    # d3pm t values: game ValidationConfig.d3pm_sample_ts_resolved
+    ts = [t0 + i * (1 - t0) / nsteps for i in range(nsteps)]
+    ts_mx = mx.array(ts, dtype=mx.float32)
+
+    writers = {}
+    if "mid" in output_formats:
+        writers["mid"] = _MidiWriter(output_dir, tempo)
+
+    n_chunks = 0
+    for key, filepath in filemap.items():
+        wav = load_audio(filepath, sr=sr)
+        chunks = slicer.slice(wav)
+        for chunk in chunks:
+            chunk_wav = np.asarray(chunk["waveform"], dtype=np.float32)
+            duration = chunk_wav.shape[0] / sr
+            rng = np.random.default_rng(rng_seed) if rng_seed is not None else None
+            dur, pres, sc = m(
+                mx.array(chunk_wav[None, :]),
+                mx.array([[duration]], dtype=mx.float32),
+                boundary_threshold=seg_threshold,
+                boundary_radius=radius_frames,
+                score_threshold=est_threshold,
+                language=mx.array([language_id]),
+                t=ts_mx,
+                rng_matrices=rng,
+            )
+            mx.eval(dur, pres, sc)
+            for w in writers.values():
+                w.add_chunk(key, chunk["offset"], duration,
+                            np.asarray(dur[0]), np.asarray(sc[0]), np.asarray(pres[0]))
+            n_chunks += 1
+
+    for key in list(filemap.keys()):
+        for w in writers.values():
+            f = w.flush(key)
+            if "mid" in writers and w is writers["mid"]:
+                click.echo(f"Saved MIDI file: {f.as_posix()}")
+    click.echo(f"Done: {len(filemap)} file(s), {n_chunks} chunk(s), "
+               f"total {time.time() - t_start:.1f}s ({precision}).")
+
+
+if __name__ == "__main__":
+    main()

gamemlx/config.py

@@ -0,0 +1,113 @@
+"""Model and inference configuration, mirroring GAME config.yaml structure.
+
+Reference: GAME lib/config/schema.py (ModelConfig / InferenceConfig)
+"""
+import dataclasses
+import json
+import pathlib
+from typing import Any
+
+
+@dataclasses.dataclass
+class FeaturesConfig:
+    audio_sample_rate: int = 44100
+    hop_size: int = 441
+    fft_size: int = 2048
+    win_size: int = 2048
+    n_mels: int = 80
+    fmin: float = 0.0
+    fmax: float = 8000.0
+
+    @property
+    def timestep(self) -> float:
+        return self.hop_size / self.audio_sample_rate
+
+
+@dataclasses.dataclass
+class BackbonesConfig:
+    cls: str
+    kwargs: dict[str, Any] = dataclasses.field(default_factory=dict)
+
+
+@dataclasses.dataclass
+class ModelConfigData:
+    mode: str = "d3pm"
+    use_languages: bool = True
+    num_languages: int = 127
+    region_cycle_len: int = 3
+    in_dim: int = 80
+    embedding_dim: int = 256
+    estimator_out_dim: int = 257
+    encoder: BackbonesConfig = None
+    segmenter: BackbonesConfig = None
+    estimator: BackbonesConfig = None
+
+
+@dataclasses.dataclass
+class InferenceConfigData:
+    midi_min: float = 0.0
+    midi_max: float = 128.0
+    midi_num_bins: int = 257
+    midi_std: float = 0.5
+
+
+@dataclasses.dataclass
+class GamemlxConfig:
+    features: FeaturesConfig
+    model: ModelConfigData
+    inference: InferenceConfigData
+    raw: dict = dataclasses.field(default_factory=dict, repr=False)
+
+    @property
+    def timestep(self) -> float:
+        return self.features.timestep
+
+
+def load_config(path: str | pathlib.Path) -> GamemlxConfig:
+    """Load config.yaml (GAME format: model: / inference: sections)."""
+    import yaml
+
+    path = pathlib.Path(path)
+    with open(path, "r", encoding="utf8") as f:
+        raw = yaml.safe_load(f)
+
+    m = raw["model"]
+    i = raw["inference"]
+    feats = i["features"]
+    spec = feats["spectrogram"]
+    features = FeaturesConfig(
+        audio_sample_rate=feats["audio_sample_rate"],
+        hop_size=feats["hop_size"],
+        fft_size=feats["fft_size"],
+        win_size=feats["win_size"],
+        n_mels=spec["num_bins"],
+        fmin=float(spec["fmin"]),
+        fmax=float(spec["fmax"]),
+    )
+    model = ModelConfigData(
+        mode=m["mode"],
+        use_languages=m["use_languages"],
+        num_languages=m["num_languages"],
+        region_cycle_len=m["region_cycle_len"],
+        in_dim=m["in_dim"],
+        embedding_dim=m["embedding_dim"],
+        estimator_out_dim=m["estimator_out_dim"],
+        encoder=BackbonesConfig(cls=m["encoder"]["cls"], kwargs=dict(m["encoder"]["kwargs"])),
+        segmenter=BackbonesConfig(cls=m["segmenter"]["cls"], kwargs=dict(m["segmenter"]["kwargs"])),
+        estimator=BackbonesConfig(cls=m["estimator"]["cls"], kwargs=dict(m["estimator"]["kwargs"])),
+    )
+    inference = InferenceConfigData(
+        midi_min=i["midi_min"],
+        midi_max=i["midi_max"],
+        midi_num_bins=i["midi_num_bins"],
+        midi_std=i["midi_std"],
+    )
+    return GamemlxConfig(features=features, model=model, inference=inference, raw=raw)
+
+
+def load_lang_map(model_dir: str | pathlib.Path) -> dict[str, int] | None:
+    p = pathlib.Path(model_dir) / "lang_map.json"
+    if not p.exists():
+        return None
+    with open(p, "r", encoding="utf8") as f:
+        return json.load(f)

gamemlx/d3pm.py

@@ -0,0 +1,43 @@
+"""D3PM schedule and boundary removal, aligned with GAME modules/d3pm.py.
+
+Randomness: torch uses torch.rand_like (torch RNG). For exact alignment the
+caller can inject rng matrices (same uniform values used on both the torch
+side via monkeypatching and here). Default: numpy PCG64 generator.
+"""
+import mlx.core as mx
+import numpy as np
+
+
+def d3pm_time_schedule(t: mx.array) -> mx.array:
+    # p = (1 + cos(t * pi)) / 2
+    return (1.0 + mx.cos(t * np.pi)) / 2.0
+
+
+def remove_boundaries(boundaries: mx.array, p: mx.array, rnd: mx.array | None = None) -> mx.array:
+    """boundaries [..., T] bool; p [...]; rnd optional injected uniform [boundary shape]."""
+    q = 1.0 - p[..., None]
+    if rnd is None:
+        shape = boundaries.shape
+        rnd_np = np.random.default_rng().random(shape, dtype=np.float32)
+        rnd = mx.array(rnd_np)
+    return (rnd <= q) & boundaries
+
+
+def remove_mutable_boundaries(boundaries: mx.array, immutable: mx.array, p: mx.array,
+                              rnd: mx.array | None = None) -> mx.array:
+    """torch remove_mutable_boundaries; rnd is the uniform matrix used for the
+    mutable-boundary removal draw (shape == boundaries shape)."""
+    boundaries_mutable = boundaries & ~immutable
+    n = boundaries.astype(mx.float32).sum(axis=-1)
+    m = boundaries_mutable.astype(mx.float32).sum(axis=-1)
+    P = mx.minimum(n * p / (m + 1e-8), 1.0)
+    boundaries_mutable_remain = remove_boundaries(boundaries_mutable, P, rnd=rnd)
+    return boundaries_mutable_remain | immutable
+
+
+def make_rng_matrices(rng: np.random.Generator, shapes: list[tuple]) -> list[mx.array]:
+    """Pre-generate uniform random matrices for a whole d3pm loop (step-major)."""
+    out = []
+    for shape in shapes:
+        out.append(mx.array(rng.random(shape, dtype=np.float32)))
+    return out

gamemlx/decoding.py

@@ -0,0 +1,60 @@
+"""Decoding functions aligned with GAME modules/decoding.py."""
+import math
+
+import mlx.core as mx
+
+
+def find_local_extremum(x: mx.array, threshold: float | None = None, radius: int = 2,
+                        maxima: bool = True) -> mx.array:
+    """x: [..., T]. argmax within window [i-r, i+r] == radius (first occurrence),
+    matching torch unfold+argmax semantics."""
+    T = x.shape[-1]
+    infinite = float("+inf") if maxima else float("-inf")
+    x_pad = mx.pad(x, [(0, 0)] * (x.ndim - 1) + [(radius, radius)],
+                   constant_values=infinite)
+    wins = mx.stack([x_pad[..., i:i + T] for i in range(2 * radius + 1)], axis=0)  # [W, ..., T]
+    if maxima:
+        points = mx.argmax(wins, axis=0) == radius
+        if threshold is not None:
+            points = points & (x >= threshold)
+    else:
+        points = mx.argmin(wins, axis=0) == radius
+        if threshold is not None:
+            points = points & (x <= threshold)
+    return points
+
+
+def decode_soft_boundaries(boundaries: mx.array, barriers: mx.array = None,
+                           mask: mx.array = None, threshold: float = 0.5,
+                           radius: int = 2) -> mx.array:
+    if mask is not None:
+        boundaries = mx.where(mask, boundaries, float("+inf"))
+    if barriers is not None:
+        boundaries = mx.where(barriers, float("+inf"), boundaries)
+    maxima = find_local_extremum(boundaries, threshold=threshold, radius=radius, maxima=True)
+    if mask is not None:
+        maxima = maxima & mask
+    return maxima
+
+
+def decode_gaussian_blurred_probs(probs: mx.array, min_val: float, max_val: float,
+                                  deviation: float, threshold: float = 0.1):
+    """probs: [..., N] -> values [...], presence [...]. Aligned with torch version."""
+    ndim = probs.ndim
+    N = probs.shape[-1]
+    width = math.ceil(deviation / (max_val - min_val) * (N - 1))
+    idx = mx.arange(N)  # [N]
+    center_values = mx.linspace(min_val, max_val, N)  # torch.linspace(float32)
+    centers = mx.argmax(probs, axis=-1, keepdims=True)
+    start = mx.maximum(centers - width, 0)
+    end = mx.minimum(centers + width + 1, N)
+    # broadcast idx to probs shape: [1,...,N]
+    bshape = (1,) * (ndim - 1) + (N,)
+    idx_b = mx.broadcast_to(idx.reshape(bshape), probs.shape)
+    idx_masks = (idx_b >= start) & (idx_b < end)
+    weights = probs * idx_masks
+    product_sum = mx.sum(weights * center_values, axis=-1)
+    weight_sum = mx.sum(weights, axis=-1)
+    values = product_sum / (weight_sum + 1e-8)
+    presence = mx.max(probs, axis=-1) >= threshold
+    return values, presence

gamemlx/download.py

@@ -0,0 +1,103 @@
+"""Model download: HuggingFace primary, ModelScope fallback, local cache.
+
+Sources (verified reachable):
+  primary  HF:        baicai1145/GAME-MLX-{small,medium,large}
+  fallback ModelScope: https://modelscope.cn/models/baicai1145/GAME-MLX-{...}
+
+Files per model: config.yaml, lang_map.json,
+                 weights.safetensors (fp32), weights.fp16.safetensors (fp16).
+fp16 weights are fetched lazily; fp32 only when needed.
+"""
+import os
+import pathlib
+
+HF_REPOS = {
+    "small": "baicai1145/GAME-MLX-small",
+    "medium": "baicai1145/GAME-MLX-medium",
+    "large": "baicai1145/GAME-MLX-large",
+}
+MODELSCOPE_BASE = "https://www.modelscope.cn/models/{repo}/resolve/master/{file}"
+REQUIRED_FILES = ["config.yaml", "lang_map.json"]
+
+CACHE_DIR = pathlib.Path(os.environ.get("GAMEMLX_CACHE",
+                                        pathlib.Path.home() / ".cache/gamemlx"))
+
+
+def weights_file(precision: str) -> str:
+    return "weights.fp16.safetensors" if precision == "fp16" else "weights.safetensors"
+
+
+def _hf_download(repo_id: str, filename: str, cache_dir: pathlib.Path) -> pathlib.Path:
+    from huggingface_hub import hf_hub_download
+    return pathlib.Path(hf_hub_download(
+        repo_id=repo_id, filename=filename, cache_dir=str(cache_dir / "hf"),
+    ))
+
+
+def _modelscope_download(repo_id: str, filename: str, dest: pathlib.Path) -> pathlib.Path:
+    import urllib.request
+    url = MODELSCOPE_BASE.format(repo=repo_id, file=filename)
+    dest.parent.mkdir(parents=True, exist_ok=True)
+    tmp = dest.with_suffix(dest.suffix + ".part")
+    with urllib.request.urlopen(url) as r, open(tmp, "wb") as f:
+        while True:
+            chunk = r.read(1 << 20)
+            if not chunk:
+                break
+            f.write(chunk)
+    tmp.rename(dest)
+    return dest
+
+
+def download_model(name_or_path: str, precision: str = "fp16") -> pathlib.Path:
+    """Download (or reuse cached) model; returns the local model directory.
+
+    name_or_path: 'small' | 'medium' | 'large' | HF repo id | local path.
+    precision: which weights to fetch ('fp16' default, 'fp32' optional, 'both').
+    """
+    p = pathlib.Path(name_or_path).expanduser()
+    if (p / "config.yaml").is_file():
+        return p  # local dir
+    if name_or_path not in HF_REPOS:
+        # allow full HF repo ids like 'baicai1145/GAME-MLX-medium'
+        alias = name_or_path.split("/")[-1].lower()
+        if alias not in HF_REPOS.values():
+            raise ValueError(
+                f"Unknown model '{name_or_path}'. Aliases: {list(HF_REPOS)}, "
+                "or pass a local directory / full repo id.")
+        repo_id = name_or_path
+    else:
+        alias = name_or_path
+        repo_id = HF_REPOS[alias]
+
+    model_dir = CACHE_DIR / alias
+    model_dir.mkdir(parents=True, exist_ok=True)
+
+    def ensure(file: str) -> pathlib.Path:
+        target = model_dir / file
+        if target.is_file() and target.stat().st_size > 0:
+            return target
+        try:
+            got = _hf_download(repo_id, file, CACHE_DIR)
+            if not target.is_file():
+                try:
+                    os.link(got, target)
+                except OSError:
+                    import shutil
+                    shutil.copyfile(got, target)
+            return target
+        except Exception as e:
+            try:
+                return _modelscope_download(repo_id, file, target)
+            except Exception as e2:
+                raise RuntimeError(
+                    f"Failed to fetch {file} for {alias}: HF: {e}; ModelScope: {e2}"
+                ) from e2
+
+    for f in REQUIRED_FILES:
+        ensure(f)
+    wanted = ["weights.fp16.safetensors", "weights.safetensors"] \
+        if precision == "both" else [weights_file(precision)]
+    for f in wanted:
+        ensure(f)
+    return model_dir

gamemlx/features.py

@@ -0,0 +1,62 @@
+"""MelSpectrogram feature extraction numerically identical to GAME
+StretchableMelSpectrogram (lib/feature/mel.py) at inference defaults:
+
+- librosa mel filterbank (sr, n_fft=2048, n_mels=80, fmin, fmax, htk=False, norm='slaney')
+- torch hann_window(2048): PERIODIC window (symmetric=False), divided by nothing
+- reflect pad (win-hop)//2 and (win-hop+1)//2, center=False framing
+- rFFT -> magnitude -> mel_basis @ spec -> log(clamp(., min=1e-5))
+
+STFT computed with mlx rfft on strided frames (matches torch.stft numerics
+within float32 error; validated in alignment tests).
+"""
+import numpy as np
+
+import mlx.core as mx
+
+
+def _librosa_mel(sr, n_fft, n_mels, fmin, fmax):
+    import librosa.filters
+    return librosa.filters.mel(sr=sr, n_fft=n_fft, n_mels=n_mels,
+                               fmin=fmin, fmax=fmax, htk=False, norm="slaney")
+
+
+class MelSpectrogram:
+    def __init__(self, sample_rate=44100, n_mels=80, n_fft=2048, win_length=2048,
+                 hop_length=441, fmin=0.0, fmax=8000.0, clip_val=1e-5):
+        self.sample_rate = sample_rate
+        self.n_fft = n_fft
+        self.win_size = win_length
+        self.hop_length = hop_length
+        self.fmin = fmin
+        self.fmax = fmax
+        self.clip_val = clip_val
+        self.n_mels = n_mels
+        mel_basis = _librosa_mel(sample_rate, n_fft, n_mels, fmin, fmax)
+        self.mel_basis_np = mel_basis.astype(np.float32)
+        self.mel_basis = mx.array(self.mel_basis_np)
+        # torch hann_window(w, periodic=True) == 0.5 - 0.5*cos(2*pi*i/w), i in [0, w)
+        i = np.arange(win_length, dtype=np.float64)
+        window_np = (0.5 - 0.5 * np.cos(2.0 * np.pi * i / win_length)).astype(np.float32)
+        self.window = mx.array(window_np)
+
+    def __call__(self, y: mx.array) -> mx.array:
+        """y: [B, L] -> mel [B, n_mels, T] (like torch StretchableMelSpectrogram)."""
+        win, hop, n_fft = self.win_size, self.hop_length, self.n_fft
+        pad_l = (win - hop) // 2
+        pad_r = (win - hop + 1) // 2
+        y = mx.pad(y, [(0, 0), (pad_l, pad_r)], mode="reflect")
+        B, L = y.shape
+        n_frames = 1 + (L - win) // hop
+        if n_frames <= 0:
+            raise ValueError(f"waveform too short for one frame: {L} samples")
+        # frame via strided gather
+        starts = mx.arange(n_frames) * hop  # [T]
+        idx = starts[:, None] + mx.arange(win)[None, :]  # [T, win]
+        idx = mx.minimum(idx, L - 1)  # torch would error; guard not triggered for valid lengths
+        frames = y[:, idx.reshape(-1)].reshape(B, n_frames, win)  # [B, T, win]
+        frames = frames * self.window
+        spec = mx.fft.rfft(frames, n=n_fft, axis=-1)  # [B, T, n_fft//2+1] complex64
+        mag = mx.abs(spec)  # == torch complex abs
+        mel = self.mel_basis @ mag.transpose(0, 2, 1)  # [B, mels, T]
+        mel = mx.log(mx.maximum(mel, self.clip_val))
+        return mel