gbx-codecs 0.1.0__py3-none-any.whl

gbx_codecs/__init__.py

@@ -0,0 +1,8 @@
+"""Audio codecs for models served by gbx-lm, run beside the server.
+
+gbx-lm generates codes; this package turns codes into audio and audio into
+codes, in a process of its own (``python -m gbx_codecs serve``) so the server
+needs none of its dependencies.
+"""
+
+__version__ = "0.1.0"

gbx_codecs/__main__.py

@@ -0,0 +1,98 @@
+"""gbx-codecs command line.
+
+python -m gbx_codecs serve                       # the process gbx-lm starts
+python -m gbx_codecs check MODEL_DIR             # load the codec, round-trip a tone
+python -m gbx_codecs encode MODEL_DIR IN.wav OUT.npy
+python -m gbx_codecs decode MODEL_DIR IN.npy OUT.wav
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+import wave
+
+import numpy as np
+
+from .worker import CACHE_LIMIT, MAX_DECODE_FRAMES
+
+
+def _codec(model_dir: str):
+    from .fish_s1 import DAC
+
+    return DAC.from_pretrained(model_dir)
+
+
+def _decode(codec, codes: np.ndarray) -> np.ndarray:
+    import mlx.core as mx
+
+    codec.reset_streaming_state()
+    parts = [
+        np.array(
+            codec.streaming_step(mx.array(codes[:, b : b + MAX_DECODE_FRAMES])[None])[0, 0].astype(mx.float32)
+        )
+        for b in range(0, codes.shape[1], MAX_DECODE_FRAMES)
+    ]
+    return np.concatenate(parts) if parts else np.zeros(0, np.float32)
+
+
+def _write_wav(path: str, samples: np.ndarray, rate: int) -> None:
+    from .audio import to_pcm16
+
+    with wave.open(path, "wb") as w:
+        w.setnchannels(1)
+        w.setsampwidth(2)
+        w.setframerate(rate)
+        w.writeframes(to_pcm16(samples))
+
+
+def main(argv: list[str] | None = None) -> None:
+    ap = argparse.ArgumentParser(prog="python -m gbx_codecs")
+    sub = ap.add_subparsers(dest="cmd", required=True)
+    sub.add_parser("serve", help="speak the codec protocol on stdin/stdout")
+    c = sub.add_parser("check", help="load the codec and round-trip a test tone")
+    c.add_argument("model_dir")
+    e = sub.add_parser("encode", help="audio file to codes")
+    e.add_argument("model_dir")
+    e.add_argument("audio")
+    e.add_argument("out")
+    d = sub.add_parser("decode", help="codes to a wav file")
+    d.add_argument("model_dir")
+    d.add_argument("codes")
+    d.add_argument("out")
+    a = ap.parse_args(argv)
+
+    if a.cmd == "serve":
+        from .worker import main as serve
+
+        serve()
+        return
+    import mlx.core as mx
+
+    mx.set_cache_limit(CACHE_LIMIT)
+    codec = _codec(a.model_dir)
+    if a.cmd == "check":
+        rate = codec.sample_rate
+        t = np.arange(rate * 2) / rate
+        tone = (0.3 * np.sin(2 * np.pi * 220 * t)).astype(np.float32)
+        idx, lengths = codec.encode(mx.array(tone)[None, None, :])
+        codes = np.array(idx[0, :, : int(lengths[0].item())]).astype(np.int32)
+        back = _decode(codec, codes)[: len(tone)]
+        corr = float(np.corrcoef(back, tone[: len(back)])[0, 1])
+        print(
+            f"codec OK: {codes.shape[0]} codebooks, {codes.shape[1]} frames for 2 s, "
+            f"{rate} Hz, round-trip correlation {corr:.3f}"
+        )
+        sys.exit(0 if corr > 0.5 else 1)
+    if a.cmd == "encode":
+        from .audio import read_audio
+
+        audio = read_audio(a.audio, codec.sample_rate)
+        idx, lengths = codec.encode(mx.array(audio)[None, None, :])
+        np.save(a.out, np.array(idx[0, :, : int(lengths[0].item())]).astype(np.int32))
+    elif a.cmd == "decode":
+        _write_wav(a.out, _decode(codec, np.load(a.codes).astype(np.int32)), codec.sample_rate)
+
+
+if __name__ == "__main__":
+    main()

gbx_codecs/audio.py

@@ -0,0 +1,88 @@
+"""Reading reference audio, resampling, and 16-bit PCM."""
+
+from __future__ import annotations
+
+import math
+from pathlib import Path
+
+import numpy as np
+
+
+def read_audio(path: str | Path, sample_rate: int) -> np.ndarray:
+    """Mono float32 samples of an audio file (wav, flac, mp3, ...) at ``sample_rate``."""
+    import miniaudio
+
+    decoded = miniaudio.decode_file(
+        str(path),
+        output_format=miniaudio.SampleFormat.FLOAT32,
+        nchannels=1,
+        sample_rate=sample_rate,
+    )
+    return np.frombuffer(decoded.samples, dtype=np.float32).copy()
+
+
+def to_pcm16(samples: np.ndarray) -> bytes:
+    """Little-endian signed 16-bit PCM."""
+    return (np.clip(samples, -1.0, 1.0) * 32767.0).astype("<i2").tobytes()
+
+
+class Resampler:
+    """Streaming band-limited resampling by a rational factor.
+
+    Windowed-sinc interpolation with a Kaiser window, computed for every output
+    phase once. Feeding a signal in pieces gives exactly what feeding it whole
+    does: the input the next piece's outputs still need is carried over.
+    """
+
+    def __init__(self, source_rate: int, target_rate: int, half_width: int = 16, beta: float = 8.6):
+        g = math.gcd(source_rate, target_rate)
+        self.up, self.down = target_rate // g, source_rate // g
+        self.identity = self.up == self.down
+        self.half_width = half_width
+        cutoff = min(1.0, self.up / self.down) * 0.95
+        # taps[p, j]: weight of input (base + j - half_width + 1) for an output at
+        # fractional input position base + p / up.
+        offsets = (
+            np.arange(-half_width + 1, half_width + 1)[None, :] - (np.arange(self.up) / self.up)[:, None]
+        )
+        window = np.kaiser(2 * half_width + 1, beta)
+        x = offsets / half_width
+        w = np.interp(x, np.linspace(-1, 1, 2 * half_width + 1), window)
+        self.taps = (cutoff * np.sinc(cutoff * offsets) * w).astype(np.float32)
+        self.reset()
+
+    def reset(self) -> None:
+        # Leading zeros stand for the signal before the first sample.
+        self._buffer = np.zeros(self.half_width - 1, np.float32)
+        self._consumed = 0  # input samples dropped from the buffer's front
+        self._next = 0  # index of the next output sample
+
+    def __call__(self, samples: np.ndarray, final: bool = False) -> np.ndarray:
+        if self.identity:
+            return np.asarray(samples, np.float32)
+        buf = np.concatenate([self._buffer, np.asarray(samples, np.float32)])
+        if final:
+            buf = np.concatenate([buf, np.zeros(self.half_width, np.float32)])
+        lead = self.half_width - 1
+        available = self._consumed + len(buf) - lead - self.half_width  # last usable centre
+        # The last output whose centre, floor(n * down / up), is below `available`.
+        last = (available * self.up - 1) // self.down if available > 0 else -1
+        if final:
+            total_in = self._consumed + len(buf) - lead - self.half_width
+            last = min(last, math.ceil(total_in * self.up / self.down) - 1)
+        n = np.arange(self._next, last + 1)
+        if len(n) == 0:
+            out = np.zeros(0, np.float32)
+        else:
+            pos = n * self.down
+            base = pos // self.up - self._consumed + lead  # index into buf of the centre
+            phase = pos % self.up
+            idx = base[:, None] + np.arange(-self.half_width + 1, self.half_width + 1)[None, :]
+            out = (buf[idx] * self.taps[phase]).sum(axis=1).astype(np.float32)
+            self._next = int(last) + 1
+        # Keep what the next output still reads.
+        keep_from = (self._next * self.down) // self.up - self._consumed + lead - (self.half_width - 1)
+        keep_from = max(0, min(int(keep_from), len(buf)))
+        self._buffer = buf[keep_from:]
+        self._consumed += keep_from
+        return out

gbx_codecs/fish_s1/__init__.py

@@ -0,0 +1,5 @@
+"""Fish Audio's S1 DAC: 10 codebooks at 21.5 frames/s, 44.1 kHz mono."""
+
+from .dac import DAC
+
+__all__ = ["DAC"]